security_module.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
安全模块 - 速率限制、恶意请求检测、输入过滤

功能亮点：
✅ 多维度速率限制（会话级、IP级、全局级、VIP级）
✅ 智能输入过滤（敏感词、恶意内容检测）
✅ 异常行为检测（高频请求、相同查询、攻击模式）
✅ 实时阻塞和降级策略
✅ 详细的监控和统计
✅ 可配置的安全策略

使用示例：
    from security_module import SecurityManager
    
    security = SecurityManager()
    
    # 检查请求是否允许
    allowed, reason = security.check_request("user123", "192.168.1.100", "你好")
    
    # 过滤输入
    filtered_text, is_safe = security.filter_input("包含敏感词的文本")
    
    # 获取统计信息
    stats = security.get_stats()
"""

import os
import re
import time
import json
import random
import string
import threading
import traceback
from typing import Dict, Tuple, List, Optional, Any, Set, Union
from collections import defaultdict, Counter
from datetime import datetime, timedelta
from enum import Enum
from dataclasses import dataclass, field
from logger import logger

# ======================
# 安全配置
# ======================

class SecurityLevel(Enum):
    """安全级别枚举"""
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"

class SecurityConfig:
    """安全配置类"""
    
    # 速率限制配置
    RATE_LIMITS = {
        "global": {  # 全局限制
            "requests_per_minute": 1000,
            "burst": 200
        },
        "ip": {  # IP级别限制
            "requests_per_minute": 60,
            "burst": 15
        },
        "session": {  # 会话级别限制
            "requests_per_minute": 30,
            "burst": 10
        },
        "vip": {  # VIP用户限制
            "requests_per_minute": 120,
            "burst": 30
        }
    }
    
    # 输入过滤配置
    MAX_INPUT_LENGTH = 1000  # 最大输入长度
    MIN_INPUT_LENGTH = 1     # 最小输入长度
    ALLOWED_CHARACTERS = set(string.printable + '，。、；：\'"【】《》！？……（）') - set('<>{}[]\\|`')
    
    # 敏感词列表
    SENSITIVE_WORDS = [
        # 脏话和不文明用语
        "fuck", "shit", "damn", "bitch", "asshole", "nigger", "操", "他妈的", "狗日的", "傻逼",
        "废物", "垃圾", "混蛋", "王八蛋", "畜生", "禽兽", "无耻", "下流", "卑鄙",
        # 攻击性词汇
        "terror", "bomb", "kill", "murder", "attack", "hack", "crack", "destroy", "攻击", "破坏",
        # 系统相关词汇（防止注入攻击）
        "admin", "root", "password", "database", "select", "drop", "delete", "union", "insert", "update",
        "truncate", "exec", "execute", "system", "cmd", "command", "shell", "bash",
        # 脚本相关
        "script", "alert", "onload", "onerror", "eval", "javascript", "iframe", "src=", "href=",
        "cookie", "localStorage", "sessionStorage", "document", "window", "location",
        # 恶意软件
        "malware", "virus", "trojan", "worm", "ransomware", "spyware", "keylogger", "backdoor",
        # 其他敏感词
        "赌博", "色情", "毒品", "枪支", "爆炸", "恐怖", "分裂", "攻击", "破坏", "诈骗", "传销",
        "政治", "宗教", "色情", "暴力", "血腥", "歧视", "仇恨", "极端"
    ]
    
    # 恶意模式（正则表达式）
    MALICIOUS_PATTERNS = [
        r'<script.*?>.*?</script>',      # XSS脚本
        r'on\w+\s*=\s*["\'].*?["\']',    # 事件处理器
        r'union\s+select',               # SQL注入
        r'select\s+.*?from',             # SQL注入
        r'1\s*=\s*1',                    # SQL注入
        r'\.\./',                        # 路径遍历
        r'\${.*?}',                     # 模板注入
        r'eval\s*\(.*?\)',               # 代码执行
        r'import\s+.*?os',               # 模块导入
        r'exec\s*\(.*?\)',               # 代码执行
        r'__import__',                   # Python动态导入
        r'rm\s+-rf',                     # Linux命令
        r'cmd\.exe',                     # Windows命令
        r'powershell',                   # PowerShell
        r'<iframe.*?>',                  # iframe嵌入
        r'javascript\s*:',               # JavaScript协议
        r'data\s*:',                     # Data协议
        r'vbscript\s*:',                 # VBScript协议
        r'expression\s*\(',              # CSS表达式
        r'@import',                      # CSS导入
        r'base64,',                      # Base64编码
        r'document\.cookie',             # Cookie窃取
        r'localStorage',                 # 本地存储访问
        r'sessionStorage',               # 会话存储访问
        r'window\.location',             # 位置重定向
        r'alert\s*\(.*?\)',              # 弹窗攻击
    ]
    
    # VIP用户列表（示例，可以在运行时动态添加）
    VIP_USERS = {"vip_user_1", "premium_customer", "admin_staff", "test_vip"}
    
    # 异常行为阈值
    ABNORMAL_THRESHOLD = {
        "same_query_count": 5,          # 相同查询次数阈值
        "short_interval_requests": 10,   # 短时间请求次数阈值
        "min_interval_seconds": 0.5,     # 最小请求间隔（秒）
        "suspicious_pattern_count": 3    # 可疑模式出现次数阈值
    }
    
    # 阻塞时间配置（秒）
    BLOCK_DURATION = {
        "low": 30,      # 低风险阻塞30秒
        "medium": 300,  # 中风险阻塞5分钟
        "high": 3600    # 高风险阻塞1小时
    }
    
    # 可信IP列表（不会被阻塞）
    TRUSTED_IPS = {"127.0.0.1", "localhost", "::1"}

SECURITY_CONFIG = SecurityConfig()

# ======================
# 速率限制器
# ======================

@dataclass
class RateLimitWindow:
    """速率限制窗口"""
    count: int = 0
    start_time: float = field(default_factory=time.time)
    last_request_time: float = field(default_factory=time.time)

class RateLimitResult:
    """速率限制结果"""
    
    def __init__(self, allowed: bool, remaining: int, reset_after: float, limit: int, severity: SecurityLevel = SecurityLevel.LOW):
        self.allowed = allowed
        self.remaining = remaining
        self.reset_after = reset_after
        self.limit = limit
        self.severity = severity
    
    def to_dict(self) -> dict:
        return {
            "allowed": self.allowed,
            "remaining": self.remaining,
            "reset_after": round(self.reset_after, 2),
            "limit": self.limit,
            "severity": self.severity.value
        }

class RateLimiter:
    """多维度速率限制器"""
    
    def __init__(self, config: SecurityConfig = SECURITY_CONFIG):
        self.config = config
        self._lock = threading.RLock()
        
        # 限制窗口数据
        self.windows = {
            "global": defaultdict(lambda: RateLimitWindow()),
            "ip": defaultdict(lambda: RateLimitWindow()),
            "session": defaultdict(lambda: RateLimitWindow()),
            "vip": defaultdict(lambda: RateLimitWindow())
        }
        
        # 请求历史记录（用于异常检测）
        self.request_history = defaultdict(list)  # session_id -> [(timestamp, query, ip)]
        
        # 被限制的实体
        self.blocked_entities = {
            "ips": {},      # ip: (block_until, severity)
            "sessions": {}, # session_id: (block_until, severity)
            "users": {}     # user_id: (block_until, severity)
        }
        
        # 统计数据
        self.stats = {
            "total_requests": 0,
            "blocked_requests": 0,
            "rate_limited_requests": 0,
            "filtered_inputs": 0,
            "abnormal_detected": 0,
            "malicious_detected": 0
        }
        
        # 清理线程
        self._cleanup_thread = threading.Thread(target=self._cleanup_loop, daemon=True)
        self._cleanup_thread.start()
        
        logger.info("🔒 速率限制器初始化完成")
    
    def _cleanup_loop(self):
        """定期清理过期数据"""
        while True:
            try:
                time.sleep(60)  # 每分钟清理一次
                
                with self._lock:
                    current_time = time.time()
                    
                    # 清理过期的阻塞记录
                    for entity_type in ["ips", "sessions", "users"]:
                        expired_keys = [
                            key for key, (block_until, _) in self.blocked_entities[entity_type].items()
                            if current_time > block_until
                        ]
                        for key in expired_keys:
                            del self.blocked_entities[entity_type][key]
                    
                    # 清理过期的请求历史
                    expired_sessions = []
                    for session_id, history in self.request_history.items():
                        # 保留最近10分钟的记录
                        self.request_history[session_id] = [
                            (t, q, ip) for t, q, ip in history if current_time - t < 600
                        ]
                        if not self.request_history[session_id]:
                            expired_sessions.append(session_id)
                    
                    for session_id in expired_sessions:
                        del self.request_history[session_id]
                    
                    # 清理过期的时间窗口
                    for window_type in self.windows:
                        expired_keys = []
                        for key, window in self.windows[window_type].items():
                            if current_time - window.start_time >= 300:  # 5分钟
                                expired_keys.append(key)
                        
                        for key in expired_keys:
                            del self.windows[window_type][key]
                    
            except Exception as e:
                logger.info(f"🚨 清理过程中出错: {e}")
                time.sleep(10)
    
    def _get_window(self, key: str, window_type: str) -> RateLimitWindow:
        """获取或创建时间窗口"""
        with self._lock:
            window = self.windows[window_type][key]
            current_time = time.time()
            
            # 检查窗口是否需要重置（1分钟窗口）
            if current_time - window.start_time >= 60:  # 1分钟
                window.count = 0
                window.start_time = current_time
            
            return window
    
    def _is_blocked(self, entity_type: str, entity_id: str) -> Tuple[bool, float, SecurityLevel]:
        """检查实体是否被临时阻塞"""
        with self._lock:
            blocked_dict = self.blocked_entities[entity_type + "s"]
            if entity_id in blocked_dict:
                block_until, severity = blocked_dict[entity_id]
                current_time = time.time()
                if current_time < block_until:
                    return True, block_until - current_time, severity
                else:
                    del blocked_dict[entity_id]
            return False, 0, SecurityLevel.LOW
    
    def _block_entity(self, entity_type: str, entity_id: str, duration: int, severity: SecurityLevel = SecurityLevel.MEDIUM):
        """阻塞实体一段时间（秒）"""
        with self._lock:
            # 可信实体不阻塞
            if entity_type == "ip" and entity_id in self.config.TRUSTED_IPS:
                return
            
            blocked_dict = self.blocked_entities[entity_type + "s"]
            block_until = time.time() + duration
            
            # 如果已有阻塞记录，取较长的阻塞时间
            if entity_id in blocked_dict:
                existing_until, existing_severity = blocked_dict[entity_id]
                if existing_until > block_until and existing_severity.value >= severity.value:
                    return  # 保留现有的更严格的阻塞
            
            blocked_dict[entity_id] = (block_until, severity)
    
    def _check_abnormal_behavior(self, session_id: str, query: str, ip_address: str) -> Tuple[bool, str, SecurityLevel]:
        """检查异常行为模式"""
        with self._lock:
            history = self.request_history[session_id]
            current_time = time.time()
            
            # 1. 检查相同查询频率
            same_query_count = sum(1 for t, q, ip in history if q == query and current_time - t < 10)
            if same_query_count >= self.config.ABNORMAL_THRESHOLD["same_query_count"]:
                return True, f"相同查询频繁出现 ({same_query_count}次/10秒)", SecurityLevel.MEDIUM
            
            # 2. 检查短时间请求频率
            recent_requests = sum(1 for t, q, ip in history if current_time - t < self.config.ABNORMAL_THRESHOLD["min_interval_seconds"])
            if recent_requests >= self.config.ABNORMAL_THRESHOLD["short_interval_requests"]:
                return True, f"短时间高频请求 ({recent_requests}次/{self.config.ABNORMAL_THRESHOLD['min_interval_seconds']}秒)", SecurityLevel.HIGH
            
            # 3. 检查相同IP的多个会话
            same_ip_sessions = defaultdict(int)
            for sess_id, hist in self.request_history.items():
                if sess_id == session_id:
                    continue
                for t, q, ip in hist:
                    if ip == ip_address and current_time - t < 60:
                        same_ip_sessions[sess_id] += 1
            
            if len(same_ip_sessions) > 5:  # 同一个IP有超过5个活跃会话
                return True, f"单IP多会话攻击 (IP: {ip_address}, 会话数: {len(same_ip_sessions)})", SecurityLevel.HIGH
            
            return False, "", SecurityLevel.LOW
    
    def check_rate_limit(self, session_id: str, ip_address: str = None, user_id: str = None) -> RateLimitResult:
        """检查速率限制"""
        with self._lock:
            self.stats["total_requests"] += 1
            
            # 0. 检查是否被阻塞
            blocked, remaining_time, severity = self._is_blocked("session", session_id)
            if blocked:
                self.stats["blocked_requests"] += 1
                return RateLimitResult(
                    allowed=False,
                    remaining=0,
                    reset_after=remaining_time,
                    limit=0,
                    severity=severity
                )
            
            if ip_address:
                blocked, remaining_time, severity = self._is_blocked("ip", ip_address)
                if blocked:
                    self.stats["blocked_requests"] += 1
                    return RateLimitResult(
                        allowed=False,
                        remaining=0,
                        reset_after=remaining_time,
                        limit=0,
                        severity=severity
                    )
            
            if user_id:
                blocked, remaining_time, severity = self._is_blocked("user", user_id)
                if blocked:
                    self.stats["blocked_requests"] += 1
                    return RateLimitResult(
                        allowed=False,
                        remaining=0,
                        reset_after=remaining_time,
                        limit=0,
                        severity=severity
                    )
            
            current_time = time.time()
            
            # 1. 全局限制
            global_window = self._get_window("global", "global")
            global_limit = self.config.RATE_LIMITS["global"]["requests_per_minute"]
            global_burst = self.config.RATE_LIMITS["global"]["burst"]
            
            if global_window.count >= global_limit + global_burst:
                self.stats["rate_limited_requests"] += 1
                return RateLimitResult(
                    allowed=False,
                    remaining=0,
                    reset_after=60 - (current_time - global_window.start_time),
                    limit=global_limit,
                    severity=SecurityLevel.MEDIUM
                )
            
            # 2. IP限制
            if ip_address:
                ip_window = self._get_window(ip_address, "ip")
                ip_limit = self.config.RATE_LIMITS["ip"]["requests_per_minute"]
                ip_burst = self.config.RATE_LIMITS["ip"]["burst"]
                
                if ip_window.count >= ip_limit + ip_burst:
                    self.stats["rate_limited_requests"] += 1
                    return RateLimitResult(
                        allowed=False,
                        remaining=0,
                        reset_after=60 - (current_time - ip_window.start_time),
                        limit=ip_limit,
                        severity=SecurityLevel.MEDIUM
                    )
            
            # 3. 会话限制（检查是否VIP）
            is_vip = session_id in self.config.VIP_USERS or (user_id and user_id in self.config.VIP_USERS)
            limit_type = "vip" if is_vip else "session"
            session_window = self._get_window(session_id, limit_type)
            session_limit = self.config.RATE_LIMITS[limit_type]["requests_per_minute"]
            session_burst = self.config.RATE_LIMITS[limit_type]["burst"]
            
            if session_window.count >= session_limit + session_burst:
                self.stats["rate_limited_requests"] += 1
                return RateLimitResult(
                    allowed=False,
                    remaining=0,
                    reset_after=60 - (current_time - session_window.start_time),
                    limit=session_limit,
                    severity=SecurityLevel.LOW if is_vip else SecurityLevel.MEDIUM
                )
            
            # 4. 通过所有限制，增加计数
            global_window.count += 1
            global_window.last_request_time = current_time
            
            if ip_address:
                ip_window = self._get_window(ip_address, "ip")
                ip_window.count += 1
                ip_window.last_request_time = current_time
            
            session_window.count += 1
            session_window.last_request_time = current_time
            
            # 计算剩余配额
            remaining = min(
                global_limit + global_burst - global_window.count,
                session_limit + session_burst - session_window.count
            )
            
            return RateLimitResult(
                allowed=True,
                remaining=max(0, remaining),
                reset_after=60 - (current_time - global_window.start_time),
                limit=session_limit,
                severity=SecurityLevel.LOW
            )
    
    def record_request(self, session_id: str, query: str, ip_address: str = None, user_id: str = None):
        """记录请求历史（用于异常检测）"""
        with self._lock:
            current_time = time.time()
            
            # 记录请求
            self.request_history[session_id].append((current_time, query, ip_address))
            
            # 检查异常行为
            is_abnormal, reason, severity = self._check_abnormal_behavior(session_id, query, ip_address)
            if is_abnormal:
                self.stats["abnormal_detected"] += 1
                
                # 根据严重程度阻塞
                duration = self.config.BLOCK_DURATION[severity.value]
                self._block_entity("session", session_id, duration, severity)
                
                if ip_address and severity != SecurityLevel.LOW:
                    self._block_entity("ip", ip_address, duration * 2, severity)  # IP阻塞时间更长
    
    def get_stats(self) -> dict:
        """获取速率限制统计"""
        with self._lock:
            current_time = time.time()
            
            # 计算活跃会话数
            active_sessions = sum(1 for sess_id, history in self.request_history.items() 
                                if any(current_time - t < 60 for t, _, _ in history))
            
            # 计算活跃IP数
            active_ips = set()
            for history in self.request_history.values():
                for _, _, ip in history:
                    if ip:
                        active_ips.add(ip)
            
            return {
                "total_requests": self.stats["total_requests"],
                "blocked_requests": self.stats["blocked_requests"],
                "rate_limited_requests": self.stats["rate_limited_requests"],
                "filtered_inputs": self.stats["filtered_inputs"],
                "abnormal_detected": self.stats["abnormal_detected"],
                "malicious_detected": self.stats["malicious_detected"],
                "blocked_sessions": len([k for k, (t, _) in self.blocked_entities["sessions"].items() if t > current_time]),
                "blocked_ips": len([k for k, (t, _) in self.blocked_entities["ips"].items() if t > current_time]),
                "active_sessions": active_sessions,
                "active_ips": len(active_ips),
                "current_time": current_time
            }
    
    def clear_stats(self):
        """清除统计数据"""
        with self._lock:
            self.stats = {
                "total_requests": 0,
                "blocked_requests": 0,
                "rate_limited_requests": 0,
                "filtered_inputs": 0,
                "abnormal_detected": 0,
                "malicious_detected": 0
            }
    
    def get_blocked_entities(self) -> dict:
        """获取被阻塞的实体"""
        with self._lock:
            current_time = time.time()
            return {
                "sessions": {k: (t - current_time, s.value) for k, (t, s) in self.blocked_entities["sessions"].items() if t > current_time},
                "ips": {k: (t - current_time, s.value) for k, (t, s) in self.blocked_entities["ips"].items() if t > current_time},
                "users": {k: (t - current_time, s.value) for k, (t, s) in self.blocked_entities["users"].items() if t > current_time}
            }
    
    def add_vip_user(self, user_id: str):
        """添加VIP用户"""
        with self._lock:
            self.config.VIP_USERS.add(user_id)
    
    def remove_vip_user(self, user_id: str):
        """移除VIP用户"""
        with self._lock:
            self.config.VIP_USERS.discard(user_id)

# ======================
# 输入过滤器
# ======================

class InputFilterResult:
    """输入过滤结果"""
    
    def __init__(self, allowed: bool, filtered_text: str, reason: str = None, severity: SecurityLevel = SecurityLevel.LOW):
        self.allowed = allowed
        self.filtered_text = filtered_text
        self.reason = reason
        self.severity = severity
    
    def to_dict(self) -> dict:
        return {
            "allowed": self.allowed,
            "filtered_text": self.filtered_text,
            "reason": self.reason,
            "severity": self.severity.value
        }

class InputFilter:
    """智能输入过滤器"""
    
    def __init__(self, config: SecurityConfig = SECURITY_CONFIG):
        self.config = config
        
        # 敏感词替换映射
        self.sensitive_word_map = {}
        for word in self.config.SENSITIVE_WORDS:
            replacement = '*' * len(word)
            self.sensitive_word_map[word] = replacement
        
        # 编译恶意模式
        self.malicious_patterns = [
            re.compile(pattern, re.IGNORECASE) for pattern in self.config.MALICIOUS_PATTERNS
        ]
        
        # 白名单词汇（在特定上下文中允许）
        self.whitelist_words = {
            "admin", "root", "password", "database", "select", "delete", "update",
            "管理员", "密码", "数据库", "删除", "更新", "查询"
        }
        
        logger.info("🛡️ 输入过滤器初始化完成")
    
    def _check_length(self, text: str) -> Tuple[bool, str, SecurityLevel]:
        """检查输入长度"""
        if len(text) < self.config.MIN_INPUT_LENGTH:
            return False, f"输入太短，至少需要 {self.config.MIN_INPUT_LENGTH} 个字符", SecurityLevel.LOW
        
        if len(text) > self.config.MAX_INPUT_LENGTH:
            return False, f"输入太长，最多允许 {self.config.MAX_INPUT_LENGTH} 个字符", SecurityLevel.LOW
        
        return True, "", SecurityLevel.LOW
    
    def _check_characters(self, text: str) -> Tuple[bool, str, SecurityLevel]:
        """检查字符合法性"""
        invalid_chars = []
        for char in text:
            if char not in self.config.ALLOWED_CHARACTERS and not '\u4e00' <= char <= '\u9fff':
                if char not in ['\n', '\t', '\r']:  # 允许换行和制表符
                    invalid_chars.append(char)
        
        if invalid_chars:
            unique_invalid = list(set(invalid_chars))[:10]  # 取前10个唯一的非法字符
            return False, f"包含非法字符: {''.join(unique_invalid)}", SecurityLevel.MEDIUM
        
        return True, "", SecurityLevel.LOW
    
    def _filter_sensitive_words(self, text: str) -> str:
        """过滤敏感词"""
        text_lower = text.lower()
        filtered_text = text
        
        for word, replacement in self.sensitive_word_map.items():
            if len(word) < 2:  # 跳过单字符敏感词，避免误杀
                continue
            
            # 检查是否在白名单上下文中
            if any(whitelist_word in text_lower for whitelist_word in self.whitelist_words):
                continue
            
            if word in text_lower:
                # 保留原始大小写，只替换敏感词
                pattern = re.compile(re.escape(word), re.IGNORECASE)
                filtered_text = pattern.sub(replacement, filtered_text)
        
        return filtered_text
    
    def _detect_malicious_content(self, text: str) -> Tuple[bool, str, SecurityLevel]:
        """检测恶意内容"""
        text_lower = text.lower()
        
        # 1. 检查明显的攻击模式
        for pattern in self.malicious_patterns:
            match = pattern.search(text_lower)
            if match:
                matched_text = match.group(0)[:50]  # 只显示前50个字符
                return True, f"检测到潜在攻击模式: '{matched_text}...'", SecurityLevel.HIGH
        
        # 2. 检查可疑的字符组合
        suspicious_sequences = []
        
        # 检查连续的特殊字符
        special_chars = re.findall(r'[^\w\s\u4e00-\u9fff]{3,}', text_lower)
        if special_chars:
            suspicious_sequences.extend(special_chars[:3])
        
        # 检查长数字序列（排除订单号、手机号等已知业务格式）
        # 先移除已知的业务 ID 格式：订单号 ORD+数字、手机号
        cleaned = re.sub(r'ord\d{8,}', '', text_lower)  # 订单号 (text_lower 已小写)
        cleaned = re.sub(r'1[3-9]\d{9}', '', cleaned)     # 手机号
        long_numbers = re.findall(r'\d{8,}', cleaned)
        if long_numbers:
            suspicious_sequences.extend([f"长数字: {num[:10]}..." for num in long_numbers[:3]])
        
        # 检查编码字符串
        encoded_strings = re.findall(r'(?:[A-Za-z0-9+/]{4}){10,}', text_lower)
        if encoded_strings:
            suspicious_sequences.extend([f"编码字符串: {s[:20]}..." for s in encoded_strings[:3]])
        
        if suspicious_sequences:
            return True, f"检测到可疑内容: {', '.join(suspicious_sequences)}", SecurityLevel.MEDIUM
        
        # 3. 检查SQL注入特征
        sql_keywords = ['select', 'insert', 'update', 'delete', 'drop', 'union', 'create', 'alter', 'exec']
        sql_count = sum(1 for keyword in sql_keywords if keyword in text_lower)
        if sql_count >= 2:
            return True, f"检测到多个SQL关键词 ({sql_count}个)", SecurityLevel.MEDIUM
        
        # 4. 检查XSS特征
        xss_patterns = ['<script', 'javascript:', 'onerror=', 'onload=', 'iframe', 'alert(']
        xss_count = sum(1 for pattern in xss_patterns if pattern in text_lower)
        if xss_count >= 1:
            return True, f"检测到XSS攻击特征", SecurityLevel.HIGH
        
        return False, "", SecurityLevel.LOW
    
    def filter_input(self, text: str, session_id: str = None, context: dict = None) -> InputFilterResult:
        """过滤输入"""
        original_text = text.strip()
        
        # 1. 基础验证
        length_ok, length_reason, length_severity = self._check_length(original_text)
        if not length_ok:
            return InputFilterResult(
                allowed=False,
                filtered_text=original_text[:self.config.MAX_INPUT_LENGTH],
                reason=length_reason,
                severity=length_severity
            )
        
        # 2. 字符检查
        chars_ok, chars_reason, chars_severity = self._check_characters(original_text)
        if not chars_ok:
            return InputFilterResult(
                allowed=False,
                filtered_text=original_text,
                reason=chars_reason,
                severity=chars_severity
            )
        
        # 3. 恶意内容检测
        is_malicious, malicious_reason, malicious_severity = self._detect_malicious_content(original_text)
        if is_malicious:
            # 高严重度时阻塞会话
            if malicious_severity == SecurityLevel.HIGH and session_id:
                return InputFilterResult(
                    allowed=False,
                    filtered_text="⚠️ 检测到不安全内容，请求已被阻止",
                    reason=malicious_reason,
                    severity=malicious_severity
                )
            
            return InputFilterResult(
                allowed=False,
                filtered_text="⚠️ 检测到不安全内容，请求已被阻止",
                reason=malicious_reason,
                severity=malicious_severity
            )
        
        # 4. 敏感词过滤（不阻止，只替换）
        filtered_text = self._filter_sensitive_words(original_text)
        
        # 5. 特殊处理：如果过滤后内容为空或全是*
        if not filtered_text or all(c == '*' for c in filtered_text if c not in ' *'):
            return InputFilterResult(
                allowed=False,
                filtered_text="⚠️ 内容包含不适当信息，已被过滤",
                reason="包含大量不适当内容",
                severity=SecurityLevel.MEDIUM
            )
        
        return InputFilterResult(
            allowed=True,
            filtered_text=filtered_text,
            reason="输入已过滤敏感内容" if filtered_text != original_text else "输入通过安全检查",
            severity=SecurityLevel.LOW
        )
    
    def get_filter_stats(self) -> dict:
        """获取过滤统计"""
        return {
            "sensitive_words_count": len(self.sensitive_word_map),
            "malicious_patterns_count": len(self.malicious_patterns),
            "whitelist_words_count": len(self.whitelist_words)
        }

# ======================
# 安全管理器（主类）
# ======================

class SecurityManager:
    """安全管理器 - 统一接口"""
    
    def __init__(self, config: SecurityConfig = None):
        self.config = config or SECURITY_CONFIG
        self.rate_limiter = RateLimiter(self.config)
        self.input_filter = InputFilter(self.config)
        
        self.stats = {
            "total_requests": 0,
            "allowed_requests": 0,
            "blocked_requests": 0,
            "filtered_inputs": 0
        }
        
        logger.info("🛡️  安全管理器初始化完成")
    
    def check_request(self, session_id: str, ip_address: str = None, query: str = None, user_id: str = None) -> Tuple[bool, str, SecurityLevel]:
        """
        检查请求是否允许
        
        Args:
            session_id: 会话ID
            ip_address: 客户端IP地址
            query: 用户查询
            user_id: 用户ID
            
        Returns:
            Tuple[bool, str, SecurityLevel]: (是否允许, 原因, 安全级别)
        """
        self.stats["total_requests"] += 1
        
        # 1. 速率限制检查
        rate_limit_result = self.rate_limiter.check_rate_limit(session_id, ip_address, user_id)
        if not rate_limit_result.allowed:
            return False, f"速率限制: {rate_limit_result.reset_after:.1f}秒后重试", rate_limit_result.severity
        
        # 2. 输入过滤（如果有查询）
        if query:
            filter_result = self.input_filter.filter_input(query, session_id)
            if not filter_result.allowed:
                # 记录恶意请求
                if filter_result.severity == SecurityLevel.HIGH:
                    self.rate_limiter._block_entity("session", session_id, 
                                                   self.config.BLOCK_DURATION[filter_result.severity.value],
                                                   filter_result.severity)
                
                return False, filter_result.reason, filter_result.severity
        
        self.stats["allowed_requests"] += 1
        return True, "请求允许", SecurityLevel.LOW
    
    def filter_input(self, text: str, session_id: str = None) -> Tuple[str, bool, str]:
        """
        过滤输入文本
        
        Args:
            text: 输入文本
            session_id: 会话ID
            
        Returns:
            Tuple[str, bool, str]: (过滤后的文本, 是否安全, 原因)
        """
        filter_result = self.input_filter.filter_input(text, session_id)
        
        if not filter_result.allowed:
            self.stats["filtered_inputs"] += 1
        
        return filter_result.filtered_text, filter_result.allowed, filter_result.reason
    
    def record_request(self, session_id: str, query: str, ip_address: str = None, user_id: str = None):
        """
        记录请求（用于异常检测）
        
        Args:
            session_id: 会话ID
            query: 用户查询
            ip_address: 客户端IP地址
            user_id: 用户ID
        """
        self.rate_limiter.record_request(session_id, query, ip_address, user_id)
    
    def get_stats(self) -> dict:
        """
        获取安全统计信息
        
        Returns:
            dict: 统计信息
        """
        rate_stats = self.rate_limiter.get_stats()
        filter_stats = self.input_filter.get_filter_stats()
        
        return {
            "request_stats": {
                "total_requests": self.stats["total_requests"],
                "allowed_requests": self.stats["allowed_requests"],
                "blocked_requests": self.stats["blocked_requests"],
                "filtered_inputs": self.stats["filtered_inputs"]
            },
            "rate_limit_stats": rate_stats,
            "filter_stats": filter_stats,
            "blocked_entities": self.rate_limiter.get_blocked_entities()
        }
    
    def clear_stats(self):
        """清除所有统计数据"""
        self.stats = {
            "total_requests": 0,
            "allowed_requests": 0,
            "blocked_requests": 0,
            "filtered_inputs": 0
        }
        self.rate_limiter.clear_stats()
    
    def add_vip_user(self, user_id: str):
        """添加VIP用户"""
        self.rate_limiter.add_vip_user(user_id)
    
    def remove_vip_user(self, user_id: str):
        """移除VIP用户"""
        self.rate_limiter.remove_vip_user(user_id)
    
    def get_security_headers(self) -> dict:
        """
        获取安全相关的HTTP头
        
        Returns:
            dict: HTTP头字典
        """
        return {
            "X-Content-Type-Options": "nosniff",
            "X-Frame-Options": "DENY",
            "X-XSS-Protection": "1; mode=block",
            "Strict-Transport-Security": "max-age=31536000; includeSubDomains",
            "Content-Security-Policy": "default-src 'self'; script-src 'self' 'unsafe-inline' 'unsafe-eval'; style-src 'self' 'unsafe-inline'",
            "Referrer-Policy": "strict-origin-when-cross-origin",
            "Permissions-Policy": "geolocation=(), camera=(), microphone=()",
            "X-Permitted-Cross-Domain-Policies": "none"
        }
    
    def get_rate_limit_headers(self, rate_limit_result: RateLimitResult) -> dict:
        """
        获取速率限制相关的HTTP头
        
        Args:
            rate_limit_result: 速率限制结果
            
        Returns:
            dict: HTTP头字典
        """
        headers = {
            "X-RateLimit-Limit": str(rate_limit_result.limit),
            "X-RateLimit-Remaining": str(max(0, rate_limit_result.remaining)),
            "X-RateLimit-Reset": str(round(rate_limit_result.reset_after)),
            "X-RateLimit-Severity": rate_limit_result.severity.value
        }
        
        if not rate_limit_result.allowed:
            headers["Retry-After"] = str(round(rate_limit_result.reset_after))
            headers["X-RateLimit-Blocked"] = "true"
        
        return headers

# ======================
# 实用工具函数
# ======================

def validate_session_id(session_id: str) -> bool:
    """
    验证会话ID格式
    
    Args:
        session_id: 会话ID
        
    Returns:
        bool: 是否有效
    """
    if not session_id or not isinstance(session_id, str):
        return False
    
    # 会话ID应该只包含字母、数字、下划线、连字符
    if not re.match(r'^[a-zA-Z0-9_-]{1,64}$', session_id):
        return False
    
    # 避免特殊保留字
    reserved_ids = ["admin", "system", "root", "null", "undefined", "anonymous", "guest"]
    if session_id.lower() in reserved_ids:
        return False
    
    # 避免纯数字（可能被用作暴力破解）
    if session_id.isdigit():
        return False
    
    return True

def get_client_ip(request_context: Any = None) -> str:
    """
    获取客户端IP地址
    
    Args:
        request_context: 请求上下文对象
        
    Returns:
        str: 客户端IP地址
    """
    # 模拟实现 - 在实际应用中替换为真实的IP获取逻辑
    if request_context:
        if hasattr(request_context, 'headers'):
            # 从请求头获取
            ip = request_context.headers.get('X-Forwarded-For', '').split(',')[0].strip()
            if ip and ip != '127.0.0.1':
                return ip
        
        if hasattr(request_context, 'client'):
            return request_context.client.host
    
    # 模拟测试IP
    test_ips = [
        "192.168.1.100", "10.0.0.5", "172.16.0.1", 
        "8.8.8.8", "1.1.1.1", "203.0.113.1", "198.51.100.1"
    ]
    return random.choice(test_ips)

def safe_parse_input(input_data: Any) -> str:
    """
    安全解析用户输入
    
    Args:
        input_data: 输入数据
        
    Returns:
        str: 解析后的字符串
    """
    try:
        if isinstance(input_data, str):
            return input_data.strip()
        elif isinstance(input_data, dict):
            if 'query' in input_data:
                return str(input_data['query']).strip()
            elif 'question' in input_data:
                return str(input_data['question']).strip()
            elif 'input' in input_data:
                return str(input_data['input']).strip()
        elif hasattr(input_data, 'content'):
            return str(input_data.content).strip()
        
        # 默认处理
        return str(input_data).strip()
    
    except Exception as e:
        return f"解析错误: {str(e)}"

# ======================
# 测试函数
# ======================

def run_security_tests():
    """运行安全测试"""
    logger.info("\n" + "="*50)
    logger.info("🛡️  运行安全测试套件")
    logger.info("="*50)
    
    security = SecurityManager()
    
    test_cases = [
        # 正常查询
        ("龙井茶有什么特点？", "user_123", "192.168.1.100", "正常查询"),
        
        # 速率限制测试
        ("快速查询1", "user_123", "192.168.1.100", "速率限制测试"),
        ("快速查询2", "user_123", "192.168.1.100", "速率限制测试"),
        
        # 恶意输入测试
        ("<script>alert('xss')</script>", "attacker_1", "10.0.0.1", "XSS攻击"),
        ("' OR 1=1 --", "attacker_2", "10.0.0.2", "SQL注入"),
        ("delete from users", "attacker_3", "10.0.0.3", "恶意命令"),
        
        # VIP用户测试
        ("VIP查询", "vip_user_1", "192.168.1.200", "VIP用户"),
        
        # 长文本测试
        ("a" * 1500, "user_456", "192.168.1.101", "超长输入"),
        
        # 敏感词测试
        ("这个茶叶太他妈的好喝了", "user_789", "192.168.1.102", "敏感词过滤"),
    ]
    
    for i, (query, session_id, ip, desc) in enumerate(test_cases):
        logger.info(f"\n🔍 测试用例 {i+1}: {desc}")
        logger.info(f"   查询: '{query[:50]}{'...' if len(query) > 50 else ''}'")
        logger.info(f"   会话: {session_id}, IP: {ip}")
        
        allowed, reason, severity = security.check_request(session_id, ip, query)
        logger.info(f"   允许: {allowed}, 原因: {reason}, 严重程度: {severity.value}")
        
        if allowed:
            filtered_text, is_safe, filter_reason = security.filter_input(query, session_id)
            logger.info(f"   过滤后: '{filtered_text[:50]}{'...' if len(filtered_text) > 50 else ''}'")
            logger.info(f"   安全: {is_safe}, 过滤原因: {filter_reason}")
        
        # 记录请求
        security.record_request(session_id, query, ip)
        
        time.sleep(0.1)  # 防止测试过快
    
    # 显示统计信息
    stats = security.get_stats()
    logger.info("\n" + "="*50)
    logger.info("📊  安全统计信息")
    logger.info("="*50)
    logger.info(f"总请求数: {stats['request_stats']['total_requests']}")
    logger.info(f"允许请求: {stats['request_stats']['allowed_requests']}")
    logger.info(f"过滤输入: {stats['request_stats']['filtered_inputs']}")
    logger.info(f"活跃会话: {stats['rate_limit_stats']['active_sessions']}")
    logger.info(f"活跃IP: {stats['rate_limit_stats']['active_ips']}")

if __name__ == "__main__":
    run_security_tests()