# /etc/nginx/conf.d/waf-collector.conf
log_format waf_audit '$remote_addr - $request_time - $upstream_response_time - '
                     '$http_user_agent - $http_x_forwarded_for - '
                     '$request_method $request_uri $status $body_bytes_sent - '
                     '$http_referer - $http_cookie';

access_log /var/log/nginx/waf-audit.log waf_audit;

# 所有请求经过 WAF 模块
location / {
    modsecurity on;
    modsecurity_rules_file /etc/nginx/modsecurity/waf-rules.conf;
    proxy_pass http://backend;
}

# /etc/nginx/modsecurity/waf-baseline.conf
SecEngineActivation on
SecRequestBodyAccess On
SecResponseBodyAccess On
SecRequestBodyLimit 13107200
SecRequestBodyNoFilesLimit 131072

# 基础请求扫描
SecRule REQUEST_BODY "@rx (?i)(select|insert|update|delete|drop|union|concat)" \
    "id:100001,phase:2,deny,status:403,msg:'SQL 注入检测'"

SecRule REQUEST_BODY "@rx (?i)(<script|javascript:|onerror|onload)" \
    "id:100002,phase:2,deny,status:403,msg:'XSS 注入检测'"

# LLM 注入防护
SecRule REQUEST_BODY "@rx (?i)(prompt injection|ignore previous|jailbreak|sysprompt)" \
    "id:100003,phase:2,deny,status:403,msg:'LLM 注入检测'"

import psycopg2
import json
from datetime import datetime, timedelta
from collections import defaultdict
import numpy as np

class TrafficBaselineLearner:
    """流量基线学习器"""
    
    def __init__(self, db_config):
        self.conn = psycopg2.connect(**db_config)
        self.baseline = {}
    
    def collect_traffic_fingerprint(self, days=7):
        """收集流量指纹"""
        cutoff = datetime.now() - timedelta(days=days)
        
        query = """
        SELECT 
            remote_addr,
            request_method,
            request_uri,
            http_user_agent,
            http_referer,
            DATE_TRUNC('hour', timestamp) as hour_bucket,
            COUNT(*) as request_count
        FROM waf_audit_log
        WHERE timestamp > %s
        GROUP BY remote_addr, request_method, request_uri, 
                 http_user_agent, http_referer, hour_bucket
        ORDER BY request_count DESC
        """
        
        with self.conn.cursor() as cur:
            cur.execute(query, (cutoff,))
            rows = cur.fetchall()
        
        # 构建基线模型
        self.baseline = {
            'normal_uris': self._extract_normal_patterns(rows),
            'normal_agents': self._extract_agent_patterns(rows),
            'traffic_volume': self._compute_volume_baselines(rows),
            'geo_distribution': self._compute_geo_patterns(rows),
            'time_distribution': self._compute_time_patterns(rows),
            'collected_at': datetime.now().isoformat()
        }
        
        return self.baseline
    
    def _extract_normal_patterns(self, rows):
        """提取正常 URI 模式"""
        uri_freq = defaultdict(int)
        for row in rows:
            uri_freq[row[3]] += row[6]  # request_uri, request_count
        # 返回 top 1000 正常 URI
        return dict(sorted(uri_freq.items(), 
                          key=lambda x: x[1], reverse=True)[:1000])
    
    def _extract_agent_patterns(self, rows):
        """提取正常 User-Agent 模式"""
        agents = set()
        for row in rows:
            if row[4]:  # http_user_agent
                agents.add(row[4][:100])  # 截断存储
        return list(agents)[:500]
    
    def _compute_volume_baselines(self, rows):
        """计算流量体积基线（均值、标准差、分位数）"""
        volumes = [row[6] for row in rows]
        if not volumes:
            return {}
        return {
            'mean': np.mean(volumes),
            'std': np.std(volumes),
            'p95': np.percentile(volumes, 95),
            'p99': np.percentile(volumes, 99),
            'max': np.max(volumes)
        }
    
    def save_baseline(self):
        """保存基线到数据库"""
        with self.conn.cursor() as cur:
            cur.execute("""
                INSERT INTO waf_baseline (data, created_at) 
                VALUES (%s, %s)
            """, (json.dumps(self.baseline), datetime.now()))
        self.conn.commit()

-- 基线表结构
CREATE TABLE waf_baseline (
    id SERIAL PRIMARY KEY,
    data JSONB NOT NULL,
    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
    is_active BOOLEAN DEFAULT true
);

-- 请求审计表
CREATE TABLE waf_audit_log (
    id SERIAL PRIMARY KEY,
    timestamp TIMESTAMP NOT NULL DEFAULT NOW(),
    remote_addr INET,
    request_method VARCHAR(10),
    request_uri TEXT,
    status_code INTEGER,
    body_bytes_sent INTEGER,
    http_user_agent TEXT,
    http_referer TEXT,
    request_time FLOAT,
    upstream_response_time FLOAT,
    is_blocked BOOLEAN DEFAULT false,
    block_reason TEXT,
    ai_confidence FLOAT,
    processed_by_ai BOOLEAN DEFAULT false
);

CREATE INDEX idx_waf_audit_timestamp ON waf_audit_log(timestamp);
CREATE INDEX idx_waf_audit_addr ON waf_audit_log(remote_addr);
CREATE INDEX idx_waf_audit_blocked ON waf_audit_log(is_blocked);

import asyncio
import json
import redis
import psycopg2
from datetime import datetime
from typing import Dict, List, Optional, Tuple
import numpy as np
from openai import AsyncOpenAI

class AITraficDetector:
    """AI 流量检测引擎"""
    
    def __init__(self, redis_host='localhost', db_config=None, llm_config=None):
        self.redis = redis.Redis(host=redis_host, decode_responses=True)
        self.db_config = db_config
        self.llm_client = AsyncOpenAI(
            api_key=llm_config['api_key'],
            base_url=llm_config.get('base_url', 'https://api.openai.com/v1')
        )
        self.llm_model = llm_config.get('model', 'gpt-4o-mini')
        self.baseline_cache = None
    
    async def analyze_request(self, request_data: Dict) -> Dict:
        """分析单个请求，返回检测结果"""
        
        # Layer 1: 快速规则匹配
        fast_result = self._fast_rule_check(request_data)
        if fast_result['is_threat']:
            return fast_result
        
        # Layer 2: 统计分析
        stat_result = await self._statistical_analysis(request_data)
        if stat_result['threat_score'] > 0.8:
            # 高置信度威胁，直接拦截
            return stat_result
        
        # Layer 3: AI 深度分析（仅对可疑请求）
        if stat_result['threat_score'] > 0.4:
            ai_result = await self._llm_deep_analysis(request_data, stat_result)
            return ai_result
        
        # 正常请求
        return {
            'is_threat': False,
            'threat_score': 0.0,
            'action': 'allow',
            'reason': '正常流量',
            'layer': 'baseline'
        }
    
    def _fast_rule_check(self, request: Dict) -> Dict:
        """Layer 1: 基于已知模式的快速过滤"""
        threats = []
        
        # SQL 注入模式
        sql_patterns = [
            r"(?i)(\b(SELECT|INSERT|UPDATE|DELETE|DROP|UNION|ALTER)\b.*\b(FROM|INTO|TABLE|SET)\b)",
            r"(?i)(\b(OR|AND)\b\s+\d+\s*=\s*\d+)",
            r"(?i)(--|#|/\*)",
            r"(?i)(;\s*(DROP|DELETE|INSERT|UPDATE))",
        ]
        
        uri = request.get('request_uri', '')
        body = request.get('request_body', '')
        combined = uri + ' ' + body
        
        for pattern in sql_patterns:
            if __import__('re').search(pattern, combined):
                threats.append({
                    'type': 'sql_injection',
                    'pattern': pattern,
                    'match': combined[:100]
                })
        
        # XSS 模式
        xss_patterns = [
            r"(?i)<script[^>]*>",
            r"(?i)javascript\s*:",
            r"(?i)\bon\w+\s*=",
            r"(?i)<iframe[^>]*>",
        ]
        
        for pattern in xss_patterns:
            if __import__('re').search(pattern, combined):
                threats.append({
                    'type': 'xss',
                    'pattern': pattern,
                    'match': combined[:100]
                })
        
        # LLM 注入模式
        llm_patterns = [
            r"(?i)(prompt\s+injection|ignore\s+previous|jailbreak|sysprompt|dan\s+mode)",
            r"(?i)(you\s+are\s+now|let's\s+roleplay|pretend\s+that)",
        ]
        
        for pattern in llm_patterns:
            if __import__('re').search(pattern, combined):
                threats.append({
                    'type': 'llm_injection',
                    'pattern': pattern,
                    'match': combined[:100]
                })
        
        if threats:
            return {
                'is_threat': True,
                'threat_score': 0.95,
                'action': 'block',
                'reason': f"快速规则匹配: {[t['type'] for t in threats]}",
                'threats': threats,
                'layer': 'fast_rule'
            }
        
        return {'is_threat': False, 'threat_score': 0.0, 'action': 'allow', 'layer': 'fast_rule'}
    
    async def _statistical_analysis(self, request: Dict) -> Dict:
        """Layer 2: 基于统计的异常检测"""
        uri = request.get('request_uri', '/')
        addr = str(request.get('remote_addr', '0.0.0.0'))
        
        # 获取 URI 频率
        uri_count = self.redis.get(f"waf:uri_freq:{uri}") or 0
        uri_count = int(uri_count) if uri_count else 0
        
        # 获取请求者频率（滑动窗口）
        req_key = f"waf:addr_reqs:{addr}"
        req_count = self.redis.incr(req_key)
        self.redis.expire(req_key, 60)
        
        # 获取 URI 异常分数（基于对数频率）
        uri_anomaly = 0.0
        if uri_count > 0:
            # 低频 URI 更可疑
            uri_anomaly = max(0, 1.0 - np.log1p(uri_count) / 10.0)
        
        # 请求频率异常
        freq_anomaly = min(1.0, req_count / 100.0)
        
        # 综合威胁分数
        threat_score = 0.4 * uri_anomaly + 0.3 * freq_anomaly + 0.3 * self._path_anomaly(uri)
        
        return {
            'is_threat': threat_score > 0.5,
            'threat_score': round(threat_score, 3),
            'action': 'block' if threat_score > 0.8 else 'monitor',
            'uri_freq': uri_count,
            'addr_freq': req_count,
            'uri_anomaly': round(uri_anomaly, 3),
            'freq_anomaly': round(freq_anomaly, 3),
            'layer': 'statistical'
        }
    
    def _path_anomaly(self, uri: str) -> float:
        """计算 URI 路径异常分数"""
        # 长路径、特殊字符、编码异常
        length_score = min(1.0, len(uri) / 500.0)
        special_char_score = sum(1 for c in uri if c in '?&="\'%') / max(1, len(uri))
        encoding_score = 1.0 if '%0' in uri.lower() or '%27' in uri.lower() else 0.0
        return min(1.0, length_score * 0.4 + special_char_score * 0.3 + encoding_score * 0.3)
    
    async def _llm_deep_analysis(self, request: Dict, stat_result: Dict) -> Dict:
        """Layer 3: LLM 深度语义分析"""
        
        # 构建分析上下文
        context = {
            'request_uri': request.get('request_uri', ''),
            'request_method': request.get('request_method', 'GET'),
            'request_body': request.get('request_body', '')[:500],
            'user_agent': request.get('http_user_agent', ''),
            'remote_addr': str(request.get('remote_addr', '')),
            'statistical_score': stat_result['threat_score'],
            'analyzed_at': datetime.now().isoformat()
        }
        
        # 构建 Prompt
        prompt = f"""你是一个专业的 Web 安全分析师。分析以下 HTTP 请求是否包含攻击行为。

请求信息：
- Method: {context['request_method']}
- URI: {context['request_uri'][:200]}
- Body: {context['request_body'][:300]}
- User-Agent: {context['user_agent'][:100]}
- 统计威胁分数: {context['statistical_score']}

请分析：
1. 是否存在 SQL 注入、XSS、命令注入、路径遍历等攻击？
2. 是否存在 LLM 注入攻击（提示词注入、角色扮演绕过等）？
3. 该请求是否可能是正常业务流量？

以 JSON 格式返回分析结果：
{{
  "is_threat": true/false,
  "threat_type": "sql_injection|xss|command_injection|llm_injection|none",
  "threat_score": 0.0-1.0,
  "reason": "简要说明原因",
  "suggested_action": "block|rate_limit|challenge|allow"
}}"""
        
        try:
            response = await self.llm_client.chat.completions.create(
                model=self.llm_model,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.1,
                max_tokens=300
            )
            
            content = response.choices[0].message.content
            # 解析 JSON 响应
            ai_result = self._parse_llm_response(content)
            
            # 更新威胁分数（结合统计结果）
            final_score = 0.6 * ai_result.get('threat_score', 0) + 0.4 * stat_result['threat_score']
            ai_result['threat_score'] = round(final_score, 3)
            ai_result['layer'] = 'llm_deep'
            ai_result['statistical_context'] = stat_result
            
            return ai_result
            
        except Exception as e:
            # LLM 调用失败时，回退到统计结果
            return {
                'is_threat': stat_result['threat_score'] > 0.7,
                'threat_score': stat_result['threat_score'],
                'action': 'block' if stat_result['threat_score'] > 0.8 else 'monitor',
                'reason': f'LLM 分析失败，使用统计结果: {str(e)}',
                'layer': 'llm_fallback',
                'statistical_context': stat_result
            }
    
    def _parse_llm_response(self, content: str) -> Dict:
        """解析 LLM 返回的 JSON"""
        try:
            # 提取 JSON 部分
            import re
            json_match = re.search(r'\{[^{}]*\}', content, re.DOTALL)
            if json_match:
                return json.loads(json_match.group())
        except json.JSONDecodeError:
            pass
        
        # 默认安全结果
        return {
            'is_threat': False,
            'threat_type': 'none',
            'threat_score': 0.1,
            'reason': '无法解析分析结果，默认放行',
            'suggested_action': 'allow'
        }

import aiohttp
from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
import orjson

class RequestPipeline:
    """异步请求处理流水线"""
    
    def __init__(self, config):
        self.detector = AITraficDetector(**config['detector'])
        self.rule_generator = RuleGenerator(**config['rule_gen'])
        self.redis = redis.Redis(**config['redis'])
        self.db_config = config['db']
        
    async def process_request(self, request: Dict) -> Dict:
        """处理单个请求"""
        start_time = datetime.now()
        
        # 1. 异步检测
        detection_result = await self.detector.analyze_request(request)
        
        # 2. 执行动作
        action = self._decide_action(detection_result)
        
        # 3. 异步记录日志
        asyncio.create_task(self._log_request(request, detection_result, action))
        
        # 4. 如果是新攻击模式，触发规则生成
        if action == 'block' and detection_result.get('threat_type') == 'none':
            asyncio.create_task(
                self.rule_generator.generate_rules_from_attack(request, detection_result)
            )
        
        elapsed = (datetime.now() - start_time).total_seconds() * 1000
        detection_result['processing_time_ms'] = round(elapsed, 2)
        
        return detection_result
    
    def _decide_action(self, result: Dict) -> str:
        """决定最终动作"""
        score = result.get('threat_score', 0)
        action = result.get('action', 'allow')
        
        if score >= 0.9:
            return 'block'
        elif score >= 0.7:
            return 'rate_limit'
        elif score >= 0.5:
            return 'challenge'  # CAPTCHA 或 JS 挑战
        else:
            return 'allow'
    
    async def _log_request(self, request: Dict, result: Dict, action: str):
        """异步记录请求日志"""
        conn = psycopg2.connect(**self.db_config)
        try:
            with conn.cursor() as cur:
                cur.execute("""
                    INSERT INTO waf_audit_log 
                    (remote_addr, request_method, request_uri, status_code,
                     http_user_agent, http_referer, request_time,
                     is_blocked, block_reason, ai_confidence, processed_by_ai, timestamp)
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, NOW())
                """, (
                    str(request.get('remote_addr', '')),
                    request.get('request_method', 'GET'),
                    request.get('request_uri', '/'),
                    403 if action == 'block' else 200,
                    request.get('http_user_agent', '')[:200],
                    request.get('http_referer', '')[:200],
                    request.get('request_time', 0),
                    action == 'block',
                    result.get('reason', ''),
                    result.get('threat_score', 0),
                    result.get('layer', '') != 'fast_rule'
                ))
            conn.commit()
        finally:
            conn.close()

import re
from typing import List, Dict
import yaml

class RuleGenerator:
    """AI 驱动的规则生成器"""
    
    def __init__(self, db_config, output_dir='/etc/nginx/modsecurity/generated'):
        self.db_config = db_config
        self.output_dir = output_dir
        self.rule_counter = 200000  # 从 200000 开始编号
    
    async def generate_rules_from_attack(self, attack_request: Dict, 
                                          detection_result: Dict):
        """从攻击请求生成 WAF 规则"""
        
        # 1. 提取攻击特征
        features = self._extract_attack_features(attack_request, detection_result)
        
        # 2. 生成规则表达式
        rule_expr = self._generate_rule_expression(features)
        
        # 3. 构建 ModSecurity 规则
        modsec_rule = self._build_modsec_rule(features, rule_expr)
        
        # 4. 验证规则有效性
        if self._validate_rule(modsec_rule):
            # 5. 保存规则
            self._save_rule(modsec_rule, features)
            
            # 6. 重载 Nginx
            self._reload_nginx()
            
            return {'success': True, 'rule_id': modsec_rule['id']}
        
        return {'success': False, 'reason': '规则验证失败'}
    
    def _extract_attack_features(self, request: Dict, detection: Dict) -> Dict:
        """提取攻击特征"""
        uri = request.get('request_uri', '')
        body = request.get('request_body', '')
        combined = uri + ' ' + body
        
        return {
            'attack_type': detection.get('threat_type', 'unknown'),
            'pattern': self._find_attack_pattern(combined),
            'original_payload': combined[:200],
            'source_addr': str(request.get('remote_addr', '')),
            'user_agent': request.get('http_user_agent', '')[:100],
            'severity': self._assess_severity(detection.get('threat_score', 0)),
            'detected_at': datetime.now().isoformat()
        }
    
    def _find_attack_pattern(self, text: str) -> str:
        """从攻击文本中提取正则模式"""
        patterns = []
        
        # 移除安全字符，保留可疑部分
        suspicious = re.sub(r'[a-zA-Z0-9_\-./?=:&@#+ ]', '', text)
        if suspicious:
            patterns.append(re.escape(suspicious))
        
        # 检测编码攻击
        encoded = re.findall(r'%[0-9a-fA-F]{2}', text)
        if encoded:
            patterns.append('%[0-9a-fA-F]{2,}')
        
        # 检测特殊字符序列
        special = re.findall(r'[\x00-\x1f\x7f-\xff]{2,}', text)
        if special:
            patterns.append(r'[\x00-\x1f\x7f-\xff]{2,}')
        
        return '|'.join(patterns) if patterns else text[:50]
    
    def _generate_rule_expression(self, features: Dict) -> str:
        """生成规则表达式"""
        attack_type = features['attack_type']
        
        expressions = {
            'sql_injection': (
                r'(?i)(\b(SELECT|INSERT|UPDATE|DELETE|DROP|UNION|ALTER|CREATE|TRUNCATE)\b'
                r'.*\b(FROM|INTO|TABLE|SET|WHERE|DATABASE|INDEX)\b)'
            ),
            'xss': (
                r'(?i)(<script[^>]*>|javascript\s*:|on(error|load|click|mouseover)\s*=)'
            ),
            'command_injection': (
                r'(?i)(;|\||\$\(|`)(\s*(ls|cat|id|whoami|wget|curl|nc|bash|sh|python|perl)\b)'
            ),
            'path_traversal': r'(\.\./|\.\.\\|%2e%2e)',
            'llm_injection': (
                r'(?i)(prompt\s+injection|ignore\s+(all|previous|prior)|jailbreak|'
                r'dan\s+mode|you\s+are\s+now|system\s*prompt|developer\s+mode)'
            ),
            'rce': r'(?i)(eval|exec|system|passthrough|assert|include|require)\s*\(',
            'ssrf': r'(?i)(file://|gopher://|dict://|ssh://|redis://|localhost|127\.0\.0\.1)'
        }
        
        return expressions.get(attack_type, features['pattern'])
    
    def _build_modsec_rule(self, features: Dict, expression: str) -> Dict:
        """构建 ModSecurity 规则"""
        rule_id = self.rule_counter
        self.rule_counter += 1
        
        severity_map = {'high': 'CRITICAL', 'medium': 'ERROR', 'low': 'WARNING'}
        severity = severity_map.get(features['severity'], 'WARNING')
        
        rule = {
            'id': rule_id,
            'expression': expression,
            'phase': '2',
            'severity': severity,
            'action': 'deny,status:403',
            'message': f"AI 自动生成的 {features['attack_type']} 防护规则",
            'tags': [f'ai-auto-{features['attack_type']}', f'auto-generated'],
            'created_at': features['detected_at'],
            'source': 'ai_rule_generator'
        }
        
        # 构建完整的 SecRule 语法
        rule['modsec_syntax'] = (
            f"SecRule REQUEST_BODY|REQUEST_URI|ARGS \"@rx {expression}\" "
            f"\"id:{rule_id},phase:{rule['phase']},"
            f"{rule['action']},"
            f"msg:'{rule['message']}',"
            f"tag:'{rule['tags'][0]}',severity:'{severity}'\""
        )
        
        return rule
    
    def _validate_rule(self, rule: Dict) -> bool:
        """验证规则语法有效性"""
        # 检查正则表达式是否有效
        try:
            re.compile(rule['expression'])
            return True
        except re.error:
            return False
    
    def _save_rule(self, rule: Dict, features: Dict):
        """保存规则到文件和数据库"""
        import os
        
        # 追加到规则文件
        rule_file = os.path.join(self.output_dir, 'ai-generated-rules.conf')
        with open(rule_file, 'a') as f:
            f.write(f"\n# Auto-generated AI rule: {rule['id']}\n")
            f.write(f"# Type: {features['attack_type']}\n")
            f.write(f"# Pattern: {rule['expression'][:100]}\n")
            f.write(rule['modsec_syntax'] + "\n\n")
        
        # 保存到数据库
        conn = psycopg2.connect(**self.db_config)
        try:
            with conn.cursor() as cur:
                cur.execute("""
                    INSERT INTO waf_generated_rules 
                    (rule_id, rule_expression, attack_type, severity, 
                     modsec_syntax, source, created_at)
                    VALUES (%s, %s, %s, %s, %s, %s, NOW())
                """, (
                    rule['id'], rule['expression'], features['attack_type'],
                    features['severity'], rule['modsec_syntax'], 'ai_auto'
                ))
            conn.commit()
        finally:
            conn.close()
    
    def _reload_nginx(self):
        """重载 Nginx 配置"""
        import subprocess
        try:
            subprocess.run(['nginx', '-t'], check=True, capture_output=True)
            subprocess.run(['nginx', '-s', 'reload'], check=True, capture_output=True)
        except subprocess.CalledProcessError as e:
            # 配置测试失败，回滚规则
            self._rollback_last_rule()
            raise

-- 规则版本管理表
CREATE TABLE waf_generated_rules (
    id SERIAL PRIMARY KEY,
    rule_id INTEGER UNIQUE NOT NULL,
    rule_expression TEXT NOT NULL,
    attack_type VARCHAR(50) NOT NULL,
    severity VARCHAR(20) NOT NULL,
    modsec_syntax TEXT NOT NULL,
    source VARCHAR(50) DEFAULT 'ai_auto',
    is_active BOOLEAN DEFAULT true,
    auto_blocked_count INTEGER DEFAULT 0,
    false_positive_count INTEGER DEFAULT 0,
    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
    activated_at TIMESTAMP,
    deactivated_at TIMESTAMP
);

-- 规则性能追踪
CREATE TABLE waf_rule_performance (
    id SERIAL PRIMARY KEY,
    rule_id INTEGER REFERENCES waf_generated_rules(rule_id),
    triggered_at TIMESTAMP NOT NULL DEFAULT NOW(),
    request_uri TEXT,
    matched_payload TEXT,
    is_false_positive BOOLEAN DEFAULT false,
    user_feedback TEXT
);

-- 创建索引
CREATE INDEX idx_waf_rules_type ON waf_generated_rules(attack_type);
CREATE INDEX idx_waf_rules_active ON waf_generated_rules(is_active);
CREATE INDEX idx_waf_perf_rule ON waf_rule_performance(rule_id);

class RuleQualityEvaluator:
    """规则质量评估器"""
    
    def __init__(self, db_config, redis_config):
        self.db_config = db_config
        self.redis = redis.Redis(**redis_config)
    
    def evaluate_rules(self, days=7) -> Dict:
        """评估所有自动生成规则的质量"""
        conn = psycopg2.connect(**self.db_config)
        with conn.cursor() as cur:
            cur.execute("""
                SELECT 
                    r.rule_id,
                    r.attack_type,
                    r.is_active,
                    r.auto_blocked_count,
                    r.false_positive_count,
                    COUNT(p.id) as total_triggers,
                    SUM(CASE WHEN p.is_false_positive THEN 1 ELSE 0 END) as fp_count,
                    AVG(CASE WHEN p.is_false_positive THEN 1.0 ELSE 0.0 END) as fp_rate
                FROM waf_generated_rules r
                LEFT JOIN waf_rule_performance p ON r.rule_id = p.rule_id
                WHERE r.created_at > NOW() - INTERVAL '%s days'
                GROUP BY r.rule_id, r.attack_type, r.is_active, 
                         r.auto_blocked_count, r.false_positive_count
                ORDER BY r.created_at DESC
            """, (days,))
            
            rules = cur.fetchall()
        
        evaluation = []
        for rule in rules:
            rule_id, attack_type, is_active, blocked, fp, total, fp_count, fp_rate = rule
            
            # 质量评分
            if total == 0:
                quality_score = 0.5  # 新规则，未知
            else:
                # 高质量 = 高拦截 + 低误报
                quality_score = (
                    min(1.0, blocked / max(1, total)) * 0.6 +
                    (1.0 - min(1.0, fp_rate)) * 0.4
                )
            
            # 决策
            decision = 'keep'
            if quality_score < 0.3 and total >= 10:
                decision = 'deactivate'
            elif quality_score > 0.9 and total >= 5:
                decision = 'promote'
            
            evaluation.append({
                'rule_id': rule_id,
                'attack_type': attack_type,
                'quality_score': round(quality_score, 3),
                'total_triggers': total,
                'false_positive_rate': round(fp_rate, 4) if fp_rate else 0,
                'decision': decision,
                'recommendation': self._get_recommendation(decision, quality_score)
            })
        
        conn.close()
        return {'rules': evaluation, 'evaluated_at': datetime.now().isoformat()}
    
    def _get_recommendation(self, decision: str, score: float) -> str:
        recommendations = {
            'keep': '规则表现良好，继续保持',
            'deactivate': f'误报率过高({score:.2f})，建议停用或调整',
            'promote': '规则效果优异，可提升优先级'
        }
        return recommendations.get(decision, '需要人工复核')

class DynamicIPBlocklist:
    """动态 IP 封禁管理器"""
    
    def __init__(self, redis, threshold=10, block_duration=3600):
        self.redis = redis
        self.threshold = threshold
        self.block_duration = block_duration
    
    def record_threat(self, ip: str, score: float):
        """记录威胁行为"""
        threat_key = f"waf:threats:{ip}"
        
        # 累加威胁分数
        self.redis.incrbyfloat(threat_key, score)
        self.redis.expire(threat_key, self.block_duration)
        
        # 检查是否达到封禁阈值
        current_score = self.redis.getfloat(threat_key) or 0
        if current_score >= self.threshold:
            self._block_ip(ip, current_score)
    
    def _block_ip(self, ip: str, score: float):
        """执行 IP 封禁"""
        block_key = f"waf:blocked:{ip}"
        
        # 检查是否已封禁
        if self.redis.get(block_key):
            return
        
        # 写入封禁列表
        self.redis.setex(block_key, self.block_duration, json.dumps({
            'blocked_at': datetime.now().isoformat(),
            'threat_score': score,
            'reason': f"累计威胁分数 {score:.2f} 超过阈值 {self.threshold}"
        }))
        
        # 记录封禁事件
        print(f"[WAF] IP {ip} 已被封禁，威胁分数: {score:.2f}")
    
    def is_blocked(self, ip: str) -> bool:
        """检查 IP 是否被封禁"""
        return bool(self.redis.get(f"waf:blocked:{ip}"))
    
    def unblock_ip(self, ip: str):
        """手动解封 IP"""
        self.redis.delete(f"waf:blocked:{ip}")
        self.redis.delete(f"waf:threats:{ip}")

class FalsePositiveCorrector:
    """误报自我修正器"""
    
    def __init__(self, db_config, redis):
        self.db_config = db_config
        self.redis = redis
    
    def record_user_feedback(self, request_id: str, is_false_positive: bool, 
                              user_comment: str = ''):
        """记录用户反馈"""
        conn = psycopg2.connect(**self.db_config)
        try:
            with conn.cursor() as cur:
                cur.execute("""
                    INSERT INTO waf_user_feedback 
                    (request_id, is_false_positive, comment, created_at)
                    VALUES (%s, %s, %s, NOW())
                """, (request_id, is_false_positive, user_comment))
            conn.commit()
            
            # 更新规则误报计数
            if is_false_positive:
                self._update_rule_fp_count(request_id)
        finally:
            conn.close()
    
    def _update_rule_fp_count(self, request_id: str):
        """更新规则误报计数"""
        conn = psycopg2.connect(**self.db_config)
        try:
            with conn.cursor() as cur:
                # 获取请求关联的规则
                cur.execute("""
                    SELECT rule_id FROM waf_audit_log 
                    WHERE id = %s AND is_blocked = true
                """, (request_id,))
                row = cur.fetchone()
                
                if row:
                    rule_id = row[0]
                    # 更新误报计数
                    cur.execute("""
                        UPDATE waf_generated_rules 
                        SET false_positive_count = false_positive_count + 1
                        WHERE rule_id = %s
                    """, (rule_id,))
                    
                    # 检查是否需要自动停用
                    cur.execute("""
                        SELECT auto_blocked_count, false_positive_count, 
                               is_active FROM waf_generated_rules 
                        WHERE rule_id = %s
                    """, (rule_id,))
                    rule = cur.fetchone()
                    
                    if rule and rule[0] > 0:
                        fp_rate = rule[1] / rule[0]
                        if fp_rate > 0.3 and rule[2]:  # 误报率 > 30%
                            cur.execute("""
                                UPDATE waf_generated_rules 
                                SET is_active = false, deactivated_at = NOW()
                                WHERE rule_id = %s
                            """, (rule_id,))
                            print(f"[WAF] 规则 {rule_id} 因高误报率自动停用")
                    
                    conn.commit()
        finally:
            conn.close()
    
    def auto_tune_rules(self, batch_size=50):
        """自动调优规则"""
        conn = psycopg2.connect(**self.db_config)
        try:
            with conn.cursor() as cur:
                # 获取高误报率规则
                cur.execute("""
                    SELECT rule_id, attack_type, auto_blocked_count, 
                           false_positive_count, is_active
                    FROM waf_generated_rules
                    WHERE auto_blocked_count >= 10 
                      AND false_positive_count > 0
                      AND is_active = true
                    ORDER BY (false_positive_count::float / 
                              NULLIF(auto_blocked_count, 0)) DESC
                    LIMIT %s
                """, (batch_size,))
                
                rules = cur.fetchall()
                
                for rule_id, attack_type, blocked, fp, active in rules:
                    fp_rate = fp / blocked if blocked > 0 else 0
                    
                    if fp_rate > 0.25:
                        # 建议降低拦截灵敏度
                        cur.execute("""
                            UPDATE waf_generated_rules 
                            SET severity = 'WARNING',
                                action_override = 'rate_limit'
                            WHERE rule_id = %s
                        """, (rule_id,))
                        print(f"[WAF] 规则 {rule_id} 调整拦截策略为 rate_limit")
                
                conn.commit()
        finally:
            conn.close()

# docker-compose.waf.yml
version: '3.8'

services:
  # Nginx + ModSecurity
  nginx-waf:
    image: nginx:1.25-alpine
    ports:
      - "80:80"
      - "443:443"
    volumes:
      - ./nginx/conf.d:/etc/nginx/conf.d
      - ./nginx/modsecurity:/etc/nginx/modsecurity
      - ./logs/nginx:/var/log/nginx
    depends_on:
      - ai-engine
      - redis
      - postgres
    restart: unless-stopped

  # AI 检测引擎
  ai-engine:
    build: ./ai-engine
    environment:
      - REDIS_HOST=redis
      - DB_HOST=postgres
      - DB_NAME=waf
      - DB_USER=waf_user
      - DB_PASSWORD=${WAF_DB_PASSWORD}
      - LLM_API_KEY=${LLM_API_KEY}
      - LLM_BASE_URL=${LLM_BASE_URL}
      - LEARNING_MODE=true  # 学习模式开关
    volumes:
      - ./ai-engine/rules:/app/rules
      - ./logs/ai:/app/logs
    depends_on:
      - redis
      - postgres
    restart: unless-stopped

  # Redis 缓存
  redis:
    image: redis:7-alpine
    command: redis-server --appendonly yes
    volumes:
      - redis_data:/data
    restart: unless-stopped

  # PostgreSQL 存储
  postgres:
    image: postgres:16-alpine
    environment:
      - POSTGRES_DB=waf
      - POSTGRES_USER=waf_user
      - POSTGRES_PASSWORD=${WAF_DB_PASSWORD}
    volumes:
      - postgres_data:/var/lib/postgresql/data
    restart: unless-stopped

  # 规则可视化 Dashboard
  waf-dashboard:
    build: ./waf-dashboard
    ports:
      - "8080:80"
    environment:
      - API_URL=http://ai-engine:8000
    depends_on:
      - ai-engine
      - postgres
    restart: unless-stopped

volumes:
  redis_data:
  postgres_data:

# ai-engine/Dockerfile
FROM python:3.11-slim

WORKDIR /app

# 安装系统依赖
RUN apt-get update && apt-get install -y \
    gcc \
    libpq-dev \
    && rm -rf /var/lib/apt/lists/*

# 安装 Python 依赖
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# 复制应用代码
COPY . .

# 暴露 API 端口
EXPOSE 8000

# 启动命令
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

# ai-engine/requirements.txt
fastapi==0.109.0
uvicorn[standard]==0.27.0
psycopg2-binary==2.9.9
redis==5.0.1
openai==1.12.0
numpy==1.26.3
aiohttp==3.9.1
aiokafka==0.11.0
pydantic==2.5.3
python-multipart==0.0.6
orjson==3.9.12

# 1. 设置环境变量
cp .env.example .env
# 编辑 .env，填入 LLM_API_KEY 和 WAF_DB_PASSWORD

# 2. 初始化数据库
docker-compose -f docker-compose.waf.yml run --rm ai-engine \
    python manage.py migrate

# 3. 启动学习模式（7-14 天）
docker-compose -f docker-compose.waf.yml up -d
# 设置 LEARNING_MODE=true

# 4. 学习完成后切换到防护模式
# 编辑 .env，设置 LEARNING_MODE=false
docker-compose -f docker-compose.waf.yml restart ai-engine

# 5. 查看防护效果
curl http://localhost:8080/dashboard/stats

# 推荐配置
ai_engine:
  # 学习模式持续时间（天）
  learning_days: 7
  
  # 各层检测阈值
  thresholds:
    fast_rule_block: 0.95
    statistical_block: 0.8
    llm_trigger: 0.4
  
  # 规则质量阈值
  rule_quality:
    auto_deactivate_fp_rate: 0.3
    auto_deactivate_min_triggers: 10
    promote_min_triggers: 5
  
  # IP 封禁
  ip_blocking:
    threshold_score: 10.0
    block_duration_hours: 24
  
  # LLM 调用
  llm:
    model: gpt-4o-mini  # 性价比最高的选择
    max_tokens: 300
    temperature: 0.1
    timeout_seconds: 5