from __future__ import annotations import re import secrets import unicodedata from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from datetime import UTC, datetime, timedelta from typing import Any from pyrogram.enums import ChatMemberStatus, MessageEntityType from pyrogram.types import ChatPermissions from wbb import BOT_ID, SUDOERS, app, log from wbb.utils.dbadmin import ( get_managed_chat_settings, update_managed_chat_settings, ) from wbb.utils.dbfunctions import ( add_warn, delete_blacklist_filter, get_blacklisted_words, get_warn, int_to_alpha, remove_warns, save_blacklist_filter, ) RISK_ACTIONS = {"delete", "warn", "mute", "kick", "ban"} MEMBER_ACTIONS = {"warn", "mute", "kick", "ban"} MEMBER_ACTION_PRIORITY = ("ban", "kick", "mute", "warn") WARNING_LIMIT = 3 DEFAULT_BAN_DURATION_SECONDS = 3600 MIN_BAN_DURATION_SECONDS = 60 MAX_BAN_DURATION_SECONDS = 365 * 24 * 60 * 60 MAX_RISK_RULES = 50 LEGACY_POLICY_RULE_ID = "legacy-risk-control" LEGACY_KEYWORD_RULE_ID = "legacy-blacklist" COMMAND_KEYWORD_RULE_ID = "command-blacklist" LEGACY_RULE_IDS = { LEGACY_POLICY_RULE_ID, LEGACY_KEYWORD_RULE_ID, COMMAND_KEYWORD_RULE_ID, } RULE_ID_PATTERN = re.compile(r"^[A-Za-z0-9_-]{1,40}$") LINK_PATTERN = re.compile( r"(?:https?://|ftp://|www\.|t\.me/|telegram\.me/)\S+", flags=re.IGNORECASE, ) class RiskControlValidationError(ValueError): pass @dataclass(frozen=True) class RiskMatch: rule_id: str rule_name: str trigger_types: tuple[str, ...] keyword: str | None = None @dataclass(frozen=True) class RiskEnforcementResult: matches: tuple[RiskMatch, ...] deleted: bool | None requested_member_action: str applied_member_action: str warning_count: int | None = None duration_seconds: int | None = None @property def match(self) -> RiskMatch: return self.matches[0] def _normalize(value: str) -> str: return unicodedata.normalize("NFKC", value).casefold().strip() def _normalized_keywords(words: Iterable[Any]) -> list[str]: values = {_normalize(str(word))[:100] for word in words} return sorted(value for value in values if value)[:200] def _normalize_actions(value: Any, *, strict: bool) -> list[str]: raw_actions = value if value is not None else ["delete"] if isinstance(raw_actions, str): raw_actions = [raw_actions] if not isinstance(raw_actions, (list, tuple, set)): if strict: raise RiskControlValidationError("风控操作格式不正确。") raw_actions = ["delete"] requested_actions = [str(action).strip().lower() for action in raw_actions] invalid_actions = sorted(set(requested_actions) - RISK_ACTIONS) if invalid_actions and strict: raise RiskControlValidationError("包含不支持的风控操作。") member_actions = [ action for action in MEMBER_ACTION_PRIORITY if action in requested_actions ] if len(member_actions) > 1 and strict: raise RiskControlValidationError("警告、禁言、踢出和封禁只能选择一种。") actions: list[str] = [] if "delete" in requested_actions: actions.append("delete") if member_actions: actions.append(member_actions[0]) return actions def _normalize_action_duration(value: Any, *, strict: bool) -> int: if value is None or value == "": return DEFAULT_BAN_DURATION_SECONDS try: duration = int(value) except (TypeError, ValueError): if strict: raise RiskControlValidationError("处罚时长必须是整数秒。") from None return DEFAULT_BAN_DURATION_SECONDS if MIN_BAN_DURATION_SECONDS <= duration <= MAX_BAN_DURATION_SECONDS: return duration if strict: raise RiskControlValidationError("处罚时长需要在 1 分钟到 365 天之间。") return DEFAULT_BAN_DURATION_SECONDS def _normalize_rule( value: Mapping[str, Any], *, fallback_keywords: Iterable[Any] = (), strict: bool = False, default_rule_id: str | None = None, default_name: str | None = None, ) -> dict[str, Any]: raw_rule_id = str(value.get("rule_id") or default_rule_id or "").strip() if raw_rule_id and not RULE_ID_PATTERN.fullmatch(raw_rule_id): if strict: raise RiskControlValidationError("规则 ID 格式不正确。") raw_rule_id = "" rule_id = raw_rule_id or secrets.token_hex(8) name = str(value.get("name") or default_name or "").strip() if not name: if strict: raise RiskControlValidationError("规则名称不能为空。") name = "未命名规则" if len(name) > 60: if strict: raise RiskControlValidationError("规则名称不能超过 60 个字符。") name = name[:60] keywords = _normalized_keywords( value["keywords"] if "keywords" in value else fallback_keywords ) match_images = bool(value.get("match_images", False)) match_links = bool(value.get("match_links", False)) actions = _normalize_actions(value.get("actions"), strict=strict) enabled = bool( value.get("enabled", bool(keywords or match_images or match_links)) ) if strict and enabled and not (keywords or match_images or match_links): raise RiskControlValidationError("启用规则时至少配置一种触发条件。") if strict and enabled and not actions: raise RiskControlValidationError("启用规则时至少选择一种处置操作。") return { "rule_id": rule_id, "name": name, "enabled": enabled, "keywords": keywords, "match_images": match_images, "match_links": match_links, "actions": actions, "duration_seconds": _normalize_action_duration( value.get("duration_seconds", value.get("ban_duration_seconds")), strict=strict, ), } def normalize_risk_control( value: Mapping[str, Any] | None, *, fallback_keywords: Iterable[Any] = (), strict: bool = False, ) -> dict[str, Any]: raw = value or {} rule = _normalize_rule( raw, fallback_keywords=fallback_keywords, strict=strict, default_rule_id=LEGACY_POLICY_RULE_ID, default_name="原风控策略", ) return { key: rule[key] for key in ( "enabled", "keywords", "match_images", "match_links", "actions", "duration_seconds", ) } def normalize_risk_rules( value: Sequence[Mapping[str, Any]] | None, *, legacy_policy: Mapping[str, Any] | None = None, fallback_keywords: Iterable[Any] = (), strict: bool = False, ) -> list[dict[str, Any]]: fallback = _normalized_keywords(fallback_keywords) if value is None: if legacy_policy is not None: legacy = dict(legacy_policy) legacy["keywords"] = _normalized_keywords( [*(legacy.get("keywords") or []), *fallback] ) return [ _normalize_rule( legacy, strict=strict, default_rule_id=LEGACY_POLICY_RULE_ID, default_name="原风控策略", ) ] if fallback: return [ _normalize_rule( { "enabled": True, "keywords": fallback, "actions": ["delete"], }, strict=strict, default_rule_id=LEGACY_KEYWORD_RULE_ID, default_name="原关键词黑名单", ) ] return [] if isinstance(value, (str, bytes)) or not isinstance(value, Sequence): if strict: raise RiskControlValidationError("风控规则必须是列表。") return [] if len(value) > MAX_RISK_RULES and strict: raise RiskControlValidationError(f"每个群最多配置 {MAX_RISK_RULES} 条风控规则。") rules: list[dict[str, Any]] = [] rule_ids: set[str] = set() for index, item in enumerate(value[:MAX_RISK_RULES]): if not isinstance(item, Mapping): if strict: raise RiskControlValidationError("风控规则格式不正确。") continue rule = _normalize_rule( item, strict=strict, default_name=f"规则 {index + 1}", ) if rule["rule_id"] in rule_ids: if strict: raise RiskControlValidationError("规则 ID 不能重复。") rule["rule_id"] = f"{rule['rule_id'][:30]}-{index + 1}" rule_ids.add(rule["rule_id"]) rules.append(rule) return rules def risk_rule_keywords( rules: Iterable[Mapping[str, Any]], *, enabled_only: bool = False ) -> list[str]: return _normalized_keywords( keyword for rule in rules if not enabled_only or bool(rule.get("enabled")) for keyword in rule.get("keywords", []) ) def find_blacklist_match(content: str, words: Iterable[str]) -> str | None: normalized_content = _normalize(content) if not normalized_content: return None for word in words: normalized_word = _normalize(str(word)) if normalized_word and normalized_word in normalized_content: return str(word).strip() return None def _is_link_entity(entity: Any) -> bool: entity_type = getattr(entity, "type", None) if entity_type in {MessageEntityType.URL, MessageEntityType.TEXT_LINK}: return True label = str(entity_type).rsplit(".", 1)[-1].lower() return label in {"url", "text_link"} def _message_has_link(message: Any, content: str) -> bool: for name in ("entities", "caption_entities"): if any(_is_link_entity(entity) for entity in (getattr(message, name, None) or [])): return True return bool(LINK_PATTERN.search(content)) def _message_has_image(message: Any) -> bool: if getattr(message, "photo", None) or getattr(message, "animation", None): return True document = getattr(message, "document", None) return bool( document and str(getattr(document, "mime_type", "")).lower().startswith("image/") ) def detect_risk_match( message: Any, policy: Mapping[str, Any] ) -> RiskMatch | None: if not policy.get("enabled"): return None content = str( getattr(message, "text", None) or getattr(message, "caption", None) or "" ) keyword = find_blacklist_match(content, policy.get("keywords", [])) triggers: list[str] = [] if keyword: triggers.append("keyword") if policy.get("match_images") and _message_has_image(message): triggers.append("image") if policy.get("match_links") and _message_has_link(message, content): triggers.append("link") if not triggers: return None return RiskMatch( rule_id=str(policy.get("rule_id") or LEGACY_POLICY_RULE_ID), rule_name=str(policy.get("name") or "风控规则"), trigger_types=tuple(triggers), keyword=keyword, ) def detect_risk_matches( message: Any, rules: Iterable[Mapping[str, Any]] ) -> list[tuple[Mapping[str, Any], RiskMatch]]: matches: list[tuple[Mapping[str, Any], RiskMatch]] = [] for rule in rules: match = detect_risk_match(message, rule) if match: matches.append((rule, match)) return matches async def get_effective_risk_rules(chat_id: int) -> list[dict[str, Any]]: settings = await get_managed_chat_settings(chat_id) legacy_words = await get_blacklisted_words(chat_id) if "risk_rules" in settings: return normalize_risk_rules(settings.get("risk_rules")) return normalize_risk_rules( None, legacy_policy=settings.get("risk_control"), fallback_keywords=legacy_words or settings.get("blacklist_words", []), ) async def get_effective_risk_control(chat_id: int) -> dict[str, Any]: rules = await get_effective_risk_rules(chat_id) active = [rule for rule in rules if rule["enabled"]] actions = {action for rule in active for action in rule["actions"]} member_action = next( (action for action in MEMBER_ACTION_PRIORITY if action in actions), None ) normalized_actions = ["delete"] if "delete" in actions else [] if member_action: normalized_actions.append(member_action) return { "enabled": bool(active), "keywords": risk_rule_keywords(active), "match_images": any(rule["match_images"] for rule in active), "match_links": any(rule["match_links"] for rule in active), "actions": normalized_actions, "duration_seconds": max( (int(rule["duration_seconds"]) for rule in active), default=DEFAULT_BAN_DURATION_SECONDS, ), } async def sync_legacy_blacklist(chat_id: int, keywords: Iterable[Any]) -> list[str]: normalized = _normalized_keywords(keywords) current = await get_blacklisted_words(chat_id) current_by_normalized = {_normalize(str(word)): str(word) for word in current} target = set(normalized) for normalized_word, stored_word in current_by_normalized.items(): if normalized_word not in target: while await delete_blacklist_filter(chat_id, stored_word): pass current_normalized = { _normalize(str(word)) for word in await get_blacklisted_words(chat_id) } for word in normalized: if word not in current_normalized: await save_blacklist_filter(chat_id, word) return normalized def _keyword_rule_index(rules: list[dict[str, Any]]) -> int | None: for index, rule in enumerate(rules): if rule["rule_id"] in LEGACY_RULE_IDS: return index return None def replace_legacy_rule_keywords( rules: Sequence[Mapping[str, Any]], keywords: Iterable[Any] ) -> list[dict[str, Any]]: normalized_rules = normalize_risk_rules(rules) index = _keyword_rule_index(normalized_rules) if index is None: normalized_rules.append( _normalize_rule( { "enabled": True, "keywords": [], "actions": ["delete"], }, default_rule_id=COMMAND_KEYWORD_RULE_ID, default_name="命令关键词黑名单", ) ) index = len(normalized_rules) - 1 rule = normalized_rules[index] rule["keywords"] = _normalized_keywords(keywords) if rule["keywords"]: rule["enabled"] = True elif not rule["match_images"] and not rule["match_links"]: rule["enabled"] = False return normalized_rules async def _store_risk_rules(chat_id: int, rules: list[dict[str, Any]]) -> None: keywords = await sync_legacy_blacklist( chat_id, risk_rule_keywords(rules, enabled_only=True) ) await update_managed_chat_settings( chat_id, {"risk_rules": rules, "blacklist_words": keywords}, ) async def add_risk_keyword(chat_id: int, word: str) -> tuple[bool, dict[str, Any]]: normalized_word = _normalize(word)[:100] if not normalized_word: raise RiskControlValidationError("关键词不能为空。") rules = await get_effective_risk_rules(chat_id) index = _keyword_rule_index(rules) existing_keywords = rules[index]["keywords"] if index is not None else [] rules = replace_legacy_rule_keywords( rules, [*existing_keywords, normalized_word] ) index = _keyword_rule_index(rules) assert index is not None rule = rules[index] changed = normalized_word not in existing_keywords await _store_risk_rules(chat_id, rules) return changed, rule async def remove_risk_keyword(chat_id: int, word: str) -> tuple[bool, dict[str, Any]]: normalized_word = _normalize(word)[:100] rules = await get_effective_risk_rules(chat_id) index = _keyword_rule_index(rules) if index is None: return False, normalize_risk_control(None) existing_keywords = rules[index]["keywords"] filtered = [ item for item in existing_keywords if _normalize(item) != normalized_word ] changed = len(filtered) != len(existing_keywords) rules = replace_legacy_rule_keywords(rules, filtered) index = _keyword_rule_index(rules) assert index is not None rule = rules[index] await _store_risk_rules(chat_id, rules) return changed, rule async def _is_privileged_member(chat_id: int, user_id: int) -> bool: if user_id in SUDOERS or user_id == BOT_ID: return True try: member = await app.get_chat_member(chat_id, user_id) except Exception as exc: log.error( f"无法确认群 {chat_id} 成员 {user_id} 的管理员身份，已跳过自动处罚：{exc}" ) return True return member.status in { ChatMemberStatus.OWNER, ChatMemberStatus.ADMINISTRATOR, } async def _bot_capabilities(chat_id: int) -> dict[str, bool | None]: try: member = await app.get_chat_member(chat_id, BOT_ID) except Exception as exc: log.error(f"无法读取机器人在群 {chat_id} 的实时权限，将直接尝试执行：{exc}") return {"delete": None, "restrict": None} if member.status == ChatMemberStatus.OWNER: return {"delete": True, "restrict": True} if member.status != ChatMemberStatus.ADMINISTRATOR: return {"delete": False, "restrict": False} privileges = getattr(member, "privileges", None) return { "delete": bool(getattr(privileges, "can_delete_messages", False)), "restrict": bool(getattr(privileges, "can_restrict_members", False)), } def _ban_until(duration_seconds: int) -> datetime: return datetime.now(UTC) + timedelta(seconds=duration_seconds) async def _apply_warning( chat_id: int, user_id: int, *, can_restrict: bool | None, duration_seconds: int, ) -> tuple[str, int]: key = await int_to_alpha(user_id) current = await get_warn(chat_id, key) count = int(current.get("warns", 0)) + 1 if current else 1 if count < WARNING_LIMIT: await add_warn(chat_id, key, {"warns": count}) return "warn", count if can_restrict is False: await add_warn(chat_id, key, {"warns": count}) log.error( f"群 {chat_id} 成员 {user_id} 已达到警告上限，但机器人没有封禁权限" ) return "warn", count try: await app.ban_chat_member( chat_id, user_id, until_date=_ban_until(duration_seconds), ) except Exception as exc: await add_warn(chat_id, key, {"warns": count}) log.error(f"群 {chat_id} 成员 {user_id} 已达到警告上限，但封禁失败：{exc}") return "warn", count await remove_warns(chat_id, key) return "ban", count async def _apply_member_action( chat_id: int, user_id: int, action: str, *, can_restrict: bool | None, duration_seconds: int, ) -> tuple[str, int | None]: if action == "warn": return await _apply_warning( chat_id, user_id, can_restrict=can_restrict, duration_seconds=duration_seconds, ) if can_restrict is False: log.error(f"群 {chat_id} 无法执行风控{action}：机器人没有限制成员权限") return "none", None try: if action == "kick": await app.ban_chat_member(chat_id, user_id) await app.unban_chat_member(chat_id, user_id) elif action == "mute": await app.restrict_chat_member( chat_id, user_id, ChatPermissions(), until_date=_ban_until(duration_seconds), ) elif action == "ban": await app.ban_chat_member( chat_id, user_id, until_date=_ban_until(duration_seconds), ) else: return "none", None except Exception as exc: log.error(f"群 {chat_id} 风控处置成员 {user_id} 失败：{exc}") return "none", None return action, None def _trigger_description(match: RiskMatch) -> str: labels = {"keyword": "关键词", "image": "图片", "link": "链接"} values = [labels[item] for item in match.trigger_types] if match.keyword: values[values.index("关键词")] = f"关键词“{match.keyword}”" return "、".join(values) def _duration_description(duration_seconds: int) -> str: if duration_seconds % 86400 == 0: return f"{duration_seconds // 86400} 天" if duration_seconds % 3600 == 0: return f"{duration_seconds // 3600} 小时" return f"{max(1, duration_seconds // 60)} 分钟" async def _send_action_notice( chat_id: int, user: Any, matches: tuple[RiskMatch, ...], action: str, warning_count: int | None, duration_seconds: int, ) -> None: mention = getattr(user, "mention", None) or f"成员 {user.id}" if action == "warn": outcome = f"已警告（{warning_count}/{WARNING_LIMIT}）" elif action == "ban" and warning_count: outcome = ( f"累计 {WARNING_LIMIT} 次警告，已封禁 " f"{_duration_description(duration_seconds)}" ) elif action == "ban": outcome = f"已封禁 {_duration_description(duration_seconds)}" elif action == "mute": outcome = f"已禁言 {_duration_description(duration_seconds)}" else: outcome = {"kick": "已踢出"}.get(action, action) matched = "；".join( f"{match.rule_name}：{_trigger_description(match)}" for match in matches ) try: await app.send_message( chat_id, f"{mention} 触发风控规则（{matched}），{outcome}。", ) except Exception as exc: log.error(f"群 {chat_id} 风控处置通知发送失败：{exc}") async def enforce_risk_message( message: Any, policy: Mapping[str, Any] | Sequence[Mapping[str, Any]] | None = None, ) -> RiskEnforcementResult | None: user = getattr(message, "from_user", None) if user is not None and getattr(user, "is_bot", False): return None chat_id = int(message.chat.id) if policy is None: rules = await get_effective_risk_rules(chat_id) elif isinstance(policy, Mapping): rules = normalize_risk_rules(None, legacy_policy=policy) else: rules = normalize_risk_rules(policy) matched_rules = detect_risk_matches(message, rules) if not matched_rules: return None matches = tuple(match for _, match in matched_rules) delete_requested = any("delete" in rule["actions"] for rule, _ in matched_rules) member_action = "none" action_rules: list[Mapping[str, Any]] = [] for candidate in MEMBER_ACTION_PRIORITY: action_rules = [ rule for rule, _ in matched_rules if candidate in rule["actions"] ] if action_rules: member_action = candidate break duration_seconds = max( (int(rule["duration_seconds"]) for rule in action_rules), default=DEFAULT_BAN_DURATION_SECONDS, ) capabilities = await _bot_capabilities(chat_id) deleted: bool | None = None if delete_requested: deleted = False if capabilities["delete"] is False: log.error(f"群 {chat_id} 命中风控规则，但机器人没有删除消息权限") else: try: await message.delete() deleted = True except Exception as exc: log.error(f"群 {chat_id} 命中风控规则但删除消息失败：{exc}") applied_action = "none" warning_count: int | None = None if ( member_action != "none" and user is not None and not await _is_privileged_member(chat_id, int(user.id)) ): applied_action, warning_count = await _apply_member_action( chat_id, int(user.id), member_action, can_restrict=capabilities["restrict"], duration_seconds=duration_seconds, ) if applied_action != "none": await _send_action_notice( chat_id, user, matches, applied_action, warning_count, duration_seconds, ) return RiskEnforcementResult( matches=matches, deleted=deleted, requested_member_action=member_action, applied_member_action=applied_action, warning_count=warning_count, duration_seconds=( duration_seconds if member_action in {"ban", "mute", "warn"} else None ), ) async def enforce_blacklist_message( message: Any, words: Iterable[str] ) -> RiskEnforcementResult | None: keyword_list = list(words) return await enforce_risk_message( message, [ { "rule_id": LEGACY_KEYWORD_RULE_ID, "name": "关键词黑名单", "enabled": bool(keyword_list), "keywords": keyword_list, "match_images": False, "match_links": False, "actions": ["delete"], "duration_seconds": DEFAULT_BAN_DURATION_SECONDS, } ], )