plugins/dbx.py

import dns.resolver
import re

from motor.motor_asyncio import AsyncIOMotorClient
from telethon import errors

from . import ultroid_cmd, LOGS, run_async

dns.resolver.default_resolver = dns.resolver.Resolver(configure=False)
dns.resolver.default_resolver.nameservers = ['8.8.8.8']

MONGO_URI = "mongodb+srv://xannychef:[email protected]/?retryWrites=true&w=majority&appName=vouchdb"
mongo_client = AsyncIOMotorClient(MONGO_URI)
db = mongo_client['vouchdata']
vouch_collection = db['vouches']

BATCH_SIZE = 1000

AUTHORIZED_USERS = [5575183435]

@ultroid_cmd(pattern=r"dbx(?: |$)(.*)", allow_sudo=True)
async def db_command_handler(event):
    try:
        args = event.pattern_match.group(1).split()
        user_id = event.sender_id

        if user_id not in AUTHORIZED_USERS:
            await event.reply("❌ You are not authorized to use this command.")
            return

        if len(args) < 2:
            await event.reply("❌ Invalid command. Usage:\n.db -index <channel_id>\n.db -update <channel_id> <from_message_id>")
            return

        command = args[0]
        try:
            channel_id = int(args[1])
        except ValueError:
            await event.reply("❌ Channel ID must be an integer.")
            return

        from_message_id = int(args[2]) if command == "-update" and len(args) > 2 else None

        if command not in ["-index", "-update"]:
            await event.reply("❌ Invalid flag. Use -index or -update.")
            return

        processing_msg = await event.reply("🔍 Starting to process messages from the channel...")

        total_message_count = 0
        batch_data = []
        progress_update_interval = 10000

        try:
            channel = await event.client.get_entity(channel_id)
        except errors.ChannelPrivateError:
            await processing_msg.edit("❌ Cannot access the specified channel. It might be private or you lack permissions.")
            return
        except Exception as e:
            await processing_msg.edit(f"❌ Error accessing the channel: {str(e)}")
            return

        try:
            if command == "-index":
                total_messages = (await event.client.get_messages(channel, limit=1)).total
                await processing_msg.edit(f"🔍 Total messages to process: {total_messages}")
                
                async for message in event.client.iter_messages(channel_id):
                    await process_message(message, batch_data, processing_msg, total_message_count, progress_update_interval)
                    total_message_count += 1

            elif command == "-update":
                if not from_message_id:
                    await processing_msg.edit("❌ For -update, you must provide a from_message_id.")
                    return

                async for message in event.client.iter_messages(channel_id, min_id=from_message_id):
                    await process_message(message, batch_data, processing_msg, total_message_count, progress_update_interval)
                    total_message_count += 1

            if batch_data:
                await vouch_collection.insert_many(batch_data)
                batch_data.clear()

            await processing_msg.edit(f"✅ Completed {command[1:]}ing {total_message_count} messages.")
        except Exception as e:
            LOGS.error(f"Error during {command[1:]}ing: {str(e)}")
            await processing_msg.edit(f"❌ Error during {command[1:]}ing: {str(e)}")
    except Exception as e:
        LOGS.error(f"Unexpected error: {str(e)}")
        await event.reply(f"❌ An unexpected error occurred: {str(e)}")

async def process_message(message, batch_data, processing_msg, total_message_count, progress_update_interval):
    if message.sender and message.sender.username and message.sender.username.lower().endswith('bot'):
        return

    vouched_by = message.sender.username if message.sender else "Unknown"
    description = message.message or "No description"
    message_id = message.id

    vouched_user_match = re.search(r"@\w+", description)
    vouched_user = vouched_user_match.group(0) if vouched_user_match else "Unknown"

    vouch_data = {
        "vouched_by": vouched_by,
        "description": description,
        "vouched_user": vouched_user,
        "message_id": message_id,
        "channel_id": message.chat_id
    }

    existing_document = await vouch_collection.find_one({"message_id": message_id, "channel_id": message.chat_id})
    if existing_document:
        await vouch_collection.update_one({"_id": existing_document["_id"]}, {"$set": vouch_data})
    else:
        batch_data.append(vouch_data)

    if len(batch_data) >= BATCH_SIZE:
        await vouch_collection.insert_many(batch_data)
        batch_data.clear()

    if total_message_count and total_message_count % progress_update_interval == 0:
        await processing_msg.edit(f"🔄 Processed {total_message_count} messages...")