TelegramBotAdmin/wbb/admin/bot_config.py

from __future__ import annotations

import json
import re
import tempfile
import threading
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
from uuid import uuid4

from aiohttp import ClientError, ClientSession, ClientTimeout

BOT_TOKEN_PATTERN = re.compile(r"^\d{5,}:[A-Za-z0-9_-]{20,}$")
API_HASH_PATTERN = re.compile(r"^[A-Fa-f0-9]{32}$")
PROFILE_FIELDS = {
    "sudo_users_id",
    "log_group_id",
    "gban_log_group_id",
    "message_dump_chat",
}
_STORE_LOCK = threading.RLock()


class BotConfigError(ValueError):
    def __init__(self, code: str, message: str):
        super().__init__(message)
        self.code = code


def _empty_document() -> dict[str, Any]:
    return {
        "version": 1,
        "telegram": {"api_id": 0, "api_hash": ""},
        "bots": [],
    }


def _read_document(path: str | Path) -> dict[str, Any]:
    config_path = Path(path)
    if not config_path.exists():
        return _empty_document()
    try:
        data = json.loads(config_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as exc:
        raise BotConfigError(
            "bot_config_corrupt",
            "机器人配置文件无法读取。",
        ) from exc
    if not isinstance(data, dict) or not isinstance(data.get("bots", []), list):
        raise BotConfigError("bot_config_corrupt", "机器人配置文件结构无效。")
    telegram = data.get("telegram")
    if not isinstance(telegram, dict):
        telegram = {}
    return {
        "version": 1,
        "telegram": {
            "api_id": int(telegram.get("api_id") or 0),
            "api_hash": str(telegram.get("api_hash") or ""),
        },
        "bots": [item for item in data.get("bots", []) if isinstance(item, dict)],
    }


def _write_document(path: str | Path, data: dict[str, Any]) -> None:
    config_path = Path(path).resolve()
    config_path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(
        mode="w",
        encoding="utf-8",
        dir=config_path.parent,
        delete=False,
    ) as handle:
        temporary_path = Path(handle.name)
        json.dump(data, handle, ensure_ascii=True, indent=2)
        handle.write("\n")
    try:
        temporary_path.chmod(0o600)
        temporary_path.replace(config_path)
        config_path.chmod(0o600)
    finally:
        temporary_path.unlink(missing_ok=True)


def _integer(value: Any, name: str, *, positive: bool = False) -> int:
    try:
        parsed = int(value)
    except (TypeError, ValueError) as exc:
        raise BotConfigError("invalid_bot_config", f"{name} 必须是整数。") from exc
    if positive and parsed <= 0:
        raise BotConfigError("invalid_bot_config", f"{name} 必须大于 0。")
    return parsed


def _normalize_sudoers(value: Any) -> list[int]:
    items = value if isinstance(value, list) else str(value or "").replace(",", " ").split()
    normalized = [_integer(item, "SUDOERS 用户 ID", positive=True) for item in items]
    return list(dict.fromkeys(normalized))


def _find_profile(data: dict[str, Any], bot_id: str) -> dict[str, Any]:
    profile = next(
        (item for item in data["bots"] if str(item.get("bot_id")) == str(bot_id)),
        None,
    )
    if profile is None:
        raise BotConfigError("bot_not_found", "未找到该机器人配置。")
    return profile


def _public_profile(
    profile: dict[str, Any],
    *,
    api_ready: bool,
    runtime: dict[str, Any] | None = None,
) -> dict[str, Any]:
    token_configured = bool(BOT_TOKEN_PATTERN.fullmatch(str(profile.get("bot_token") or "")))
    identity = profile.get("identity") if isinstance(profile.get("identity"), dict) else None
    return {
        "bot_id": str(profile.get("bot_id") or ""),
        "label": str(profile.get("label") or ""),
        "enabled": bool(profile.get("enabled", True)),
        "bot_token_configured": token_configured,
        "sudo_users_id": [str(item) for item in profile.get("sudo_users_id", [])],
        "log_group_id": str(profile.get("log_group_id") or 0),
        "gban_log_group_id": str(profile.get("gban_log_group_id") or 0),
        "message_dump_chat": str(profile.get("message_dump_chat") or 0),
        "identity": identity,
        "ready_to_connect": api_ready and token_configured,
        "created_at": profile.get("created_at"),
        "updated_at": profile.get("updated_at"),
        "runtime": runtime or {"state": "stopped"},
    }


def telegram_config_status(
    path: str | Path,
    *,
    runtimes: dict[str, dict[str, Any]] | None = None,
) -> dict[str, Any]:
    with _STORE_LOCK:
        data = _read_document(path)
    telegram = data["telegram"]
    api_id = int(telegram.get("api_id") or 0)
    api_hash_configured = bool(API_HASH_PATTERN.fullmatch(str(telegram.get("api_hash") or "")))
    api_ready = api_id > 0 and api_hash_configured
    runtime_map = runtimes or {}
    return {
        "api_id": api_id or None,
        "api_hash_configured": api_hash_configured,
        "api_ready": api_ready,
        "bots": [
            _public_profile(
                profile,
                api_ready=api_ready,
                runtime=runtime_map.get(str(profile.get("bot_id"))),
            )
            for profile in data["bots"]
        ],
    }


def update_telegram_config(path: str | Path, body: dict[str, Any]) -> dict[str, Any]:
    with _STORE_LOCK:
        data = _read_document(path)
        telegram = data["telegram"]
        changed = False
        if "api_id" in body:
            telegram["api_id"] = _integer(
                body.get("api_id"), "Telegram App API ID", positive=True
            )
            changed = True
        api_hash = str(body.get("api_hash") or "").strip()
        if api_hash:
            if not API_HASH_PATTERN.fullmatch(api_hash):
                raise BotConfigError(
                    "invalid_api_hash",
                    "Telegram 应用密钥必须是 my.telegram.org/apps 提供的 "
                    "32 位十六进制字符串，不能使用机器人令牌。",
                )
            telegram["api_hash"] = api_hash
            changed = True
        if not changed:
            raise BotConfigError(
                "bot_config_unchanged",
                "没有可保存的 Telegram 应用接口配置。",
            )
        _write_document(path, data)
    return telegram_config_status(path)


def create_bot_profile(path: str | Path, body: dict[str, Any]) -> dict[str, Any]:
    label = str(body.get("label") or "").strip()
    token = str(body.get("bot_token") or "").strip()
    if not label or len(label) > 60:
        raise BotConfigError(
            "invalid_bot_label",
            "机器人名称长度需要在 1 到 60 个字符之间。",
        )
    if not BOT_TOKEN_PATTERN.fullmatch(token):
        raise BotConfigError("invalid_bot_token", "机器人令牌格式无效。")
    now = datetime.now(UTC).isoformat().replace("+00:00", "Z")
    profile = {
        "bot_id": uuid4().hex,
        "label": label,
        "bot_token": token,
        "enabled": bool(body.get("enabled", True)),
        "sudo_users_id": _normalize_sudoers(body.get("sudo_users_id", [])),
        "log_group_id": _integer(body.get("log_group_id", 0), "日志群 ID"),
        "gban_log_group_id": _integer(body.get("gban_log_group_id", 0), "全局封禁日志群 ID"),
        "message_dump_chat": _integer(body.get("message_dump_chat", 0), "媒体中转群 ID"),
        "identity": None,
        "created_at": now,
        "updated_at": now,
    }
    with _STORE_LOCK:
        data = _read_document(path)
        data["bots"].append(profile)
        _write_document(path, data)
        api_ready = bool(data["telegram"]["api_id"] and data["telegram"]["api_hash"])
    return _public_profile(profile, api_ready=api_ready)


def update_bot_profile(path: str | Path, bot_id: str, body: dict[str, Any]) -> dict[str, Any]:
    with _STORE_LOCK:
        data = _read_document(path)
        profile = _find_profile(data, bot_id)
        if "label" in body:
            label = str(body.get("label") or "").strip()
            if not label or len(label) > 60:
                raise BotConfigError(
                    "invalid_bot_label",
                    "机器人名称长度需要在 1 到 60 个字符之间。",
                )
            profile["label"] = label
        token = str(body.get("bot_token") or "").strip()
        if token:
            if not BOT_TOKEN_PATTERN.fullmatch(token):
                raise BotConfigError("invalid_bot_token", "机器人令牌格式无效。")
            profile["bot_token"] = token
            profile["identity"] = None
        if "enabled" in body:
            profile["enabled"] = bool(body.get("enabled"))
        if "sudo_users_id" in body:
            profile["sudo_users_id"] = _normalize_sudoers(body.get("sudo_users_id"))
        for field in PROFILE_FIELDS - {"sudo_users_id"}:
            if field in body:
                profile[field] = _integer(body.get(field), field)
        profile["updated_at"] = datetime.now(UTC).isoformat().replace("+00:00", "Z")
        _write_document(path, data)
        api_ready = bool(data["telegram"]["api_id"] and data["telegram"]["api_hash"])
    return _public_profile(profile, api_ready=api_ready)


def delete_bot_profile(path: str | Path, bot_id: str) -> None:
    with _STORE_LOCK:
        data = _read_document(path)
        _find_profile(data, bot_id)
        data["bots"] = [
            item for item in data["bots"] if str(item.get("bot_id")) != str(bot_id)
        ]
        _write_document(path, data)


def get_bot_profile_secrets(path: str | Path, bot_id: str) -> dict[str, Any]:
    with _STORE_LOCK:
        data = _read_document(path)
        profile = dict(_find_profile(data, bot_id))
        telegram = dict(data["telegram"])
    return {**profile, **telegram}


def store_bot_identity(path: str | Path, bot_id: str, identity: dict[str, Any]) -> None:
    with _STORE_LOCK:
        data = _read_document(path)
        profile = _find_profile(data, bot_id)
        profile["identity"] = {
            "id": str(identity.get("id") or ""),
            "username": str(identity.get("username") or ""),
            "name": str(identity.get("name") or ""),
        }
        profile["updated_at"] = datetime.now(UTC).isoformat().replace("+00:00", "Z")
        _write_document(path, data)


async def test_bot_token(token: str) -> dict[str, Any]:
    if not BOT_TOKEN_PATTERN.fullmatch(token):
        raise BotConfigError("invalid_bot_token", "机器人令牌格式无效。")
    timeout = ClientTimeout(total=12)
    try:
        async with ClientSession(timeout=timeout) as session:
            async with session.post(f"https://api.telegram.org/bot{token}/getMe") as response:
                payload = await response.json(content_type=None)
    except (ClientError, TimeoutError, ValueError) as exc:
        raise BotConfigError(
            "telegram_unreachable",
            "无法连接 Telegram 机器人接口。",
        ) from exc
    if response.status != 200 or not payload.get("ok"):
        raise BotConfigError("telegram_auth_failed", "机器人令牌验证失败。")
    result = payload.get("result") or {}
    return {
        "id": str(result.get("id") or ""),
        "username": result.get("username") or "",
        "name": " ".join(
            part for part in (result.get("first_name"), result.get("last_name")) if part
        ),
    }