TelegramBotAdmin/wbb/services/external_risk_rules.py

from __future__ import annotations

import hashlib
import re
import unicodedata
from datetime import UTC, datetime
from secrets import token_hex
from typing import Any
from urllib.parse import urlsplit

from aiohttp import ClientSession, ClientTimeout
from pymongo import ASCENDING

from wbb import BOT_PROFILE_ID, db, log

SOURCES = {
    "yt_spam_domains": {
        "name": "YT-Spam-Lists 域名",
        "url": "https://raw.githubusercontent.com/ThioJoe/YT-Spam-Lists/main/SpamDomainsList.txt",
        "repository": "https://github.com/ThioJoe/YT-Spam-Lists",
        "license": "MIT",
        "kind": "domain",
        "max_bytes": 1_000_000,
    },
    "phishing_domains": {
        "name": "Phishing.Database 活跃域名",
        "url": "https://raw.githubusercontent.com/Phishing-Database/Phishing.Database/master/phishing-domains-ACTIVE.txt",
        "repository": "https://github.com/Phishing-Database/Phishing.Database",
        "license": "MIT",
        "kind": "domain",
        "max_bytes": 20_000_000,
    },
    "zapper_phrases": {
        "name": "zapper-TGAB 短语候选",
        "url": "https://raw.githubusercontent.com/RGB-Outl4w/zapper-TGAB/master/spam_phrases.txt",
        "repository": "https://github.com/RGB-Outl4w/zapper-TGAB",
        "license": "MIT",
        "kind": "phrase",
        "max_bytes": 100_000,
    },
}

snapshotsdb = db.external_risk_snapshots
entriesdb = db.external_risk_entries
statesdb = db.external_risk_sources
policiesdb = db.external_risk_chat_policies
eventsdb = db.external_risk_events
DOMAIN_RE = re.compile(r"^(?:[a-z0-9-]{1,63}\.)+[a-z0-9-]{2,63}$")
LINK_RE = re.compile(r"(?:(?:https?://|www\.)[^\s<>]+)", re.IGNORECASE)
_indexes_ready = False


async def ensure_indexes() -> None:
    global _indexes_ready
    if _indexes_ready:
        return
    await snapshotsdb.create_index([("snapshot_id", ASCENDING)], unique=True)
    await snapshotsdb.create_index([("bot_id", ASCENDING), ("source_id", ASCENDING), ("created_at", -1)])
    await entriesdb.create_index([("snapshot_id", ASCENDING), ("value", ASCENDING)], unique=True)
    await statesdb.create_index([("bot_id", ASCENDING), ("source_id", ASCENDING)], unique=True)
    await policiesdb.create_index([("bot_id", ASCENDING), ("chat_id", ASCENDING)], unique=True)
    await eventsdb.create_index([("bot_id", ASCENDING), ("chat_id", ASCENDING), ("created_at", -1)])
    _indexes_ready = True


def normalize_domain(raw: str) -> str | None:
    value = raw.strip().lower().lstrip("*.").rstrip(".")
    if value.startswith(("http://", "https://")):
        value = urlsplit(value).hostname or ""
    value = value.split("/")[0].split(":")[0]
    try:
        value = value.encode("idna").decode("ascii")
    except UnicodeError:
        return None
    if not DOMAIN_RE.fullmatch(value) or ".." in value:
        return None
    return value


def parse_entries(source_id: str, content: bytes) -> list[str]:
    source = SOURCES[source_id]
    values: set[str] = set()
    for line in content.decode("utf-8-sig", errors="replace").splitlines():
        line = line.strip()
        if not line or line.startswith(("#", "//")):
            continue
        if source["kind"] == "domain":
            domain = normalize_domain(line)
            if domain:
                values.add(domain)
        else:
            phrase = unicodedata.normalize("NFKC", line).casefold().strip()
            if 5 <= len(phrase) <= 120:
                values.add(phrase)
    return sorted(values)


async def stage_snapshot(source_id: str) -> dict[str, Any]:
    await ensure_indexes()
    if source_id not in SOURCES:
        raise ValueError("未知外部规则来源。")
    source = SOURCES[source_id]
    async with ClientSession(timeout=ClientTimeout(total=60)) as session:
        async with session.get(source["url"], allow_redirects=False) as response:
            response.raise_for_status()
            chunks: list[bytes] = []
            size = 0
            async for chunk in response.content.iter_chunked(65536):
                chunks.append(chunk)
                size += len(chunk)
                if size > int(source["max_bytes"]):
                    raise ValueError("外部规则文件超过大小限制。")
            content = b"".join(chunks)
    if len(content) > int(source["max_bytes"]):
        raise ValueError("外部规则文件超过大小限制。")
    values = parse_entries(source_id, content)
    if not values:
        raise ValueError("外部规则文件没有可导入的条目。")
    digest = hashlib.sha256(content).hexdigest()
    existing = await snapshotsdb.find_one(
        {"bot_id": BOT_PROFILE_ID, "source_id": source_id, "sha256": digest}
    )
    if existing and existing.get("status") in {"pending", "published"}:
        return existing
    snapshot_id = token_hex(12)
    now = datetime.now(UTC)
    snapshot = {
        "snapshot_id": snapshot_id, "bot_id": BOT_PROFILE_ID,
        "source_id": source_id, "source_url": source["url"],
        "repository": source["repository"], "license": source["license"],
        "kind": source["kind"], "sha256": digest, "entry_count": len(values),
        "status": "pending", "created_at": now,
    }
    await snapshotsdb.insert_one(snapshot)
    try:
        for offset in range(0, len(values), 1000):
            await entriesdb.insert_many(
                [{"snapshot_id": snapshot_id, "value": value} for value in values[offset:offset + 1000]]
            )
    except Exception:
        await snapshotsdb.update_one(
            {"snapshot_id": snapshot_id}, {"$set": {"status": "failed"}}
        )
        raise
    return snapshot


async def list_snapshots(source_id: str | None = None) -> list[dict[str, Any]]:
    await ensure_indexes()
    query: dict[str, Any] = {"bot_id": BOT_PROFILE_ID}
    if source_id:
        query["source_id"] = source_id
    return [item async for item in snapshotsdb.find(query).sort("created_at", -1).limit(100)]


async def list_active_sources() -> list[dict[str, Any]]:
    await ensure_indexes()
    return [item async for item in statesdb.find({"bot_id": BOT_PROFILE_ID})]


async def list_snapshot_entries(
    snapshot_id: str, *, page: int = 1, page_size: int = 50, query: str = ""
) -> tuple[list[dict[str, Any]], int]:
    await ensure_indexes()
    snapshot = await snapshotsdb.find_one({"snapshot_id": snapshot_id, "bot_id": BOT_PROFILE_ID})
    if not snapshot:
        raise ValueError("规则快照不存在。")
    filters: dict[str, Any] = {"snapshot_id": snapshot_id}
    if query.strip():
        filters["value"] = {"$regex": re.escape(query.strip()[:100]), "$options": "i"}
    page = max(1, int(page))
    page_size = max(1, min(100, int(page_size)))
    total = await entriesdb.count_documents(filters)
    cursor = entriesdb.find(filters).sort("value", ASCENDING).skip((page - 1) * page_size).limit(page_size)
    return [item async for item in cursor], total


async def publish_snapshot(snapshot_id: str) -> dict[str, Any]:
    await ensure_indexes()
    snapshot = await snapshotsdb.find_one({"snapshot_id": snapshot_id, "bot_id": BOT_PROFILE_ID})
    if not snapshot or snapshot["status"] not in {"pending", "published"}:
        raise ValueError("快照不存在或不可发布。")
    await statesdb.update_one(
        {"bot_id": BOT_PROFILE_ID, "source_id": snapshot["source_id"]},
        {"$set": {"active_snapshot_id": snapshot_id, "updated_at": datetime.now(UTC)}},
        upsert=True,
    )
    await snapshotsdb.update_one(
        {"snapshot_id": snapshot_id}, {"$set": {"status": "published", "published_at": datetime.now(UTC)}}
    )
    return snapshot


async def get_policy(chat_id: int) -> dict[str, Any]:
    await ensure_indexes()
    stored = await policiesdb.find_one({"bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id)})
    return {"chat_id": int(chat_id), "mode": (stored or {}).get("mode", "observe")}


async def set_policy(chat_id: int, mode: str) -> dict[str, Any]:
    if mode not in {"off", "observe", "delete"}:
        raise ValueError("风控模式只能是关闭、观察或删消息。")
    await ensure_indexes()
    await policiesdb.update_one(
        {"bot_id": BOT_PROFILE_ID, "chat_id": int(chat_id)},
        {"$set": {"mode": mode, "updated_at": datetime.now(UTC)}}, upsert=True,
    )
    return {"chat_id": int(chat_id), "mode": mode}


def _message_domains(message: Any) -> set[str]:
    text = str(getattr(message, "text", None) or getattr(message, "caption", None) or "")
    links = LINK_RE.findall(text)
    for entity in list(getattr(message, "entities", None) or []) + list(getattr(message, "caption_entities", None) or []):
        if getattr(entity, "url", None):
            links.append(str(entity.url))
    result: set[str] = set()
    for link in links:
        domain = normalize_domain(link.rstrip(".,!?;:)]}"))
        if domain:
            labels = domain.split(".")
            result.update(".".join(labels[index:]) for index in range(len(labels) - 1))
    return result


async def detect_external(message: Any) -> list[dict[str, Any]]:
    await ensure_indexes()
    active = [item async for item in statesdb.find({"bot_id": BOT_PROFILE_ID, "active_snapshot_id": {"$exists": True}})]
    if not active:
        return []
    domains = _message_domains(message)
    text = unicodedata.normalize(
        "NFKC", str(getattr(message, "text", None) or getattr(message, "caption", None) or "")
    ).casefold()
    matches: list[dict[str, Any]] = []
    for state in active:
        snapshot = await snapshotsdb.find_one({"snapshot_id": state["active_snapshot_id"]})
        if not snapshot:
            continue
        if snapshot["kind"] == "domain" and domains:
            entry = await entriesdb.find_one(
                {"snapshot_id": snapshot["snapshot_id"], "value": {"$in": list(domains)}}
            )
            if entry:
                matches.append({"source_id": snapshot["source_id"], "snapshot_id": snapshot["snapshot_id"], "value": entry["value"]})
        elif snapshot["kind"] == "phrase" and text:
            entries = [item async for item in entriesdb.find({"snapshot_id": snapshot["snapshot_id"]})]
            hit = next((item for item in entries if item["value"] in text), None)
            if hit:
                matches.append({"source_id": snapshot["source_id"], "snapshot_id": snapshot["snapshot_id"], "value": hit["value"]})
    return matches


async def enforce_external(message: Any, *, already_deleted: bool = False) -> None:
    chat_id = int(message.chat.id)
    policy = await get_policy(chat_id)
    if policy["mode"] == "off" or getattr(getattr(message, "from_user", None), "is_bot", False):
        return
    matches = await detect_external(message)
    if not matches:
        return
    deleted = False
    if policy["mode"] == "delete" and not already_deleted:
        from wbb.services.blacklist_enforcement import _bot_capabilities, _is_privileged_member

        user = getattr(message, "from_user", None)
        if user and not await _is_privileged_member(chat_id, int(user.id)):
            capabilities = await _bot_capabilities(chat_id)
            if capabilities["delete"] is True:
                try:
                    await message.delete()
                    deleted = True
                except Exception as exc:
                    log.error(f"外部风控删除消息失败 chat_id={chat_id}: {exc}")
    await eventsdb.insert_one({
        "bot_id": BOT_PROFILE_ID, "chat_id": chat_id,
        "message_id": int(message.id),
        "user_id": int(message.from_user.id) if getattr(message, "from_user", None) else None,
        "matches": matches, "mode": policy["mode"], "deleted": deleted,
        "created_at": datetime.now(UTC),
    })