from __future__ import annotations import asyncio import re from datetime import UTC, datetime from secrets import token_hex from typing import Any from pymongo import ASCENDING, DESCENDING from wbb import BOT_PROFILE_ID, control_db, db admin_usersdb = control_db.admin_users admin_sessionsdb = control_db.admin_sessions auditdb = db.admin_audit_logs managed_chatsdb = db.managed_chats managed_chat_settingsdb = db.managed_chat_settings invite_linksdb = db.admin_invite_links recent_chat_membersdb = db.recent_chat_members member_identity_changesdb = db.member_identity_changes _index_lock = asyncio.Lock() _indexes_ready = False def utc_now() -> datetime: return datetime.now(UTC) async def ensure_admin_indexes() -> None: global _indexes_ready if _indexes_ready: return async with _index_lock: if _indexes_ready: return await admin_usersdb.create_index([("username", ASCENDING)], unique=True) await admin_sessionsdb.create_index([("token_hash", ASCENDING)], unique=True) await admin_sessionsdb.create_index("expires_at", expireAfterSeconds=0) await auditdb.create_index([("created_at", DESCENDING)]) await auditdb.create_index( [("chat_id", ASCENDING), ("created_at", DESCENDING)] ) await managed_chatsdb.create_index( [("bot_id", ASCENDING), ("chat_id", ASCENDING)], unique=True ) await managed_chatsdb.create_index([("last_seen_at", DESCENDING)]) await managed_chat_settingsdb.create_index( [("bot_id", ASCENDING), ("chat_id", ASCENDING)], unique=True ) await invite_linksdb.create_index( [("bot_id", ASCENDING), ("chat_id", ASCENDING), ("created_at", DESCENDING)] ) await recent_chat_membersdb.create_index( [ ("bot_id", ASCENDING), ("chat_id", ASCENDING), ("user_id", ASCENDING), ], unique=True, ) await recent_chat_membersdb.create_index( [("bot_id", ASCENDING), ("chat_id", ASCENDING), ("last_seen_at", DESCENDING)] ) await member_identity_changesdb.create_index( [("bot_id", ASCENDING), ("chat_id", ASCENDING), ("observed_at", DESCENDING)] ) await member_identity_changesdb.create_index( [ ("bot_id", ASCENDING), ("chat_id", ASCENDING), ("user_id", ASCENDING), ("observed_at", DESCENDING), ] ) _indexes_ready = True async def ensure_default_admin( *, username: str, password_hash: str ) -> dict[str, Any]: await ensure_admin_indexes() now = utc_now() await admin_usersdb.update_one( {"username": username}, { "$setOnInsert": { "username": username, "password_hash": password_hash, "must_change_password": True, "failed_login_count": 0, "created_at": now, "updated_at": now, } }, upsert=True, ) return await admin_usersdb.find_one({"username": username}) async def get_admin_user(username: str) -> dict[str, Any] | None: await ensure_admin_indexes() return await admin_usersdb.find_one({"username": username}) async def record_login_failure(username: str) -> None: await ensure_admin_indexes() await admin_usersdb.update_one( {"username": username}, { "$inc": {"failed_login_count": 1}, "$set": {"last_failed_login_at": utc_now(), "updated_at": utc_now()}, }, ) async def record_login_success(username: str) -> None: await ensure_admin_indexes() await admin_usersdb.update_one( {"username": username}, { "$set": { "failed_login_count": 0, "last_login_at": utc_now(), "updated_at": utc_now(), } }, ) async def update_admin_password(username: str, password_hash: str) -> bool: await ensure_admin_indexes() result = await admin_usersdb.update_one( {"username": username}, { "$set": { "password_hash": password_hash, "must_change_password": False, "password_updated_at": utc_now(), "updated_at": utc_now(), } }, ) if result.modified_count: await admin_sessionsdb.delete_many({"username": username}) return result.modified_count == 1 async def create_admin_session( *, username: str, token_hash: str, csrf_token: str, expires_at: datetime, remote_address: str, user_agent: str, ) -> None: await ensure_admin_indexes() now = utc_now() await admin_sessionsdb.insert_one( { "username": username, "token_hash": token_hash, "csrf_token": csrf_token, "remote_address": remote_address[:200], "user_agent": user_agent[:500], "created_at": now, "last_seen_at": now, "expires_at": expires_at, } ) async def get_admin_session(token_hash: str) -> dict[str, Any] | None: await ensure_admin_indexes() session = await admin_sessionsdb.find_one( {"token_hash": token_hash, "expires_at": {"$gt": utc_now()}} ) if not session: return None user = await get_admin_user(session["username"]) if not user: return None await admin_sessionsdb.update_one( {"_id": session["_id"]}, {"$set": {"last_seen_at": utc_now()}} ) return {**session, "user": user} async def revoke_admin_session(token_hash: str) -> None: await ensure_admin_indexes() await admin_sessionsdb.delete_one({"token_hash": token_hash}) async def revoke_admin_sessions(username: str) -> None: await ensure_admin_indexes() await admin_sessionsdb.delete_many({"username": username}) async def upsert_managed_chat( *, chat_id: int, title: str | None, username: str | None, chat_type: str, member_count: int | None = None, accessible: bool = True, bot_status: str | None = None, bot_privileges: list[str] | None = None, ) -> None: await ensure_admin_indexes() now = utc_now() values: dict[str, Any] = { "title": title or str(chat_id), "username": username, "type": chat_type, "accessible": accessible, "last_seen_at": now, "updated_at": now, } if member_count is not None: values["member_count"] = int(member_count) if bot_status is not None: values["bot_status"] = bot_status if bot_privileges is not None: values["bot_privileges"] = bot_privileges await managed_chatsdb.update_one( {"bot_id": BOT_PROFILE_ID, "chat_id": chat_id}, { "$set": values, "$setOnInsert": {"bot_id": BOT_PROFILE_ID, "created_at": now}, }, upsert=True, ) async def mark_managed_chat_unavailable(chat_id: int, reason: str) -> None: await ensure_admin_indexes() await managed_chatsdb.update_one( {"bot_id": BOT_PROFILE_ID, "chat_id": chat_id}, { "$set": { "accessible": False, "last_error": reason[:500], "updated_at": utc_now(), } }, upsert=True, ) async def delete_unavailable_managed_chat(chat_id: int) -> bool: """Remove an unavailable chat from this bot's management index only.""" await ensure_admin_indexes() result = await managed_chatsdb.delete_one( { "bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id), "accessible": False, } ) return result.deleted_count == 1 async def managed_chat_exists(chat_id: int) -> bool: await ensure_admin_indexes() return ( await managed_chatsdb.find_one( {"bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id)}, {"_id": 1} ) is not None ) async def list_managed_chats( *, query: str = "", page: int = 1, page_size: int = 20 ) -> tuple[list[dict[str, Any]], int]: await ensure_admin_indexes() filters: dict[str, Any] = {"bot_id": BOT_PROFILE_ID} query = query.strip() if query: if query.lstrip("-").isdigit(): filters["chat_id"] = int(query) else: value = re.escape(query.lstrip("@")) filters["$or"] = [ {"title": {"$regex": value, "$options": "i"}}, {"username": {"$regex": value, "$options": "i"}}, ] page = max(1, page) page_size = max(1, min(page_size, 100)) total = await managed_chatsdb.count_documents(filters) cursor = ( managed_chatsdb.find(filters) .sort([("accessible", DESCENDING), ("last_seen_at", DESCENDING)]) .skip((page - 1) * page_size) .limit(page_size) ) return [doc async for doc in cursor], total async def record_audit( *, source: str, actor_id: int | str | None, actor_name: str, action: str, chat_id: int | None = None, target_id: int | str | None = None, summary: str = "", success: bool = True, error: str = "", metadata: dict[str, Any] | None = None, ) -> None: await ensure_admin_indexes() await auditdb.insert_one( { "source": source, "actor_id": actor_id, "actor_name": actor_name, "action": action, "chat_id": chat_id, "target_id": target_id, "summary": summary[:1000], "success": bool(success), "error": error[:1000], "metadata": metadata or {}, "bot_id": BOT_PROFILE_ID, "created_at": utc_now(), } ) async def list_audit_logs( *, chat_id: int | None = None, action: str | None = None, page: int = 1, page_size: int = 20, ) -> tuple[list[dict[str, Any]], int]: await ensure_admin_indexes() filters: dict[str, Any] = {} if chat_id is not None: filters["chat_id"] = chat_id if action: filters["action"] = action page = max(1, page) page_size = max(1, min(page_size, 100)) total = await auditdb.count_documents(filters) cursor = ( auditdb.find(filters) .sort("created_at", DESCENDING) .skip((page - 1) * page_size) .limit(page_size) ) return [doc async for doc in cursor], total async def get_managed_chat_settings(chat_id: int) -> dict[str, Any]: await ensure_admin_indexes() doc = await managed_chat_settingsdb.find_one( {"bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id)} ) if doc is not None: doc.setdefault("captcha_enabled", False) doc.setdefault("interaction_settings", { "cleanup_seconds": 10, "commands": {key: True for key in ("giveaway", "checkin", "points", "history", "shortcuts")}, }) return doc return { "chat_id": int(chat_id), "auto_replies": [], "blacklist_words": [], "welcome": {"enabled": False, "text": "", "media": None}, "identity_monitor": {"enabled": False, "notify_in_chat": False}, "captcha_enabled": False, "chatbot_enabled": False, "antiflood_enabled": True, "interaction_settings": { "cleanup_seconds": 10, "commands": {key: True for key in ("giveaway", "checkin", "points", "history", "shortcuts")}, }, } async def update_managed_chat_settings( chat_id: int, values: dict[str, Any] ) -> dict[str, Any]: await ensure_admin_indexes() allowed = { "auto_replies", "blacklist_words", "risk_control", "risk_rules", "identity_monitor", "welcome", "captcha_enabled", "chatbot_enabled", "antiflood_enabled", "interaction_settings", } update = {key: value for key, value in values.items() if key in allowed} now = utc_now() await managed_chat_settingsdb.update_one( {"bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id)}, { "$set": {**update, "updated_at": now}, "$setOnInsert": { "bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id), "created_at": now, }, }, upsert=True, ) return await get_managed_chat_settings(chat_id) async def upsert_recent_chat_member( *, chat_id: int, user_id: int, username: str | None, first_name: str | None, last_name: str | None, is_bot: bool, ) -> None: await ensure_admin_indexes() now = utc_now() await recent_chat_membersdb.update_one( { "bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id), "user_id": int(user_id), }, { "$set": { "username": username, "first_name": first_name, "last_name": last_name, "is_bot": bool(is_bot), "last_seen_at": now, }, "$inc": {"message_count": 1}, "$setOnInsert": { "bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id), "user_id": int(user_id), "created_at": now, }, }, upsert=True, ) async def get_recent_chat_member(chat_id: int, user_id: int) -> dict[str, Any] | None: await ensure_admin_indexes() return await recent_chat_membersdb.find_one( { "bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id), "user_id": int(user_id), } ) async def record_member_identity_change( *, chat_id: int, user_id: int, before: dict[str, Any], after: dict[str, Any], changed_fields: list[str], ) -> dict[str, Any]: await ensure_admin_indexes() event = { "change_id": token_hex(10), "bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id), "user_id": int(user_id), "before": before, "after": after, "changed_fields": changed_fields, "observed_at": utc_now(), } await member_identity_changesdb.insert_one(event) return event async def list_member_identity_changes( chat_id: int, *, page: int = 1, page_size: int = 20, ) -> tuple[list[dict[str, Any]], int]: await ensure_admin_indexes() filters = {"bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id)} page = max(1, int(page)) page_size = max(1, min(int(page_size), 100)) total = await member_identity_changesdb.count_documents(filters) cursor = ( member_identity_changesdb.find(filters) .sort("observed_at", DESCENDING) .skip((page - 1) * page_size) .limit(page_size) ) return [item async for item in cursor], total async def list_recent_chat_members( chat_id: int, *, query: str = "", limit: int = 30, ) -> list[dict[str, Any]]: await ensure_admin_indexes() filters: dict[str, Any] = { "bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id), "is_bot": {"$ne": True}, } normalized_query = query.strip().lstrip("@") if normalized_query: if normalized_query.isdigit(): filters["user_id"] = int(normalized_query) else: pattern = re.escape(normalized_query) filters["$or"] = [ {"username": {"$regex": pattern, "$options": "i"}}, {"first_name": {"$regex": pattern, "$options": "i"}}, {"last_name": {"$regex": pattern, "$options": "i"}}, ] cursor = ( recent_chat_membersdb.find(filters) .sort("last_seen_at", DESCENDING) .limit(max(1, min(int(limit), 50))) ) return [item async for item in cursor] async def store_invite_link( *, chat_id: int, invite_link: str, name: str, expires_at: datetime | None ) -> None: await ensure_admin_indexes() await invite_linksdb.update_one( { "bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id), "invite_link": invite_link, }, { "$set": { "name": name, "expires_at": expires_at, "revoked": False, "updated_at": utc_now(), }, "$setOnInsert": {"bot_id": BOT_PROFILE_ID, "created_at": utc_now()}, }, upsert=True, ) async def revoke_stored_invite_link(chat_id: int, invite_link: str) -> None: await ensure_admin_indexes() await invite_linksdb.update_one( { "bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id), "invite_link": invite_link, }, {"$set": {"revoked": True, "revoked_at": utc_now()}}, ) async def list_stored_invite_links(chat_id: int) -> list[dict[str, Any]]: await ensure_admin_indexes() cursor = invite_linksdb.find( {"bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id)} ).sort( "created_at", DESCENDING ) return [doc async for doc in cursor] async def dashboard_counts() -> dict[str, int]: await ensure_admin_indexes() return { "chats": await managed_chatsdb.count_documents({"accessible": {"$ne": False}}), "audit_events": await auditdb.count_documents({}), }