ForcePilot/backend/package/yuxi/channels/adapters/qqbot/security.py

from __future__ import annotations

import logging

from yuxi.channels.models import ChannelMessage
from yuxi.channels.policy.security_policy import (
    AccessResult,
    BaseSecurityPolicy,
    DmPolicy,
    GroupPolicy,
    RejectReason,
    WildcardAllowlistMatcher,
)

logger = logging.getLogger(__name__)


class QQBotSecurityPolicy(BaseSecurityPolicy):
    def __init__(self, config: dict):
        normalized = self._normalize_config_ids(config)
        super().__init__(normalized)
        self._paired_users_matcher = WildcardAllowlistMatcher.from_config(
            self._format_ids(config.get("paired_users", []))
        )
        self._group_require_mention = config.get("group_require_mention", True)
        self._pairing_enabled = config.get("pairing_enabled", config.get("enable_pairing", False))

    @staticmethod
    def _normalize_config_ids(config: dict) -> dict:
        normalized = dict(config)
        for key in ("allowFrom", "allow_from"):
            if key in normalized:
                normalized[key] = QQBotSecurityPolicy._format_ids(normalized[key])
        for key in ("groupAllowFrom", "group_allow_from"):
            if key in normalized:
                normalized[key] = QQBotSecurityPolicy._format_ids(normalized[key])
        return normalized

    def check_dm_access(self, sender_id: str) -> AccessResult:
        if self._dm_policy == DmPolicy.DISABLED:
            return AccessResult(False, RejectReason.DM_DISABLED)
        if self._dm_policy == DmPolicy.OPEN:
            return AccessResult(True, RejectReason.DM_OPEN_PASS)
        if self._dm_policy == DmPolicy.PAIRING:
            if not self._pairing_enabled:
                return AccessResult(True, RejectReason.DM_OPEN_PASS)
            if self._paired_users_matcher.match(sender_id):
                return AccessResult(True, RejectReason.DM_ALLOWLISTED)
            return AccessResult(False, RejectReason.DM_PAIRING_REQUIRED, f"sender '{sender_id}' not paired")
        if self._dm_policy == DmPolicy.ALLOWLIST:
            if self._dm_matcher.match(sender_id):
                return AccessResult(True, RejectReason.DM_ALLOWLISTED)
            return AccessResult(False, RejectReason.DM_NOT_ALLOWLISTED, f"sender '{sender_id}' not in allowFrom")
        return AccessResult(True, None)

    def check_group_access(self, group_id: str) -> AccessResult:
        if self._group_policy == GroupPolicy.DISABLED:
            return AccessResult(False, RejectReason.GROUP_DISABLED)
        if self._group_policy == GroupPolicy.OPEN:
            return AccessResult(True, RejectReason.GROUP_OPEN_PASS)
        if self._group_policy == GroupPolicy.ALLOWLIST:
            if self._group_matcher.match(group_id):
                return AccessResult(True, RejectReason.GROUP_ALLOWLISTED)
            return AccessResult(False, RejectReason.GROUP_NOT_ALLOWLISTED, f"group '{group_id}' not in groupAllowFrom")
        return AccessResult(False, RejectReason.GROUP_DISABLED)

    def check_mention_required(
        self,
        chat_id: str,
        msg: ChannelMessage,
        bot_names: list[str] | None = None,
    ) -> bool:
        if not self._group_require_mention:
            return True

        groups_config = self._config.get("groups", {})
        chat_cfg = groups_config.get(chat_id, {})
        require_mention = chat_cfg.get("require_mention", True)

        if not require_mention:
            return True

        if msg.mentions and msg.mentions.is_bot_mentioned:
            return True

        content = msg.content or ""
        for name in bot_names or []:
            if f"@{name}" in content:
                return True

        return False

    def _resolve_sender_id(self, event_data: dict) -> str:
        author = event_data.get("author", {})
        return author.get("id", author.get("member_openid", ""))

    def _resolve_group_id(self, event_data: dict) -> str:
        return event_data.get("group_openid", event_data.get("group_id", event_data.get("guild_id", "")))

    @staticmethod
    def _format_ids(raw_ids: list[str]) -> list[str]:
        result: list[str] = []
        for raw in raw_ids:
            raw = raw.strip()
            if raw.startswith("qq:"):
                result.append(raw[3:])
            else:
                result.append(raw)
        return result


def verify_webhook_ed25519(headers: dict, body: bytes, bot_secret: str) -> bool:
    sig = headers.get("x-signature-ed25519", "")
    timestamp_str = headers.get("x-signature-timestamp", "")
    if not sig or not timestamp_str:
        return False

    try:
        from cryptography.exceptions import InvalidSignature
        from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

        if len(bot_secret) != 64:
            logger.warning("[QQBot] Invalid bot_secret length, expected 64 hex chars")
            return False
        seed = bytes.fromhex(bot_secret)
        private_key = Ed25519PrivateKey.from_private_bytes(seed)
        public_key = private_key.public_key()
        message = timestamp_str.encode() + body
        public_key.verify(bytes.fromhex(sig), message)
        return True
    except InvalidSignature:
        return False
    except Exception:
        logger.exception("[QQBot] Ed25519 verification error")
        return False


__all__ = [
    "QQBotSecurityPolicy",
    "verify_webhook_ed25519",
    "check_dm_policy",
    "check_group_policy",
    "check_mention_required",
    "AccessResult",
    "RejectReason",
    "BaseSecurityPolicy",
    "DmPolicy",
    "GroupPolicy",
    "WildcardAllowlistMatcher",
]


async def check_dm_policy(user_id: str, config: dict) -> bool:
    normalized_config = _normalize_config(config)
    policy = QQBotSecurityPolicy(normalized_config)
    clean_id = user_id.strip().removeprefix("qq:")
    result = policy.check_dm_access(clean_id)
    return result.allowed


def _normalize_config(config: dict) -> dict:
    normalized = dict(config)
    for key in ("allowFrom", "allow_from"):
        if key in normalized:
            normalized[key] = QQBotSecurityPolicy._format_ids(normalized[key])
    for key in ("groupAllowFrom", "group_allow_from"):
        if key in normalized:
            normalized[key] = QQBotSecurityPolicy._format_ids(normalized[key])
    return normalized


async def check_group_policy(chat_id: str, user_id: str, config: dict) -> bool:
    normalized_config = _normalize_config(config)
    policy = QQBotSecurityPolicy(normalized_config)
    group_id = chat_id.replace("group_", "").replace("dm_", "")
    result = policy.check_group_access(group_id)
    if not result.allowed:
        group_allow = config.get("group_allow_from", [])
        if f"qq:{user_id}" in group_allow:
            return True
        groups_config = config.get("groups", {})
        chat_cfg = groups_config.get(chat_id, {})
        per_group_allow = chat_cfg.get("allow_from", [])
        if f"qq:{user_id}" in per_group_allow:
            return True
    return result.allowed


async def check_mention_required(
    chat_id: str,
    msg,
    config: dict,
    bot_names: list[str] | None = None,
) -> bool:
    policy = QQBotSecurityPolicy(config)
    return policy.check_mention_required(chat_id, msg, bot_names)