#!/usr/bin/env python3
"""
安全控制模块
负责权限检查、内容过滤、速率限制
"""

import re
import time
import logging
from typing import Dict, List, Set, Tuple
from collections import defaultdict

logger = logging.getLogger(__name__)


class SecurityChecker:
    """安全检查器"""
    
    def __init__(self, config: Dict):
        self.config = config
        
        # 允许的工具
        self.allowed_tools: Set[str] = set(
            config.get('allowed_tools', ['read', 'write', 'edit', 'exec'])
        )
        
        # 禁止的命令
        self.blocked_commands: List[str] = config.get(
            'blocked_commands',
            ['rm -rf /', 'dd if=', 'mkfs', 'shutdown', 'reboot', 'halt']
        )
        
        # 敏感关键词
        self.sensitive_keywords: List[str] = [
            'password', 'secret', 'token', 'api_key',
            '私钥', '密码', '密钥', '.ssh/id_rsa'
        ]
        
        # 速率限制
        self.max_requests = config.get('max_requests_per_minute', 10)
        self.request_records: Dict[str, List[float]] = defaultdict(list)
    
    def check_rate_limit(self, user_id: str) -> bool:
        """
        检查速率限制
        
        Args:
            user_id: 用户标识
            
        Returns:
            是否允许请求
        """
        now = time.time()
        window = 60  # 60秒窗口
        
        # 清理过期记录
        self.request_records[user_id] = [
            t for t in self.request_records[user_id]
            if now - t < window
        ]
        
        # 检查是否超过限制
        if len(self.request_records[user_id]) >= self.max_requests:
            logger.warning(f"用户 {user_id} 触发速率限制")
            return False
        
        # 记录本次请求
        self.request_records[user_id].append(now)
        return True
    
    def check_content(self, content: str) -> bool:
        """
        检查内容是否合法
        
        Args:
            content: 消息内容
            
        Returns:
            是否合法
        """
        # 检查敏感关键词
        content_lower = content.lower()
        for keyword in self.sensitive_keywords:
            if keyword.lower() in content_lower:
                logger.warning(f"内容包含敏感词: {keyword}")
                # 不直接拒绝，只记录警告
                # 实际使用时可根据需要调整策略
        
        # 检查是否有明显的恶意命令
        for cmd in self.blocked_commands:
            if cmd in content:
                logger.warning(f"检测到禁止命令: {cmd}")
                return False
        
        return True
    
    def check_tool_permission(self, tool_name: str) -> bool:
        """
        检查工具使用权限
        
        Args:
            tool_name: 工具名称
            
        Returns:
            是否允许使用
        """
        return tool_name in self.allowed_tools
    
    def check_command_safety(self, command: str) -> Tuple[bool, str]:
        """
        检查命令是否安全
        
        Args:
            command: 命令字符串
            
        Returns:
            (是否安全, 原因)
        """
        command_lower = command.lower().strip()
        
        # 检查危险命令模式
        dangerous_patterns = [
            (r'^\s*rm\s+-rf\s*/', "禁止删除根目录"),
            (r'^\s*dd\s+if=', "禁止直接使用 dd 命令"),
            (r'^\s*mkfs', "禁止格式化文件系统"),
            (r'^\s*shutdown', "禁止关机命令"),
            (r'^\s*reboot', "禁止重启命令"),
            (r'^\s*halt', "禁止 halt 命令"),
            (r'>\s*/dev/sd[a-z]', "禁止直接操作磁盘设备"),
            (r'>\s*/dev/null.*<\s*/dev/random', "禁止危险的重定向操作"),
        ]
        
        for pattern, reason in dangerous_patterns:
            if re.search(pattern, command_lower):
                return False, reason
        
        # 检查敏感文件访问
        sensitive_paths = [
            '/etc/passwd',
            '/etc/shadow',
            '~/.ssh/id_rsa',
            '~/.aws/credentials',
        ]
        
        for path in sensitive_paths:
            if path in command:
                return False, f"禁止访问敏感文件: {path}"
        
        return True, "安全"
    
    def sanitize_output(self, output: str) -> str:
        """
        清理输出内容
        
        Args:
            output: 原始输出
            
        Returns:
            清理后的输出
        """
        # 移除可能的敏感信息
        patterns = [
            (r'[a-zA-Z0-9_-]{20,}\.[a-zA-Z0-9_-]{20,}', '[TOKEN_REDACTED]'),
            (r'sk-[a-zA-Z0-9]{20,}', '[API_KEY_REDACTED]'),
            (r'password\s*[:=]\s*\S+', 'password: [REDACTED]'),
        ]
        
        result = output
        for pattern, replacement in patterns:
            result = re.sub(pattern, replacement, result, flags=re.IGNORECASE)
        
        return result
    
    def get_blocked_commands_list(self) -> List[str]:
        """获取禁止命令列表"""
        return self.blocked_commands.copy()


class SimpleSecurityChecker:
    """简化版安全检查器（用于测试）"""
    
    def __init__(self, config: Dict):
        pass
    
    def check_rate_limit(self, user_id: str) -> bool:
        return True
    
    def check_content(self, content: str) -> bool:
        return True
    
    def check_tool_permission(self, tool_name: str) -> bool:
        return True
    
    def check_command_safety(self, command: str) -> Tuple[bool, str]:
        return True, "安全"
    
    def sanitize_output(self, output: str) -> str:
        return output