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