TelegramBotAdmin/wbb/utils/dbadmin.py

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)
        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,
    }


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",
    }
    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({}),
    }