"""
安全增强模块 - 解决专家测试发现的安全漏洞
包括：SQL注入防护、CSRF防护、权限控制
"""
import re
import secrets
import hashlib
import time
from typing import Dict, List, Optional, Tuple
from functools import wraps

class SecurityEnhancer:
    """安全增强器 - 防护各种安全漏洞"""
    
    def __init__(self):
        # CSRF令牌存储
        self.csrf_tokens = {}
        # 会话管理
        self.sessions = {}
        # 权限映射
        self.permissions = self._load_default_permissions()
        
        # 安全配置
        self.config = {
            "csrf_token_expiry": 3600,  # 1小时
            "session_expiry": 86400,    # 24小时
            "max_login_attempts": 5,
            "password_min_length": 8,
            "rate_limit_requests": 100,
            "rate_limit_window": 3600,  # 1小时
        }
    
    def _load_default_permissions(self) -> Dict:
        """加载默认权限配置"""
        return {
            "guest": ["read_gift_recommendations"],
            "user": ["read_gift_recommendations", "save_preferences", "view_history"],
            "premium": ["read_gift_recommendations", "save_preferences", "view_history", 
                       "premium_recommendations", "export_data"],
            "admin": ["read_gift_recommendations", "save_preferences", "view_history", 
                     "premium_recommendations", "export_data", "manage_users", "system_config"]
        }
    
    # ==================== SQL注入防护 ====================
    
    def sanitize_sql_input(self, input_text: str) -> str:
        """清理SQL输入，防止注入"""
        if not input_text:
            return ""
        
        # 移除SQL关键词
        sql_keywords = ["SELECT", "INSERT", "UPDATE", "DELETE", "DROP", "UNION", 
                       "OR", "AND", "WHERE", "FROM", "TABLE", "DATABASE", 
                       "EXEC", "EXECUTE", "TRUNCATE", "ALTER", "CREATE"]
        
        sanitized = input_text
        for keyword in sql_keywords:
            pattern = re.compile(re.escape(keyword), re.IGNORECASE)
            sanitized = pattern.sub("", sanitized)
        
        # 移除单引号和双引号
        sanitized = sanitized.replace("'", "").replace('"', '')
        
        # 移除分号和注释
        sanitized = sanitized.replace(";", "").replace("--", "").replace("#", "")
        
        # 限制长度
        if len(sanitized) > 500:
            sanitized = sanitized[:500]
        
        return sanitized.strip()
    
    def validate_sql_safe(self, input_text: str) -> Tuple[bool, str]:
        """验证输入是否安全，防止SQL注入"""
        if not input_text:
            return True, ""
        
        # SQL注入模式检测
        injection_patterns = [
            r"'.*--",        # 单引号加注释
            r"'.*#",         # 单引号加注释
            r";.*--",        # 分号加注释
            r"'.*OR.*'1'='1",  # OR 1=1 模式
            r"'.*OR.*'1'='1'", # OR '1'='1' 模式
            r"UNION.*SELECT",  # UNION SELECT
            r"'.*DROP.*TABLE", # DROP TABLE
            r"'.*DELETE.*FROM", # DELETE FROM
            r"'.*INSERT.*INTO", # INSERT INTO
            r"'.*UPDATE.*SET",  # UPDATE SET
        ]
        
        for pattern in injection_patterns:
            if re.search(pattern, input_text, re.IGNORECASE):
                return False, f"检测到SQL注入模式: {pattern}"
        
        return True, ""
    
    # ==================== CSRF防护 ====================
    
    def generate_csrf_token(self, user_id: str) -> str:
        """生成CSRF令牌"""
        token = secrets.token_urlsafe(32)
        expiry = time.time() + self.config["csrf_token_expiry"]
        
        self.csrf_tokens[user_id] = {
            "token": token,
            "expiry": expiry,
            "created": time.time()
        }
        
        return token
    
    def validate_csrf_token(self, user_id: str, token: str) -> bool:
        """验证CSRF令牌"""
        if user_id not in self.csrf_tokens:
            return False
        
        token_data = self.csrf_tokens[user_id]
        
        # 检查令牌是否匹配
        if token_data["token"] != token:
            return False
        
        # 检查令牌是否过期
        if time.time() > token_data["expiry"]:
            del self.csrf_tokens[user_id]
            return False
        
        return True
    
    def csrf_protect(self, user_id: str):
        """CSRF保护装饰器"""
        def decorator(func):
            @wraps(func)
            def wrapper(*args, **kwargs):
                # 从请求中获取CSRF令牌
                csrf_token = kwargs.get('csrf_token') or args[-1] if args else None
                
                if not csrf_token:
                    return {"success": False, "error": "缺少CSRF令牌"}
                
                if not self.validate_csrf_token(user_id, csrf_token):
                    return {"success": False, "error": "无效或过期的CSRF令牌"}
                
                return func(*args, **kwargs)
            return wrapper
        return decorator
    
    # ==================== 权限控制 ====================
    
    def check_permission(self, user_role: str, permission: str) -> bool:
        """检查用户是否有指定权限"""
        if user_role not in self.permissions:
            return False
        
        user_permissions = self.permissions[user_role]
        return permission in user_permissions
    
    def require_permission(self, permission: str):
        """权限检查装饰器"""
        def decorator(func):
            @wraps(func)
            def wrapper(*args, **kwargs):
                # 获取用户角色
                user_role = kwargs.get('user_role') or 'guest'
                
                if not self.check_permission(user_role, permission):
                    return {
                        "success": False, 
                        "error": f"权限不足，需要 '{permission}' 权限"
                    }
                
                return func(*args, **kwargs)
            return wrapper
        return decorator
    
    # ==================== 会话管理 ====================
    
    def create_session(self, user_id: str, user_role: str = "user") -> str:
        """创建用户会话"""
        session_id = secrets.token_urlsafe(32)
        expiry = time.time() + self.config["session_expiry"]
        
        self.sessions[session_id] = {
            "user_id": user_id,
            "user_role": user_role,
            "expiry": expiry,
            "created": time.time(),
            "last_activity": time.time(),
            "login_attempts": 0
        }
        
        return session_id
    
    def validate_session(self, session_id: str) -> Tuple[bool, Optional[Dict]]:
        """验证会话"""
        if session_id not in self.sessions:
            return False, None
        
        session = self.sessions[session_id]
        
        # 检查会话是否过期
        if time.time() > session["expiry"]:
            del self.sessions[session_id]
            return False, None
        
        # 更新最后活动时间
        session["last_activity"] = time.time()
        
        return True, session
    
    def revoke_session(self, session_id: str):
        """撤销会话"""
        if session_id in self.sessions:
            del self.sessions[session_id]
    
    # ==================== 输入验证增强 ====================
    
    def validate_input_length(self, input_text: str, max_length: int = 500) -> Tuple[bool, str]:
        """验证输入长度"""
        if len(input_text) > max_length:
            return False, f"输入过长，最多 {max_length} 字符"
        return True, ""
    
    def validate_email_format(self, email: str) -> Tuple[bool, str]:
        """验证邮箱格式"""
        pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$'
        if not re.match(pattern, email):
            return False, "邮箱格式不正确"
        return True, ""
    
    def validate_password_strength(self, password: str) -> Tuple[bool, List[str]]:
        """验证密码强度"""
        errors = []
        
        if len(password) < self.config["password_min_length"]:
            errors.append(f"密码长度至少为 {self.config['password_min_length']} 字符")
        
        if not re.search(r'[A-Z]', password):
            errors.append("密码应包含至少一个大写字母")
        
        if not re.search(r'[a-z]', password):
            errors.append("密码应包含至少一个小写字母")
        
        if not re.search(r'[0-9]', password):
            errors.append("密码应包含至少一个数字")
        
        if not re.search(r'[!@#$%^&*(),.?\":{}|<>]', password):
            errors.append("密码应包含至少一个特殊字符")
        
        return len(errors) == 0, errors
    
    # ==================== 速率限制 ====================
    
    def check_rate_limit(self, user_id: str, endpoint: str) -> Tuple[bool, int]:
        """检查速率限制"""
        # 这里实现简单的速率限制
        # 实际项目应该使用Redis等缓存
        key = f"{user_id}:{endpoint}"
        
        current_time = int(time.time())
        window_start = current_time - self.config["rate_limit_window"]
        
        # 模拟请求计数
        request_count = 0  # 这里应该从缓存中获取
        
        if request_count >= self.config["rate_limit_requests"]:
            return False, self.config["rate_limit_window"] - (current_time - window_start)
        
        return True, 0
    
    # ==================== XSS防护 ====================
    
    def sanitize_html(self, html: str) -> str:
        """清理HTML，防止XSS"""
        if not html:
            return ""
        
        # 移除危险标签
        dangerous_tags = ["script", "iframe", "object", "embed", "frame"]
        for tag in dangerous_tags:
            pattern = re.compile(f'<{tag}[^>]*>.*?</{tag}>', re.IGNORECASE | re.DOTALL)
            html = pattern.sub("", html)
        
        # 移除危险属性
        dangerous_attrs = ["onclick", "onload", "onerror", "onmouseover", "onmouseout",
                          "onkeydown", "onkeypress", "onkeyup", "onfocus", "onblur"]
        for attr in dangerous_attrs:
            pattern = re.compile(f'{attr}=["\'][^"\']*["\']', re.IGNORECASE)
            html = pattern.sub("", html)
        
        # 转义特殊字符
        html = (html.replace("&", "&amp;")
                     .replace("<", "&lt;")
                     .replace(">", "&gt;")
                     .replace('"', "&quot;")
                     .replace("'", "&#x27;"))
        
        return html
    
    # ==================== 安全头信息 ====================
    
    def get_security_headers(self) -> Dict[str, str]:
        """获取安全头信息"""
        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'; style-src 'self' 'unsafe-inline'",
            "Referrer-Policy": "strict-origin-when-cross-origin",
            "Permissions-Policy": "geolocation=(), microphone=(), camera=()"
        }
    
    # ==================== 日志记录 ====================
    
    def log_security_event(self, event_type: str, user_id: str, details: str):
        """记录安全事件"""
        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
        log_entry = f"[{timestamp}] [{event_type}] [{user_id}] {details}\n"
        
        # 这里应该写入文件或日志系统
        print(f"🔒 安全日志: {log_entry.strip()}")
        
        # 保存到文件
        with open("security.log", "a", encoding="utf-8") as f:
            f.write(log_entry)

# 全局安全增强器实例
security = SecurityEnhancer()

# 便捷函数
def protect_sql_input(input_text: str) -> str:
    """保护SQL输入"""
    return security.sanitize_sql_input(input_text)

def validate_sql_safety(input_text: str) -> Tuple[bool, str]:
    """验证SQL安全"""
    return security.validate_sql_safe(input_text)

def protect_html_content(html: str) -> str:
    """保护HTML内容，防止XSS"""
    return security.sanitize_html(html)

def check_user_permission(user_role: str, permission: str) -> bool:
    """检查用户权限"""
    return security.check_permission(user_role, permission)

# 测试
if __name__ == "__main__":
    print("=== 安全增强模块测试 ===")
    
    # 测试1: SQL注入防护
    print("\\n1. SQL注入防护测试:")
    test_inputs = [
        "正常输入",
        "SELECT * FROM users",
        "' OR '1'='1",
        "'; DROP TABLE users; --",
        "admin' --"
    ]
    
    for test in test_inputs:
        safe, msg = security.validate_sql_safe(test)
        print(f"  '{test[:20]}...' -> {'✅' if safe else '❌'} {msg}")
    
    # 测试2: CSRF防护
    print("\\n2. CSRF防护测试:")
    user_id = "test_user"
    token = security.generate_csrf_token(user_id)
    print(f"  生成的令牌: {token[:20]}...")
    print(f"  验证令牌: {'✅' if security.validate_csrf_token(user_id, token) else '❌'}")
    
    # 测试3: 权限控制
    print("\\n3. 权限控制测试:")
    test_cases = [
        ("guest", "read_gift_recommendations", True),
        ("guest", "manage_users", False),
        ("admin", "manage_users", True),
    ]
    
    for role, perm, expected in test_cases:
        result = security.check_permission(role, perm)
        status = "✅" if result == expected else "❌"
        print(f"  {role} 有权限 {perm}: {result} {status}")
    
    # 测试4: XSS防护
    print("\\n4. XSS防护测试:")
    xss_input = '<script>alert("xss")</script><p>正常内容</p>'
    safe_html = security.sanitize_html(xss_input)
    print(f"  原始: {xss_input}")
    print(f"  安全处理: {safe_html}")
    
    print("\\n✅ 安全模块测试完成")