diff --git a/.env.example b/.env.example index 22138eb..9eb0b50 100644 --- a/.env.example +++ b/.env.example @@ -4,6 +4,7 @@ # AI服务配置 SSLOGS_AI_API_KEY=your_api_key_here +# cloud_provider 可选：deepseek、custom（任意 OpenAI 兼容端点）或 zhipu（智谱 GLM，支持 GLM Coding Plan） SSLOGS_AI_CLOUD_PROVIDER=deepseek SSLOGS_AI_LOCAL_PROVIDER=ollama SSLOGS_AI_TYPE=cloud diff --git a/.gitignore b/.gitignore index 617f6f6..c11d146 100644 --- a/.gitignore +++ b/.gitignore @@ -112,6 +112,8 @@ coverage.xml api_keys.txt secrets.json private_config.yaml +# config.yaml 含本地真实 API Key（如智谱），不提交；从 config.example.yaml 复制生成 +config.yaml # MaxMind GeoIP configuration with real credentials config/GeoIP.conf.local diff --git a/config.example.yaml b/config.example.yaml index a84f477..92baa32 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -73,6 +73,8 @@ ai_analysis: ai: type: cloud local_provider: ollama + # cloud_provider 可选：deepseek（读 deepseek: 段）、custom（读 custom: 段，任意 OpenAI 兼容端点）、 + # 或 zhipu（读 zhipu: 段，智谱 GLM；支持 GLM Coding Plan 编程套餐端点） cloud_provider: deepseek max_retries: 3 retry_delay: 1 @@ -93,6 +95,33 @@ ollama: base_url: http://localhost:11434/api/chat timeout: 60 +# 自定义 OpenAI 兼容端点配置 +# 当 ai.cloud_provider 设置为 custom 时启用，可指向任意 OpenAI 兼容 API +# （OpenAI / 通义千问 / Moonshot / SiliconFlow 自选模型 / 自建 vLLM / Ollama 的 OpenAI 兼容口 等）。 +# +# base_url 为 OpenAI 兼容的 base 地址（通常以 /v1 结尾，无需包含 /chat/completions）。 +# api_key 建议留空并使用环境变量：export SSLOGS_AI_API_KEY="your-actual-api-key" +custom: + base_url: https://api.openai.com/v1 # 例：https://api.openai.com/v1 或 http://localhost:11434/v1 + api_key: "" # 留空则读取环境变量 SSLOGS_AI_API_KEY + model: gpt-4o-mini # 任意模型名，可在 GUI 中从 /models 拉取 + timeout: 30 + max_tokens: 2048 + +# 智谱 GLM 配置（当 ai.cloud_provider 设置为 zhipu 时启用） +# 智谱提供 OpenAI 兼容端点，国内直连。base_url 留空时按 coding_plan 自动派生： +# coding_plan=false -> https://open.bigmodel.cn/api/paas/v4 （标准按量计费 API） +# coding_plan=true -> https://open.bigmodel.cn/api/coding/paas/v4（GLM Coding Plan 编程套餐） +# api_key 建议留空并使用环境变量：export SSLOGS_AI_API_KEY="your-zhipu-api-key" +# 模型可选：glm-4.6 / glm-4.5 / glm-4-plus / glm-4-flash / glm-4.7-FlashX / glm-5.2 等 +zhipu: + base_url: "" # 留空则按 coding_plan 派生默认端点 + api_key: "" # 留空则读取环境变量 SSLOGS_AI_API_KEY + model: glm-4.6 + coding_plan: false # 开启后走 GLM Coding Plan 专属端点（与套餐 Key 配合） + timeout: 120 # GLM-5 等推理模型较慢，需要更长超时 + max_tokens: 4096 # GLM-5 等推理模型需要更大空间 + # LM Studio 配置 (v3.1新增) lm_studio: base_url: http://localhost:1234/v1/chat/completions diff --git a/config.yaml b/config.yaml deleted file mode 100644 index 8b99085..0000000 --- a/config.yaml +++ /dev/null @@ -1,164 +0,0 @@ -log_path: logs/*.log -server: - ip: 192.168.1.100 -log_format: - type: web - timestamp_format: '%d/%b/%Y:%H:%M:%S %z' - fields: - src_ip: (\d+\.\d+\.\d+\.\d+) - timestamp: \[(.*?)\] - request_line: '"([A-Z]+\s+[^\s]+\s+HTTP/[\d\.]+)"' - request_method: '"([A-Z]+)\s' - request_path: '"[A-Z]+\s+([^\s]+)\s+' - http_version: '(HTTP/\d\.\d)"' - status_code: '"\s+(\d{3})\s+' - response_size: '\s+(\d+)\s+' - processing_time: '\s+([\d.]+)\s+' - user_agent: '"user_agent":"([^"]*)"' - client_info: '{(.*?)}' -rule_dir: rules -rules: - path: rules/ - default_severity: medium -<<<<<<< HEAD - case_sensitive: false - max_matches: 50 -======= - path: rules/ ->>>>>>> dbcfc95695db2065f726934b2d7b14ab8dd244c4 - auto_reload: true -analysis: - batch_size: 1000 - max_events: 100 - memory_limit_mb: 500 - dynamic_batch_size: true - memory_optimization: aggressive - rule_preloading: true -context_lines: 5 -ai_analysis: - high_risk_only: true - successful_attacks_only: true - # 攻击成功的HTTP状态码定义：只保留2xx/3xx成功响应。 - # 注意：403/404/405/500表示请求被拒绝或失败，不应视为"攻击成功"。 - success_status_codes: - - '200' - - '201' - - '202' - - '204' - - '301' - - '302' - - '304' - max_ai_analysis: 5 - high_risk_severity: high -output_dir: output -report_type: html -report: - template: templates/report_template.html - output_format: html - output_directory: reports -# ⚠️ 安全提示：生产环境中建议使用环境变量存储API密钥 -# 示例：export SSLOGS_AI_API_KEY="your-actual-api-key" -deepseek: - api_key: demo_key_for_testing # 示例密钥，生产环境请使用环境变量 - model: deepseek-ai/DeepSeek-V3 - base_url: https://api.siliconflow.cn/v1/chat/completions - timeout: 30 - max_tokens: 2048 -ai: - type: cloud - local_provider: ollama - cloud_provider: deepseek - max_retries: 3 - retry_delay: 1 - retry_backoff: 2 - default_timeout: 30 - -# 白名单配置：上下文分析阶段会用这些列表做误报抑制。 -# 缺省值见 core/rule_engine.py 的 _DEFAULT_* 常量；此处可按环境覆盖。 -# 注意：纯按 User-Agent 过滤爬虫存在被伪造绕过的风险，建议仅作辅助。 -whitelist: - # 安全爬虫/合法机器人 User-Agent 关键字（命中即跳过该条告警） - safe_agents: - - googlebot - - bingbot - - slurp - - duckduckbot - - baiduspider - - yandexbot - - facebookexternalhit - - twitterbot - - linkedinbot - - pinterest - - applebot - - semrushbot - - mj12bot - - ahrefsbot - - dotbot - - archive.org_bot - # 静态资源扩展名（请求以这些结尾时通常非攻击目标） - static_extensions: - - .css - - .js - - .jpg - - .jpeg - - .png - - .gif - - .svg - - .ico - - .woff - - .woff2 - - .ttf - - .eot - - .otf - - .mp4 - - .mp3 - - .avi - - .mov - - .pdf - - .xml - # 健康检查/监控端点 - health_paths: - - /health - - /healthcheck - - /ping - - /status - - /metrics - - /actuator/health - - /ready - - /.well-known/ - - /robots.txt - - /favicon.ico - # 内部服务请求来源（命中后降低优先级但不完全过滤） - internal_ips: - - 127.0.0.1 - - ::1 - - localhost - -# 增强的规则引擎配置 -rule_engine: - enable_context_analysis: true - threat_scoring: true - false_positive_filter: true - adaptive_threshold: true - rule_cache_size: 1000 - enable_precompilation: true - -# 性能监控配置 -performance: - enable_monitoring: true - track_memory: true - track_time: true - error_rate_threshold: 10 - memory_warning_threshold: 500 - performance_log_interval: 100 -ollama: - model: deepseek-r1:14b - base_url: http://localhost:11434/api/chat - timeout: 60 -geoip_db_path: config/GeoLite2-Country.mmdb - -# 日志配置 -logging: - level: INFO - format: '%(asctime)s - %(name)s - %(levelname)s - %(message)s' - file: logs/sslogs.log diff --git a/core/ai_analyzer.py b/core/ai_analyzer.py index 342d22c..097167d 100644 --- a/core/ai_analyzer.py +++ b/core/ai_analyzer.py @@ -3,12 +3,73 @@ import logging import yaml import time -from typing import Dict, Any, Optional +from typing import Dict, Any, Optional, List from requests.adapters import HTTPAdapter from urllib3.util.retry import Retry from .exceptions import AIServiceError, AIServiceUnavailableError, AIAuthenticationError, AIRateLimitError +def _derive_chat_url(base_url: str) -> str: + """从 base_url 派生 chat completions 端点。 + + 兼容两种写法： + - 完整地址（含 /chat/completions，如旧版 deepseek 段）-> 原样返回 + - 仅 base（如 https://api.openai.com/v1）-> 追加 /chat/completions + """ + base = (base_url or '').strip() + if not base: + return '' + if base.rstrip('/').endswith('/chat/completions'): + return base + return base.rstrip('/') + '/chat/completions' + + +def _derive_models_url(base_url: str) -> str: + """从 base_url 派生 GET /models 端点（去掉可能的 /chat/completions 后缀）。""" + base = (base_url or '').strip() + if not base: + return '' + if base.rstrip('/').endswith('/chat/completions'): + base = base.rstrip('/')[:-len('/chat/completions')] + return base.rstrip('/') + '/models' + + +def fetch_available_models(base_url: str, api_key: str = '', timeout: int = 10) -> List[str]: + """从任意 OpenAI 兼容端点拉取可用模型列表（GET /models）。 + + base_url 可以是 base（如 https://api.openai.com/v1）或完整 chat completions 地址， + 函数会自动派生 /models 端点。失败时返回空列表（不抛异常），便于 GUI 下拉容错。 + + Args: + base_url: OpenAI 兼容 API 的 base 地址 + api_key: 可选的 API 密钥（Bearer 认证） + timeout: 请求超时秒数 + + Returns: + 可用模型 id 列表；失败时返回空列表 + """ + logger = logging.getLogger('core.ai_analyzer') + models_url = _derive_models_url(base_url) + if not models_url: + logger.warning("fetch_available_models: base_url 为空") + return [] + + headers = {'Content-Type': 'application/json'} + if api_key: + headers['Authorization'] = f'Bearer {api_key}' + + try: + response = requests.get(models_url, headers=headers, timeout=timeout) + response.raise_for_status() + data = response.json() + models = [m.get('id') for m in data.get('data', []) if m.get('id')] + logger.info(f"从 {models_url} 获取到 {len(models)} 个模型") + return models + except Exception as e: + logger.warning(f"拉取模型列表失败 ({models_url}): {e}") + return [] + + class AIAnalyzer: """AI分析器 - 支持云端和本地AI模型""" @@ -36,17 +97,8 @@ def __init__(self, config_path: str = 'config.yaml'): # 初始化HTTP会话（连接池优化） self._init_http_session() - # 加载云端模型配置 - if self.cloud_provider == 'deepseek': - self.deepseek_config = self.config.get('deepseek', {}) - # 优先从环境变量获取API密钥，提高安全性 - self.api_key = self._get_secure_api_key() - self.cloud_model = self.deepseek_config.get('model', 'deepseek-ai/DeepSeek-V3') - self.cloud_base_url = self.deepseek_config.get('base_url', 'https://api.siliconflow.cn/v1/chat/completions') - self.cloud_headers = { - 'Authorization': f'Bearer {self.api_key}', - 'Content-Type': 'application/json' - } + # 加载云端模型配置（支持 deepseek 与任意 OpenAI 兼容的 custom 端点） + self._load_cloud_config() # 加载本地模型配置 if self.local_provider == 'ollama': @@ -88,23 +140,88 @@ def _get_config_with_env(self, section: str, key: str, env_var: str, default: st # 其次使用配置文件 return self.config.get(section, {}).get(key, default) - def _get_secure_api_key(self) -> str: - """安全地获取API密钥，优先从环境变量读取""" + def _get_secure_api_key(self, config_section: Optional[Dict[str, Any]] = None) -> str: + """安全地获取API密钥，优先从环境变量读取 + + Args: + config_section: 当前 provider 的配置段（如 deepseek/custom）。 + 为 None 时回退到 self.deepseek_config，保持向后兼容。 + """ # 1. 首先检查环境变量 env_api_key = os.environ.get(self.ENV_API_KEY) if env_api_key: self.logger.debug("从环境变量获取API密钥") return env_api_key - + # 2. 其次从配置文件读取 - config_api_key = self.deepseek_config.get('api_key', '') + if config_section is None: + config_section = getattr(self, 'deepseek_config', {}) or {} + config_api_key = config_section.get('api_key', '') if config_api_key and config_api_key not in ['your-api-key-here', 'YOUR_API_KEY', 'demo_key_for_testing', '']: self.logger.debug("从配置文件获取API密钥") return config_api_key - + self.logger.warning("API密钥未配置，请设置环境变量 SSLOGS_AI_API_KEY 或修改配置文件") return '' + def _load_cloud_config(self): + """加载云端 provider 配置。 + + 支持三种云端 provider，三者均使用 OpenAI Chat Completions 协议： + - 'deepseek'：读取 deepseek: 段（向后兼容） + - 'custom'：读取 custom: 段，可指向任意 OpenAI 兼容端点（base_url + api_key + model） + - 'zhipu'：读取 zhipu: 段（智谱 GLM）；coding_plan=True 时默认走编程套餐端点 + + 三种 provider 都填充同一组属性：self.cloud_model / cloud_base_url / + cloud_headers / api_key / cloud_max_tokens / cloud_timeout。 + """ + if self.cloud_provider == 'deepseek': + section = self.config.get('deepseek', {}) + default_base_url = 'https://api.siliconflow.cn/v1/chat/completions' + default_model = 'deepseek-ai/DeepSeek-V3' + self.deepseek_config = section + elif self.cloud_provider == 'custom': + section = self.config.get('custom', {}) + default_base_url = '' # custom 必须显式配置 base_url + default_model = '' + self.custom_config = section + elif self.cloud_provider == 'zhipu': + section = self.config.get('zhipu', {}) + coding_plan = bool(section.get('coding_plan', False)) + # 智谱均为 OpenAI 兼容端点；编程套餐(coding_plan)走专属 coding 端点 + default_base_url = ('https://open.bigmodel.cn/api/coding/paas/v4' + if coding_plan else 'https://open.bigmodel.cn/api/paas/v4') + default_model = 'glm-4.6' + self.zhipu_config = section + else: + # 未支持的云端 provider：置空云端配置，避免误用 + self.api_key = '' + self.cloud_model = '' + self.cloud_base_url = '' + self.cloud_headers = {'Content-Type': 'application/json'} + self.cloud_max_tokens = 1024 + self.cloud_timeout = self.default_timeout + return + + self.api_key = self._get_secure_api_key(section) + self.cloud_model = section.get('model', default_model) + # base_url 为空字符串或缺省时回退到默认端点（智谱据此按 coding_plan 派生默认端点） + self.cloud_base_url = section.get('base_url') or default_base_url + self.cloud_headers = { + 'Authorization': f'Bearer {self.api_key}', + 'Content-Type': 'application/json' + } + self.cloud_max_tokens = section.get('max_tokens', 1024) + self.cloud_timeout = section.get('timeout', self.default_timeout) + + def _chat_url(self) -> str: + """云端 chat completions 地址（兼容 base_url 含/不含 /chat/completions）。""" + return _derive_chat_url(self.cloud_base_url) + + def _models_url(self) -> str: + """云端 GET /models 地址。""" + return _derive_models_url(self.cloud_base_url) + def _load_config(self, config_path: str) -> Dict[str, Any]: try: with open(config_path, 'r', encoding='utf-8') as f: @@ -550,7 +667,7 @@ def _analyze_with_cloud(self, prompt: str) -> str: payload = { "model": self.cloud_model, "stream": False, - "max_tokens": self.deepseek_config.get('max_tokens', 1024), + "max_tokens": self.cloud_max_tokens, "temperature": 0.7, "top_p": 0.7, "messages": [{"role": "user", "content": prompt}] @@ -558,17 +675,18 @@ def _analyze_with_cloud(self, prompt: str) -> str: try: response = self._make_request_with_retry( - self.cloud_base_url, + self._chat_url(), self.cloud_headers, payload, - self.deepseek_config.get('timeout', 30) + self.cloud_timeout ) result = response.json() # 处理云端API响应格式 if 'choices' in result and len(result['choices']) > 0: message = result['choices'][0].get('message', {}) - content = message.get('content', '') + # 部分推理模型（如智谱 GLM-5 系列）把输出放在 reasoning_content，content 可能为空 + content = message.get('content') or message.get('reasoning_content') or '' if content: return content else: @@ -646,12 +764,12 @@ def _generate_rules_with_cloud(self, prompt: str) -> str: "model": self.cloud_model, "messages": [{"role": "user", "content": prompt}] } - + response = self._make_request_with_retry( - self.cloud_base_url, + self._chat_url(), self.cloud_headers, payload, - self.deepseek_config.get('timeout', 60) + self.cloud_timeout ) result = response.json() return result.get("choices", [{}])[0].get("message", {}).get("content", "") diff --git a/core/config_manager.py b/core/config_manager.py index 4828bfd..ddcb20e 100644 --- a/core/config_manager.py +++ b/core/config_manager.py @@ -104,6 +104,19 @@ def _validate_config(self) -> List[str]: api_key = deepseek_config.get('api_key', '') if not api_key or api_key in ['your-api-key-here', 'YOUR_API_KEY', 'demo_key_for_testing']: errors.append("DeepSeek API密钥未正确配置 (当前使用占位符)") + elif cloud_provider == 'custom': + # 自定义 OpenAI 兼容端点：必须有 base_url（api_key 允许留空走环境变量） + custom_config = self._config.get('custom', {}) + if not custom_config.get('base_url'): + errors.append("自定义AI配置(custom)缺少 base_url") + if not custom_config.get('model'): + errors.append("自定义AI配置(custom)缺少 model") + elif cloud_provider == 'zhipu': + # 智谱 GLM：要求 model；base_url 允许留空（由 coding_plan 派生默认）， + # api_key 允许留空（走环境变量 SSLOGS_AI_API_KEY） + zhipu_config = self._config.get('zhipu', {}) + if not zhipu_config.get('model'): + errors.append("智谱AI配置(zhipu)缺少 model") # 4. 验证路径字段 for path_field in ['rule_dir', 'output_dir']: @@ -119,7 +132,7 @@ def _validate_config(self) -> List[str]: # 5. 验证枚举类型字段 enum_validations = { 'ai.type': ['cloud', 'local'], - 'ai.cloud_provider': ['deepseek', 'openai'], + 'ai.cloud_provider': ['deepseek', 'openai', 'custom', 'zhipu'], 'ai.local_provider': ['ollama', 'lm_studio'], 'report_type': ['html', 'json', 'markdown'], } @@ -146,6 +159,8 @@ def _validate_config(self) -> List[str]: 'ai_analysis.max_ai_analysis': (1, 100, "AI分析数量应在1-100之间"), 'deepseek.timeout': (5, 300, "DeepSeek超时应在5-300秒之间"), 'ollama.timeout': (5, 600, "Ollama超时应在5-600秒之间"), + 'custom.timeout': (5, 600, "自定义AI超时应在5-600秒之间"), + 'zhipu.timeout': (5, 600, "智谱AI超时应在5-600秒之间"), } for field_path, (min_val, max_val, error_msg) in range_validations.items(): @@ -158,6 +173,7 @@ def _validate_config(self) -> List[str]: url_validations = [ 'deepseek.base_url', 'ollama.base_url', + 'custom.base_url', ] for field_path in url_validations: @@ -300,6 +316,25 @@ def _set_defaults(self): ollama.setdefault('base_url', 'http://localhost:11434/api/chat') ollama.setdefault('timeout', 60) + # 自定义 OpenAI 兼容端点默认值（cloud_provider=custom 时启用） + self._config.setdefault('custom', {}) + custom = self._config['custom'] + custom.setdefault('base_url', '') + custom.setdefault('api_key', '') + custom.setdefault('model', '') + custom.setdefault('timeout', 30) + custom.setdefault('max_tokens', 2048) + + # 智谱 GLM 默认值（cloud_provider=zhipu 时启用） + self._config.setdefault('zhipu', {}) + zhipu = self._config['zhipu'] + zhipu.setdefault('base_url', '') # 留空则由 coding_plan 派生默认端点 + zhipu.setdefault('api_key', '') # 留空则读取环境变量 SSLOGS_AI_API_KEY + zhipu.setdefault('model', 'glm-4.6') + zhipu.setdefault('coding_plan', False) + zhipu.setdefault('timeout', 120) + zhipu.setdefault('max_tokens', 4096) # GLM-5 等推理模型需要更大空间产出正式回答 + # 重试配置默认值 ai.setdefault('max_retries', 3) ai.setdefault('retry_delay', 1) @@ -321,13 +356,15 @@ def get_safe_config(self) -> Dict[str, Any]: """获取安全的配置（隐藏敏感信息）""" config = self.get_config().copy() - # 隐藏敏感信息 - if 'deepseek' in config and 'api_key' in config['deepseek']: - api_key = config['deepseek']['api_key'] - if api_key and len(api_key) > 8: - config['deepseek']['api_key'] = api_key[:4] + '*' * (len(api_key) - 8) + api_key[-4:] - else: - config['deepseek']['api_key'] = '***' + # 隐藏敏感信息（deepseek、智谱与自定义端点的 api_key 均需脱敏） + for section_name in ('deepseek', 'custom', 'zhipu'): + section = config.get(section_name) + if isinstance(section, dict) and 'api_key' in section: + api_key = section['api_key'] + if api_key and len(api_key) > 8: + config[section_name]['api_key'] = api_key[:4] + '*' * (len(api_key) - 8) + api_key[-4:] + else: + config[section_name]['api_key'] = '***' return config diff --git a/core/event_bus.py b/core/event_bus.py index 4dda68a..8772510 100644 --- a/core/event_bus.py +++ b/core/event_bus.py @@ -45,6 +45,25 @@ class EventHandler: retry_on_error: bool = True +def _normalize_priority(priority: Any) -> EventPriority: + """将优先级归一化为 EventPriority 枚举。 + + 兼容调用方传入字符串（如 'high'/'normal'/'low'/'critical'，大小写不敏感； + 'medium'/'default' 视为 NORMAL）。避免后续 handler.priority.value 在字符串上 + 触发 'str' object has no attribute 'value'。 + """ + if isinstance(priority, EventPriority): + return priority + if isinstance(priority, str): + key = priority.strip().upper() + alias = {'MEDIUM': 'NORMAL', 'DEFAULT': 'NORMAL'} + try: + return EventPriority[alias.get(key, key)] + except KeyError: + return EventPriority.NORMAL + return EventPriority.NORMAL + + class EventBus: """事件总线 - 实现模块间的解耦通信""" @@ -67,6 +86,7 @@ def subscribe(self, event_name: str, handler_func: Callable, filter_func: Optional[Callable[[Event], bool]] = None, timeout: Optional[float] = None, retry_on_error: bool = True) -> str: """订阅事件""" + priority = _normalize_priority(priority) handler_name = name or f"{handler_func.__module__}.{handler_func.__name__}" handler = EventHandler( @@ -100,6 +120,7 @@ def subscribe_global(self, handler_func: Callable, name: Optional[str] = None, priority: EventPriority = EventPriority.NORMAL, filter_func: Optional[Callable[[Event], bool]] = None) -> str: """订阅所有事件（全局处理器）""" + priority = _normalize_priority(priority) handler_name = name or f"{handler_func.__module__}.{handler_func.__name__}" handler = EventHandler( diff --git a/core/unified_config_manager.py b/core/unified_config_manager.py index 91a70f5..da7f7a4 100644 --- a/core/unified_config_manager.py +++ b/core/unified_config_manager.py @@ -273,6 +273,19 @@ def validate_config(self, config: Dict[str, Any]) -> List[str]: deepseek_config = ai_config.get('deepseek', {}) if not deepseek_config.get('api_key'): errors.append("DeepSeek配置缺少api_key") + elif ai_config['cloud_provider'] == 'custom': + # 自定义 OpenAI 兼容端点：要求 base_url 与 model + # (api_key 允许留空，走环境变量 SSLOGS_AI_API_KEY) + custom_config = ai_config.get('custom', {}) + if not custom_config.get('base_url'): + errors.append("自定义AI配置缺少base_url") + if not custom_config.get('model'): + errors.append("自定义AI配置缺少model") + elif ai_config['cloud_provider'] == 'zhipu': + # 智谱 GLM：要求 model（base_url 可由 coding_plan 派生默认，api_key 可走环境变量） + zhipu_config = config.get('zhipu', {}) + if not zhipu_config.get('model'): + errors.append("智谱AI配置缺少model") # 规则配置验证 if 'rules' in config: diff --git a/gui_optimized.py b/gui_optimized.py index 4ee97ab..2490b74 100644 --- a/gui_optimized.py +++ b/gui_optimized.py @@ -6,9 +6,12 @@ import sys import os +import re import asyncio import time import threading +import yaml +import requests from pathlib import Path # 添加项目根目录到Python路径 @@ -22,7 +25,7 @@ QTableWidgetItem, QHeaderView, QSplitter, QFrame, QScrollArea, QDialog, QDialogButtonBox, QFormLayout ) -from PyQt6.QtCore import Qt, QThread, pyqtSignal, QTimer, QPropertyAnimation, QEasingCurve +from PyQt6.QtCore import Qt, QThread, pyqtSignal, pyqtSlot, QTimer, QPropertyAnimation, QEasingCurve from PyQt6.QtGui import QFont, QPixmap, QIcon, QColor, QPalette # 导入优化后的核心模块 @@ -34,6 +37,19 @@ validate_input, sanitize_input, handle_exceptions, create_error_context, cache_result, CachePolicy ) +from core.ai_analyzer import fetch_available_models + + +# 从完整访问日志行中抽取「请求目标」(method + path + query)，仅用于安全校验。 +# 这样可避免把后端 IP(如 192.168.x)、Cookie 分隔符(;)、UA 中的 HTML 实体(&#)等 +# 误判为攻击载荷，显著降低对整行日志校验时的误报。 +_REQUEST_TARGET_RE = re.compile(r'"([A-Z]+)\s+([^"\s]+)\s+HTTP/[\d.]+"') + + +def _extract_request_target(log_line: str) -> str: + """提取日志行里的请求方法+路径(含查询串)；提取失败则回退为原行。""" + m = _REQUEST_TARGET_RE.search(log_line or '') + return f"{m.group(1)} {m.group(2)}" if m else (log_line or '') class AnalysisWorker(QThread): @@ -72,6 +88,8 @@ def run(self): } self.progress_updated.emit(100, "✅ 分析完成！", performance_data) + # 将统计写回 results，供 analysis_completed 读取（否则汇总恒为 0） + results['performance'] = performance_data self.analysis_finished.emit(True, "分析成功完成", results) except Exception as e: @@ -139,8 +157,9 @@ def process_log_entry(log_line): try: # 安全验证 if security_validation: - validation = validate_input(log_line) - if not validation.is_valid: + validation = validate_input(_extract_request_target(log_line)) + # 按「是否检测到威胁」判定（NORMAL 级下 is_valid 恒为 True，会漏计威胁） + if validation.threats or not validation.is_valid: threat_data = { 'line': log_line[:100] + "...", 'threats': validation.threats, @@ -176,38 +195,81 @@ def process_log_entry(log_line): f"📝 已处理 {processed_count} 行日志，内存使用: {memory_usage:.1%}", {'memory_usage': memory_usage, 'processed_count': processed_count}) - # 简化的AI分析（模拟） - if ai_enabled and results['processed_logs']: + # 真实AI分析（读 config.yaml，支持智谱等 provider） + if ai_enabled and (results['threats'] or results['processed_logs']): self.progress_updated.emit(85, "🤖 启动AI分析...", {}) - self._simulate_ai_analysis(results) + self._run_ai_analysis(results) return results - def _simulate_ai_analysis(self, results): - """模拟AI分析（简化版）""" + def _run_ai_analysis(self, results): + """真实AI分析：用 AIAnalyzer（读 config.yaml，支持智谱等 provider）对检测到的威胁做深度分析。 + + 结果写入 results['threats'][i]['ai_analysis'] 并实时输出到日志面板。 + AIAnalyzer.analyze_log 内部已含降级（AI 不可用时返回结构化备用分析），不会中断流程。 + 依赖：分析前请先在 GUI 点「保存AI配置」，让 config.yaml 的 provider/端点/coding_plan 生效。 + """ try: - # 准备AI分析数据 - log_count = min(len(results['processed_logs']), 10) # 限制数量 + from core.ai_analyzer import AIAnalyzer + config_path = str(Path(__file__).parent / 'config.yaml') + analyzer = AIAnalyzer(config_path=config_path) + except Exception as e: + self.log_output.emit(f"⚠️ AI分析器初始化失败: {e}") + return - for i in range(log_count): - if self.is_interrupted: - break + # 界面上刚填的 Key/模型优先（无需先点保存），端点/coding_plan 沿用 config.yaml 已保存配置 + gui_key = (self.config.get('api_key') or '').strip() + if gui_key: + analyzer.api_key = gui_key + analyzer.cloud_headers['Authorization'] = f'Bearer {gui_key}' + gui_model = (self.config.get('model_name') or '').strip() + if gui_model: + analyzer.cloud_model = gui_model + + self.log_output.emit( + f"🧠 真实AI分析 provider={getattr(analyzer, 'cloud_provider', '?')} " + f"model={analyzer.cloud_model or '?'} ..." + ) - # 模拟AI分析结果 - log_entry = results['processed_logs'][i] - log_entry['ai_analysis'] = f"AI分析结果 #{i+1}: 威胁等级低，建议持续监控" + threats = results.get('threats', []) + if not threats: + self.log_output.emit("ℹ️ 未检测到威胁，跳过AI分析") + return - # 模拟处理延迟 - time.sleep(0.1) # 短暂延迟模拟 + # 去重 + 限量，避免对相同攻击重复消耗额度 + seen, targets = set(), [] + for th in threats: + key = th.get('line', '') + if key and key not in seen: + seen.add(key) + targets.append(th) + if len(targets) >= 10: + break + + total = len(targets) + for i, th in enumerate(targets): + if self.is_interrupted: + break + line = th.get('line', '') + names = [getattr(t, 'value', str(t)) for t in th.get('threats', [])] + attack_name = ', '.join(names) if names else None + try: + ai_result = analyzer.analyze_log( + line, attack_name=attack_name, threat_score=th.get('risk_score') + ) + except Exception as e: + ai_result = f"⚠️ AI分析异常: {e}" + th['ai_analysis'] = ai_result - if i % 3 == 0: # 每3条更新一次进度 - progress = 85 + (i * 15 // log_count) - self.progress_updated.emit(progress, f"🧠 AI分析进度: {i+1}/{log_count}", {}) + # 实时把分析结果输出到日志面板（截断过长内容） + preview = ai_result.replace('\n', ' ')[:200] + tail = '...' if len(ai_result) > 200 else '' + self.log_output.emit(f"🤖 [{attack_name or '威胁'}] {preview}{tail}") - self.log_output.emit(f"🧠 AI分析完成，处理了 {log_count} 条日志") + progress = 85 + int((i + 1) / total * 14) + self.progress_updated.emit(min(progress, 99), f"🧠 AI分析进度: {i+1}/{total}", {}) - except Exception as e: - self.log_output.emit(f"⚠️ AI分析失败: {e}") + self.log_output.emit(f"🧠 AI分析完成，共分析 {total} 条威胁") def _get_memory_usage(self): """获取当前内存使用率""" @@ -238,14 +300,42 @@ def interrupt(self): self.log_output.emit("⏹️ 用户中断分析...") +class _BackgroundAIWorker(QThread): + """在后台线程运行 AI 网络任务（拉取模型 / 测试连接），避免阻塞 UI。 + + task 是一个无参 callable，返回 (ok: bool, message: str, payload)。 + 结果通过 result_ready 信号回传到主线程。 + """ + result_ready = pyqtSignal(bool, str, object) + + def __init__(self, task): + super().__init__() + self._task = task + + def run(self): + try: + ok, message, payload = self._task() + except Exception as e: + ok, message, payload = False, f"请求异常: {e}", None + self.result_ready.emit(bool(ok), str(message), payload) + + class OptimizedLogAnalyzerGUI(QMainWindow): """SSlogs v3.1 优化版GUI界面""" + # AI 模型类型下拉项与 cloud_provider 的映射 + _AI_TYPE_CLOUD = "🌐 自定义 OpenAI 兼容" + def __init__(self): super().__init__() self.setWindowTitle("🚀 SSlogs v3.1 - 企业级智能安全日志分析平台") self.setGeometry(100, 100, 1200, 900) + # 配置文件路径（项目根目录下的 config.yaml） + self.config_path = str(Path(__file__).parent / 'config.yaml') + # 后台 AI 网络任务句柄（防止被 GC 回收） + self._ai_worker = None + # 设置应用图标和样式 self._setup_appearance() @@ -266,10 +356,15 @@ def __init__(self): # 初始化变量 self.worker = None self.start_time = 0 + self.last_results = None # 最近一次分析结果，供「导出报告」使用 + self.last_log_path = '' # 最近一次分析的日志路径 # 初始化定时器 self._setup_timers() + # 从 config.yaml 回填 AI 配置到界面 + self._load_ai_config_from_yaml() + def _setup_appearance(self): """设置界面外观""" # 设置应用样式 @@ -495,10 +590,27 @@ def _create_ai_config(self, parent_layout): model_layout = QHBoxLayout() model_layout.addWidget(QLabel("AI模型类型:")) self.ai_model_type = QComboBox() - self.ai_model_type.addItems(["云端 DeepSeek", "本地 LM Studio", "本地 Ollama"]) + self.ai_model_type.addItems(["云端 DeepSeek", "本地 LM Studio", "本地 Ollama", "🌐 自定义 OpenAI 兼容", "🇨🇳 智谱 GLM (Zhipu)"]) + self.ai_model_type.currentTextChanged.connect(self._on_ai_type_changed) model_layout.addWidget(self.ai_model_type) service_layout.addLayout(model_layout) + # API基础URL（自定义/OpenAI 兼容端点） + base_url_layout = QHBoxLayout() + base_url_layout.addWidget(QLabel("API地址:")) + self.base_url_input = QLineEdit() + self.base_url_input.setPlaceholderText("OpenAI 兼容 base，如 https://api.openai.com/v1") + base_url_layout.addWidget(self.base_url_input) + service_layout.addLayout(base_url_layout) + + # 编程套餐开关（仅「智谱 GLM」provider 时可见） + # 开启后 base_url 切到智谱专属 coding 端点（/api/coding/paas/v4），与套餐 Key 配合使用。 + self.zhipu_coding_plan_cb = QCheckBox("🧑‍💻 编程套餐 (Coding Plan) — 走智谱专属端点") + self.zhipu_coding_plan_cb.setChecked(False) + self.zhipu_coding_plan_cb.setVisible(False) + self.zhipu_coding_plan_cb.stateChanged.connect(self._on_zhipu_coding_plan_toggled) + service_layout.addWidget(self.zhipu_coding_plan_cb) + # 模型选择 model_select_layout = QHBoxLayout() model_select_layout.addWidget(QLabel("模型名称:")) @@ -516,7 +628,7 @@ def _create_ai_config(self, parent_layout): api_layout.addWidget(QLabel("API密钥:")) self.api_key_input = QLineEdit() self.api_key_input.setEchoMode(QLineEdit.EchoMode.Password) - self.api_key_input.setPlaceholderText("云端服务API密钥") + self.api_key_input.setPlaceholderText("云端服务API密钥（可改用环境变量 SSLOGS_AI_API_KEY）") api_layout.addWidget(self.api_key_input) service_layout.addLayout(api_layout) @@ -534,12 +646,18 @@ def _create_ai_config(self, parent_layout): parent_layout.addWidget(service_group) # 测试区域 - test_group = QGroupBox("🧪 连接测试") + test_group = QGroupBox("🧪 连接与配置") test_layout = QVBoxLayout() + test_action_layout = QHBoxLayout() test_btn = QPushButton("测试AI连接") test_btn.clicked.connect(self.test_ai_connection) - test_layout.addWidget(test_btn) + test_action_layout.addWidget(test_btn) + + save_ai_btn = QPushButton("💾 保存AI配置") + save_ai_btn.clicked.connect(self._save_ai_config_to_yaml) + test_action_layout.addWidget(save_ai_btn) + test_layout.addLayout(test_action_layout) self.ai_status_label = QLabel("状态: 未测试") test_layout.addWidget(self.ai_status_label) @@ -800,49 +918,316 @@ def browse_output_dir(self): if directory: self.output_dir_input.setText(directory) + # 下拉文本 -> (cloud_provider/local_provider, ai.type, 是否需要 base_url) + _PROVIDER_MAP = { + "云端 DeepSeek": ("deepseek", "cloud", True), + "本地 LM Studio": ("lm_studio", "local", True), + "本地 Ollama": ("ollama", "local", True), + "🌐 自定义 OpenAI 兼容": ("custom", "cloud", True), + "🇨🇳 智谱 GLM (Zhipu)": ("zhipu", "cloud", True), + } + + # 智谱 GLM 端点与预置模型（均为 OpenAI 兼容） + _ZHIPU_LABEL = "🇨🇳 智谱 GLM (Zhipu)" + _ZHIPU_BASE_STANDARD = "https://open.bigmodel.cn/api/paas/v4" + _ZHIPU_BASE_CODING = "https://open.bigmodel.cn/api/coding/paas/v4" + _ZHIPU_MODELS_STANDARD = [ + "glm-4.6", "glm-4.5", "glm-4.5-air", "glm-4.7", + "glm-5", "glm-5-turbo", "glm-5.1", "glm-5.2", + ] + _ZHIPU_MODELS_CODING = ["glm-4.6", "glm-4.5", "glm-4.7-FlashX", "glm-5.2"] + + def _is_zhipu_selected(self) -> bool: + """当前下拉是否选中智谱 provider。""" + return self.ai_model_type.currentText() == self._ZHIPU_LABEL + + def _zhipu_default_base_url(self) -> str: + """按当前 Coding Plan 复选框状态返回智谱默认 base_url。""" + coding = getattr(self, 'zhipu_coding_plan_cb', None) and self.zhipu_coding_plan_cb.isChecked() + return self._ZHIPU_BASE_CODING if coding else self._ZHIPU_BASE_STANDARD + + def _zhipu_curated_models(self) -> list: + """按 Coding Plan 状态返回预置模型列表。""" + coding = getattr(self, 'zhipu_coding_plan_cb', None) and self.zhipu_coding_plan_cb.isChecked() + return list(self._ZHIPU_MODELS_CODING if coding else self._ZHIPU_MODELS_STANDARD) + + def _populate_zhipu_models(self): + """用预置模型列表填充下拉（可编辑，保留用户已填模型名）。""" + current = self.model_name_combo.currentText() + self.model_name_combo.clear() + self.model_name_combo.addItems(self._zhipu_curated_models()) + # 跳过加载占位符，仅保留真实模型名 + if current and current not in ("⏳ 正在拉取模型列表...", "正在加载模型列表..."): + self.model_name_combo.setCurrentText(current) + + def _on_zhipu_coding_plan_toggled(self, _state: int): + """切换 Coding Plan：用户显式勾选/取消，直接切到对应官方端点并刷新模型列表。""" + if not self._is_zhipu_selected(): + return + # 勾选编程套餐是显式动作：直接采用对应官方端点（coding 或 standard） + self.base_url_input.setText(self._zhipu_default_base_url()) + self._populate_zhipu_models() + + def _on_ai_type_changed(self, text: str): + """切换 AI 模型类型时，更新 base_url 输入的可用性、占位提示，及智谱 Coding Plan 开关可见性。""" + mapping = self._PROVIDER_MAP.get(text) + self.base_url_input.setEnabled(bool(mapping and mapping[2])) + placeholders = { + "本地 Ollama": "如 http://localhost:11434/api/chat", + "本地 LM Studio": "如 http://localhost:1234/v1", + "云端 DeepSeek": "如 https://api.siliconflow.cn/v1/chat/completions", + "🌐 自定义 OpenAI 兼容": "OpenAI 兼容 base，如 https://api.openai.com/v1", + self._ZHIPU_LABEL: "如 https://open.bigmodel.cn/api/paas/v4", + } + self.base_url_input.setPlaceholderText(placeholders.get(text, "API 地址")) + + # Coding Plan 开关仅对智谱可见 + is_zhipu = text == self._ZHIPU_LABEL + if hasattr(self, 'zhipu_coding_plan_cb'): + self.zhipu_coding_plan_cb.setVisible(is_zhipu) + if is_zhipu: + current_url = self.base_url_input.text().strip() + # 切到智谱时，若当前地址不是智谱端点（可能继承自 deepseek 等其它 provider）， + # 重置为智谱默认端点；已是智谱端点则保留。 + if current_url not in (self._ZHIPU_BASE_STANDARD, self._ZHIPU_BASE_CODING): + self.base_url_input.setText(self._zhipu_default_base_url()) + self._populate_zhipu_models() + + def _current_ai_provider(self): + """返回当前下拉选择的 (provider_name, ai_type)。""" + text = self.ai_model_type.currentText() + mapping = self._PROVIDER_MAP.get(text) + if mapping: + return mapping[0], mapping[1] + return "deepseek", "cloud" + def refresh_models(self): - """刷新AI模型列表""" + """从当前配置的 OpenAI 兼容端点拉取可用模型列表（后台线程，不阻塞 UI）。""" + base_url = self.base_url_input.text().strip() + api_key = self.api_key_input.text().strip() + + if not base_url: + QMessageBox.warning(self, "缺少地址", "请先填写 API 地址 (base_url)") + return + + # 智谱 /models 需鉴权：未填 API 密钥时直接提示，避免 401 失败把已选模型清空 + if self._is_zhipu_selected() and not api_key: + QMessageBox.warning( + self, "缺少 API 密钥", + "智谱需要先在「API 密钥」框填入有效 Key 才能拉取模型列表。\n" + "也可直接从下拉选择内置模型（glm-4.6 等）或手动输入模型名。" + ) + return + + self.model_name_combo.clear() + self.model_name_combo.addItem("⏳ 正在拉取模型列表...") + + def task(): + models = fetch_available_models(base_url, api_key, timeout=10) + if models: + return True, f"成功获取 {len(models)} 个模型", models + return False, "未能获取模型（端点可能不支持 /models，可手动填写模型名）", [] + + self._ai_worker = _BackgroundAIWorker(task) + self._ai_worker.result_ready.connect(self._on_models_fetched) + self._ai_worker.start() + + @pyqtSlot(bool, str, object) + def _on_models_fetched(self, ok: bool, message: str, models): + """模型拉取完成回调（主线程）。""" + current = self.model_name_combo.currentText() self.model_name_combo.clear() - self.model_name_combo.addItem("正在加载模型...") + if ok and models: + self.model_name_combo.addItems(models) + # 尽量保留用户已填写的模型名 + if current and current not in ("⏳ 正在拉取模型列表...",): + self.model_name_combo.setCurrentText(current) + self.ai_status_label.setText(f"✅ {message}") + self.ai_status_label.setStyleSheet("color: #10b981;") + else: + # 拉取失败：智谱恢复内置模型列表，避免下拉被清空；其它 provider 仍可手填 + if self._is_zhipu_selected(): + self._populate_zhipu_models() + self.ai_status_label.setText( + "⚠️ 拉取失败（多为 API 密钥为空或无效）；已恢复内置模型列表，可直接选择或手动填写" + ) + else: + self.ai_status_label.setText(f"⚠️ {message}") + self.ai_status_label.setStyleSheet("color: #d97706;") + + def test_ai_connection(self): + """测试与 AI 端点的连接（后台线程，真实请求 GET /models）。""" + base_url = self.base_url_input.text().strip() + api_key = self.api_key_input.text().strip() + + if not base_url: + QMessageBox.warning(self, "缺少地址", "请先填写 API 地址 (base_url)") + return - # 模拟模型加载 - QTimer.singleShot(1000, self._load_models) + self.ai_status_label.setText("🔄 正在测试连接...") + self.ai_status_label.setStyleSheet("color: #6b7280;") + + def task(): + headers = {'Content-Type': 'application/json'} + if api_key: + headers['Authorization'] = f'Bearer {api_key}' + # 派生 /models 端点：去掉可能的 /chat/completions 后缀 + base = base_url.rstrip('/') + if base.endswith('/chat/completions'): + base = base[:-len('/chat/completions')] + models_url = base + '/models' + try: + resp = requests.get(models_url, headers=headers, timeout=10) + except requests.exceptions.Timeout: + return False, "请求超时，请检查地址与网络", 0 + except requests.exceptions.ConnectionError: + return False, "无法连接到服务器，请检查地址与端口", 0 + status = resp.status_code + if status == 200: + return True, "连接成功", status + if status == 401: + return False, "认证失败：API 密钥无效或缺失", status + if status == 404: + return False, "地址错误：/models 返回 404", status + return False, f"连接失败 (HTTP {status})", status + + self._ai_worker = _BackgroundAIWorker(task) + self._ai_worker.result_ready.connect(self._on_connection_tested) + self._ai_worker.start() + + @pyqtSlot(bool, str, object) + def _on_connection_tested(self, ok: bool, message: str, status): + """连接测试结果回调（主线程）。""" + if ok: + self.ai_status_label.setText(f"✅ {message}") + self.ai_status_label.setStyleSheet("color: #10b981;") + else: + self.ai_status_label.setText(f"❌ {message}") + self.ai_status_label.setStyleSheet("color: #ef4444;") - def _load_models(self): - """加载模型列表""" + def _load_ai_config_from_yaml(self): + """启动时从 config.yaml 回填 AI 配置到界面控件。""" try: - # 这里应该调用实际的模型加载逻辑 - models = [ - "deepseek-r1:1.5b", - "deepseek-r1:7b", - "deepseek-r1:14b", - "qwen2.5-coder:7b", - "llama3.1:8b", - "mistral:7b" - ] - - self.model_name_combo.clear() - self.model_name_combo.addItems(models) - self.model_name_combo.setCurrentIndex(0) + with open(self.config_path, 'r', encoding='utf-8') as f: + config = yaml.safe_load(f) or {} + except FileNotFoundError: + return + except Exception: + return # 配置不可读时静默，使用界面默认值 + + ai_cfg = config.get('ai', {}) or {} + cloud_provider = ai_cfg.get('cloud_provider') or ai_cfg.get('local_provider') or 'deepseek' + + # 反向映射 provider 名 -> 下拉文本 + label = "云端 DeepSeek" + for lbl, (prov, _typ, _url) in self._PROVIDER_MAP.items(): + if prov == cloud_provider: + label = lbl + break + idx = self.ai_model_type.findText(label) + if idx >= 0: + self.ai_model_type.setCurrentIndex(idx) + + # 按当前 provider 读取对应配置段 + section = config.get(cloud_provider, {}) or {} + # 智谱：先同步 Coding Plan 开关状态（影响默认 base_url 与模型列表的派生） + if cloud_provider == 'zhipu' and hasattr(self, 'zhipu_coding_plan_cb'): + self.zhipu_coding_plan_cb.setChecked(bool(section.get('coding_plan', False))) + if section.get('base_url'): + self.base_url_input.setText(str(section['base_url'])) + elif cloud_provider == 'zhipu': + # base_url 留空：按 coding_plan 状态填默认端点 + self.base_url_input.setText(self._zhipu_default_base_url()) + if section.get('api_key'): + self.api_key_input.setText(str(section['api_key'])) + if section.get('model'): + self.model_name_combo.setCurrentText(str(section['model'])) + + def _save_ai_config_to_yaml(self): + """把界面上的 AI 配置写回 config.yaml（按所选 provider 更新对应段，保留注释）。""" + provider, ai_type = self._current_ai_provider() + base_url = self.base_url_input.text().strip() + api_key = self.api_key_input.text().strip() + model = self.model_name_combo.currentText().strip() + + if not base_url and provider in ('custom', 'deepseek'): + QMessageBox.warning(self, "缺少地址", "请先填写 API 地址 (base_url)") + return - except Exception as e: - self.model_name_combo.clear() - self.model_name_combo.addItem(f"加载失败: {str(e)[:30]}...") + # 优先用 ruamel.yaml 做注释保留的 round-trip；不可用时回退到 PyYAML（会丢注释） + use_ruamel = False + try: + from ruamel.yaml import YAML + yaml_rt = YAML() + yaml_rt.preserve_quotes = True + with open(self.config_path, 'r', encoding='utf-8') as f: + config = yaml_rt.load(f) or {} + use_ruamel = True + except Exception: + try: + with open(self.config_path, 'r', encoding='utf-8') as f: + config = yaml.safe_load(f) or {} + except FileNotFoundError: + config = {} + except Exception as e: + QMessageBox.critical(self, "读取配置失败", str(e)) + return + + # 更新 ai.type / ai.cloud_provider / ai.local_provider + ai_cfg = config.setdefault('ai', {}) + ai_cfg['type'] = ai_type + if ai_type == 'cloud': + ai_cfg['cloud_provider'] = provider + else: + ai_cfg['local_provider'] = provider + + # 按所选 provider 更新对应配置段 + if provider in ('custom', 'deepseek'): + section = config.setdefault(provider, {}) + section['base_url'] = base_url + section['api_key'] = api_key + if model: + section['model'] = model + section.setdefault('timeout', 30) + section.setdefault('max_tokens', 2048) + elif provider == 'zhipu': + section = config.setdefault('zhipu', {}) + section['base_url'] = base_url + section['api_key'] = api_key + if model: + section['model'] = model + section['coding_plan'] = bool(self.zhipu_coding_plan_cb.isChecked()) + # GLM-5 等推理模型需要更长超时与更大 max_tokens；用户若已设更大值则保留 + section['timeout'] = max(int(section.get('timeout') or 0), 120) + section['max_tokens'] = max(int(section.get('max_tokens') or 0), 4096) + elif provider == 'ollama': + section = config.setdefault('ollama', {}) + if base_url: + section['base_url'] = base_url + if model: + section['model'] = model + section.setdefault('timeout', 60) + elif provider == 'lm_studio': + section = config.setdefault('lm_studio', {}) + if base_url: + section['base_url'] = base_url + if model: + section['model'] = model - def test_ai_connection(self): - """测试AI连接""" try: - self.ai_status_label.setText("🔄 正在测试...") - # 模拟连接测试 - QTimer.singleShot(2000, self._ai_test_result) + with open(self.config_path, 'w', encoding='utf-8') as f: + if use_ruamel: + yaml_rt.dump(config, f) + else: + yaml.safe_dump(config, f, default_flow_style=False, allow_unicode=True, sort_keys=False) except Exception as e: - self.ai_status_label.setText(f"❌ 测试失败: {str(e)}") + QMessageBox.critical(self, "保存配置失败", str(e)) + return - def _ai_test_result(self): - """AI测试结果""" - # 模拟测试成功 - self.ai_status_label.setText("✅ 连接测试成功") - self.ai_status_label.setStyleSheet("color: #10b981;") + QMessageBox.information( + self, "保存成功", + f"AI 配置已保存到 config.yaml\n(provider={provider}, type={ai_type})" + ) def clear_cache(self): """清空缓存""" @@ -913,6 +1298,9 @@ def start_analysis(self): QMessageBox.warning(self, "配置错误", "启用AI分析时需要指定模型名称") return + # 记录本次分析的日志路径，供导出报告使用 + self.last_log_path = config['log_path'] + # 禁用按钮 self.start_button.setEnabled(False) self.stop_button.setEnabled(True) @@ -1014,6 +1402,7 @@ def analysis_completed(self, success, message, results): self.stop_button.setEnabled(False) if success: + self.last_results = results # 保存结果，供「导出报告」使用 elapsed = time.time() - self.start_time self.append_log(f"✅ 分析完成！总耗时: {elapsed:.1f}秒") @@ -1032,18 +1421,98 @@ def analysis_completed(self, success, message, results): QMessageBox.critical(self, "分析失败", message) def export_report(self): - """导出报告""" + """把最近一次分析结果导出为 HTML 报告。""" + output_dir = self.output_dir_input.text().strip() + if not output_dir: + QMessageBox.warning(self, "导出失败", "请先配置输出目录") + return + if not self.last_results: + QMessageBox.warning(self, "导出失败", "请先运行一次分析，再导出报告") + return try: - output_dir = self.output_dir_input.text() - if not output_dir: - QMessageBox.warning(self, "导出失败", "请先配置输出目录") - return - - # 这里应该调用实际的报告导出功能 - QMessageBox.information(self, "导出成功", f"报告已导出到: {output_dir}") - + os.makedirs(output_dir, exist_ok=True) + html_doc = self._render_report_html(self.last_results, self.last_log_path) + fname = "sslogs-report-" + time.strftime("%Y%m%d-%H%M%S") + ".html" + path = os.path.join(output_dir, fname) + with open(path, "w", encoding="utf-8") as f: + f.write(html_doc) + n = len(self.last_results.get('threats', [])) + QMessageBox.information(self, "导出成功", f"已导出 {n} 条威胁的报告：\n{path}") + self.append_log(f"📄 报告已导出: {path}") except Exception as e: - QMessageBox.critical(self, "导出失败", f"导出报告时发生错误: {str(e)}") + QMessageBox.critical(self, "导出失败", f"导出报告时发生错误: {e}") + + def _render_report_html(self, results, log_path): + """根据分析结果渲染一份自包含的 HTML 报告。""" + import html as _html + from collections import Counter + esc = _html.escape + perf = results.get('performance', {}) or {} + threats = results.get('threats', []) or [] + ts = time.strftime("%Y-%m-%d %H:%M:%S") + + # 威胁类型统计 + type_count = Counter() + for t in threats: + for x in t.get('threats', []): + type_count[getattr(x, 'value', str(x))] += 1 + type_rows = "".join(f"{esc(k)}{v}" + for k, v in type_count.most_common()) or '无' + + # 威胁详情卡片 + cards = [] + for i, t in enumerate(threats, 1): + names = ", ".join(getattr(x, 'value', str(x)) for x in t.get('threats', [])) or "未知" + risk = float(t.get('risk_score', 0) or 0) + line = esc(t.get('line', '')) + ai = t.get('ai_analysis') + ai_html = (f'
{esc(ai)}
' if ai + else '
（未进行AI深度分析）
') + color = '#ef4444' if risk >= 8 else ('#f59e0b' if risk >= 5 else '#3b82f6') + cards.append( + f'
' + f'
#{i}' + f'{esc(names)}' + f'风险 {risk:.1f}
' + f'
{line}
{ai_html}
' + ) + detail_html = "".join(cards) or '

未检测到威胁 🎉

' + + return f""" + + +SSlogs 安全分析报告 + +

🛡️ SSlogs 安全分析报告

+
日志: {esc(log_path or '未知')}　|　生成时间: {ts}
+
+
{perf.get('processed_count', 0)}
处理日志条数
+
{len(threats)}
威胁条数
+
{perf.get('total_time', 0):.1f}s
耗时
+
{perf.get('memory_peak', 0):.1f}%
内存峰值
+
+

威胁类型分布

+{type_rows}
类型次数
+

威胁详情（{len(threats)} 条）

+{detail_html} +""" def main(): diff --git a/tests/unit/test_ai_analyzer.py b/tests/unit/test_ai_analyzer.py index 82ecc8d..9b38120 100644 --- a/tests/unit/test_ai_analyzer.py +++ b/tests/unit/test_ai_analyzer.py @@ -2,8 +2,9 @@ AI分析器单元测试 """ import pytest +import requests from unittest.mock import Mock, patch, MagicMock -from core.ai_analyzer import AIAnalyzer +from core.ai_analyzer import AIAnalyzer, fetch_available_models from core.exceptions import ( AIServiceError, AIServiceUnavailableError, @@ -460,3 +461,190 @@ def test_generic_prompt(self, analyzer): ) assert '攻击技术分析' in prompt assert '影响范围评估' in prompt + + +class TestCustomCloudProvider: + """自定义 OpenAI 兼容云端 provider 测试 (cloud_provider='custom')""" + + @pytest.fixture + def analyzer(self): + """创建配置为 custom provider 的分析器""" + config = { + 'ai': {'type': 'cloud', 'cloud_provider': 'custom'}, + 'custom': { + 'api_key': 'custom-key', + 'base_url': 'https://api.openai.com/v1', + 'model': 'gpt-4o-mini', + 'timeout': 15, + 'max_tokens': 500 + } + } + with patch('core.ai_analyzer.AIAnalyzer._load_config', return_value=config): + with patch('core.ai_analyzer.AIAnalyzer._init_http_session'): + return AIAnalyzer(config_path='dummy') + + def test_custom_provider_config_loaded(self, analyzer): + """custom provider 正确加载 custom: 配置段""" + assert analyzer.cloud_provider == 'custom' + assert analyzer.cloud_model == 'gpt-4o-mini' + assert analyzer.cloud_base_url == 'https://api.openai.com/v1' + assert analyzer.api_key == 'custom-key' + assert analyzer.cloud_max_tokens == 500 + assert analyzer.cloud_timeout == 15 + + def test_custom_chat_and_models_url_derivation(self, analyzer): + """base_url 为 base 形式时派生正确的 chat / models 端点""" + assert analyzer._chat_url() == 'https://api.openai.com/v1/chat/completions' + assert analyzer._models_url() == 'https://api.openai.com/v1/models' + + @patch('core.ai_analyzer.requests.post') + def test_analyze_with_custom_success(self, mock_post, analyzer): + """自定义 provider 成功调用云端分析（mock requests.post，避免真实网络）""" + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.raise_for_status.return_value = None + mock_response.json.return_value = { + 'choices': [{'message': {'content': '自定义分析结果'}}] + } + mock_post.return_value = mock_response + + result = analyzer._analyze_with_cloud('测试日志内容') + + assert result == '自定义分析结果' + mock_post.assert_called_once() + # 请求 URL 应为派生的 chat 端点 + assert mock_post.call_args.args[0] == 'https://api.openai.com/v1/chat/completions' + # payload 应使用 custom 的 model 与 max_tokens + payload = mock_post.call_args.kwargs['json'] + assert payload['model'] == 'gpt-4o-mini' + assert payload['max_tokens'] == 500 + + @patch('core.ai_analyzer.requests.post') + def test_analyze_with_custom_no_api_key(self, mock_post, analyzer): + """custom provider 缺少 api_key 时抛出认证错误""" + analyzer.api_key = '' + with pytest.raises(AIAuthenticationError): + analyzer._analyze_with_cloud('测试日志内容') + mock_post.assert_not_called() + + +class TestZhipuCloudProvider: + """智谱 GLM 云端 provider 测试 (cloud_provider='zhipu')""" + + @staticmethod + def _make_analyzer(coding_plan=False, base_url='', api_key='zhipu-key'): + """构造一个智谱 provider 分析器（patch 掉配置加载与 HTTP 会话初始化）""" + config = { + 'ai': {'type': 'cloud', 'cloud_provider': 'zhipu'}, + 'zhipu': { + 'api_key': api_key, + 'model': 'glm-4.6', + 'coding_plan': coding_plan, + 'base_url': base_url, + 'timeout': 20, + 'max_tokens': 800 + } + } + with patch('core.ai_analyzer.AIAnalyzer._load_config', return_value=config): + with patch('core.ai_analyzer.AIAnalyzer._init_http_session'): + return AIAnalyzer(config_path='dummy') + + @pytest.fixture + def analyzer(self): + """标准端点（coding_plan=False）的智谱分析器""" + return self._make_analyzer(coding_plan=False) + + def test_zhipu_provider_config_loaded(self, analyzer): + """智谱 provider 正确加载 zhipu: 配置段（标准端点）""" + assert analyzer.cloud_provider == 'zhipu' + assert analyzer.cloud_model == 'glm-4.6' + assert analyzer.api_key == 'zhipu-key' + assert analyzer.cloud_max_tokens == 800 + assert analyzer.cloud_timeout == 20 + # base_url 留空 + coding_plan=False -> 派生标准端点 + assert analyzer.cloud_base_url == 'https://open.bigmodel.cn/api/paas/v4' + # OpenAI 兼容的 Bearer 鉴权 + assert analyzer.cloud_headers['Authorization'] == 'Bearer zhipu-key' + + def test_zhipu_standard_chat_url(self, analyzer): + """标准端点派生正确的 chat / models 地址""" + assert analyzer._chat_url() == 'https://open.bigmodel.cn/api/paas/v4/chat/completions' + assert analyzer._models_url() == 'https://open.bigmodel.cn/api/paas/v4/models' + + def test_zhipu_coding_plan_endpoint(self): + """coding_plan=True 时 base_url 派生为编程套餐专属端点""" + a = self._make_analyzer(coding_plan=True) + assert a.cloud_base_url == 'https://open.bigmodel.cn/api/coding/paas/v4' + assert a._chat_url() == 'https://open.bigmodel.cn/api/coding/paas/v4/chat/completions' + + def test_zhipu_explicit_base_url_wins(self): + """显式配置的 base_url 优先于 coding_plan 派生值""" + a = self._make_analyzer(coding_plan=True, base_url='https://example.com/v1') + assert a.cloud_base_url == 'https://example.com/v1' + + @patch('core.ai_analyzer.requests.post') + def test_analyze_with_zhipu_success(self, mock_post, analyzer): + """智谱 provider 成功调用云端分析（mock requests.post，避免真实网络）""" + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.raise_for_status.return_value = None + mock_response.json.return_value = { + 'choices': [{'message': {'content': '智谱分析结果'}}] + } + mock_post.return_value = mock_response + + result = analyzer._analyze_with_cloud('测试日志内容') + + assert result == '智谱分析结果' + mock_post.assert_called_once() + # 请求 URL 为派生的标准端点 chat 地址 + assert mock_post.call_args.args[0] == 'https://open.bigmodel.cn/api/paas/v4/chat/completions' + payload = mock_post.call_args.kwargs['json'] + assert payload['model'] == 'glm-4.6' + + +class TestFetchAvailableModels: + """fetch_available_models 模块函数测试""" + + @patch('core.ai_analyzer.requests.get') + def test_fetch_models_success(self, mock_get): + """成功拉取模型列表，并带 Bearer 认证""" + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.raise_for_status.return_value = None + mock_response.json.return_value = { + 'data': [{'id': 'gpt-4o'}, {'id': 'gpt-4o-mini'}, {'no_id': True}] + } + mock_get.return_value = mock_response + + models = fetch_available_models('https://api.openai.com/v1', api_key='sk-xxx') + + assert models == ['gpt-4o', 'gpt-4o-mini'] + # URL 派生为 /models + assert mock_get.call_args.args[0] == 'https://api.openai.com/v1/models' + # 请求头带 Bearer + assert mock_get.call_args.kwargs['headers']['Authorization'] == 'Bearer sk-xxx' + + @patch('core.ai_analyzer.requests.get') + def test_fetch_models_from_full_chat_url(self, mock_get): + """base_url 为完整 chat 地址时也能正确派生 /models""" + mock_response = MagicMock() + mock_response.raise_for_status.return_value = None + mock_response.json.return_value = {'data': [{'id': 'm1'}]} + mock_get.return_value = mock_response + + models = fetch_available_models('https://api.siliconflow.cn/v1/chat/completions') + + assert models == ['m1'] + assert mock_get.call_args.args[0] == 'https://api.siliconflow.cn/v1/models' + + def test_fetch_models_empty_url(self): + """空 base_url 返回空列表""" + assert fetch_available_models('') == [] + + @patch('core.ai_analyzer.requests.get') + def test_fetch_models_failure_returns_empty(self, mock_get): + """网络/HTTP 错误时返回空列表，不抛异常""" + mock_get.side_effect = requests.exceptions.ConnectionError('boom') + assert fetch_available_models('https://api.openai.com/v1') == [] + diff --git a/tests/unit/test_config_manager.py b/tests/unit/test_config_manager.py index ebf141b..4cad189 100644 --- a/tests/unit/test_config_manager.py +++ b/tests/unit/test_config_manager.py @@ -207,6 +207,43 @@ def test_get_safe_config(self, valid_config_file): finally: os.unlink(temp_path) + def test_get_safe_config_custom(self): + """测试自定义端点 custom.api_key 同样被脱敏""" + config_content = """ +log_path: logs/*.log +log_format: + type: web + fields: + src_ip: (\\d+\\.\\d+\\.\\d+\\.\\d+) + +rule_dir: rules +output_dir: output + +custom: + api_key: sk-1234567890abcdef + base_url: https://api.openai.com/v1 + model: gpt-4o-mini +""" + with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f: + f.write(config_content) + temp_path = f.name + + try: + manager = ConfigManager(temp_path) + manager.load_config() + safe_config = manager.get_safe_config() + + # 自定义端点的 API 密钥也应被脱敏（保留首4/末4，中间以 * 替换） + masked = safe_config['custom']['api_key'] + assert masked != 'sk-1234567890abcdef' # 已脱敏，不再是明文 + assert masked.startswith('sk-1') and masked.endswith('cdef') + assert '*' in masked + # base_url 等非敏感字段保持不变 + assert 'https://api.openai.com/v1' == safe_config['custom']['base_url'] + finally: + os.unlink(temp_path) + + class TestConfigValidation: """配置验证测试""" @@ -280,6 +317,175 @@ def test_validate_enum_values(self): finally: os.unlink(temp_path) + def test_custom_cloud_provider_accepted(self): + """测试 cloud_provider=custom 且 custom 配置完整时通过校验""" + config_content = """ +log_path: logs/*.log +log_format: + type: web + fields: + src_ip: (\\d+\\.\\d+\\.\\d+\\.\\d+) + +rule_dir: rules +output_dir: output + +ai: + type: cloud + cloud_provider: custom + +custom: + base_url: https://api.openai.com/v1 + api_key: sk-1234567890abcdef + model: gpt-4o-mini + timeout: 30 +""" + with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f: + f.write(config_content) + temp_path = f.name + + try: + manager = ConfigManager(temp_path) + # 配置完整时不应抛出校验异常 + manager.load_config() + assert manager.get_config()['ai']['cloud_provider'] == 'custom' + finally: + os.unlink(temp_path) + + def test_custom_cloud_provider_missing_base_url_rejected(self): + """测试 cloud_provider=custom 但缺少 base_url 时被校验拦截""" + config_content = """ +log_path: logs/*.log +log_format: + type: web + fields: + src_ip: (\\d+\\.\\d+\\.\\d+\\.\\d+) + +rule_dir: rules +output_dir: output + +ai: + type: cloud + cloud_provider: custom + +custom: + base_url: "" + model: gpt-4o-mini +""" + with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f: + f.write(config_content) + temp_path = f.name + + try: + manager = ConfigManager(temp_path) + with pytest.raises(ConfigurationError) as exc_info: + manager.load_config() + assert 'base_url' in str(exc_info.value) + finally: + os.unlink(temp_path) + + def test_zhipu_cloud_provider_accepted(self): + """测试 cloud_provider=zhipu 且 zhipu 含 model 时通过校验（base_url 可留空）""" + config_content = """ +log_path: logs/*.log +log_format: + type: web + fields: + src_ip: (\\d+\\.\\d+\\.\\d+\\.\\d+) + +rule_dir: rules +output_dir: output + +ai: + type: cloud + cloud_provider: zhipu + +zhipu: + base_url: "" + api_key: zhipu-secret-key + model: glm-4.6 + coding_plan: false + timeout: 30 +""" + with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f: + f.write(config_content) + temp_path = f.name + + try: + manager = ConfigManager(temp_path) + manager.load_config() + assert manager.get_config()['ai']['cloud_provider'] == 'zhipu' + finally: + os.unlink(temp_path) + + def test_zhipu_cloud_provider_missing_model_rejected(self): + """测试 cloud_provider=zhipu 但缺少 model 时被校验拦截""" + config_content = """ +log_path: logs/*.log +log_format: + type: web + fields: + src_ip: (\\d+\\.\\d+\\.\\d+\\.\\d+) + +rule_dir: rules +output_dir: output + +ai: + type: cloud + cloud_provider: zhipu + +zhipu: + api_key: zhipu-key +""" + with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f: + f.write(config_content) + temp_path = f.name + + try: + manager = ConfigManager(temp_path) + with pytest.raises(ConfigurationError) as exc_info: + manager.load_config() + assert 'model' in str(exc_info.value) + finally: + os.unlink(temp_path) + + def test_zhipu_defaults_and_safe_config(self): + """_set_defaults 生成 zhipu 段，且 get_safe_config 对 zhipu.api_key 脱敏""" + config_content = """ +log_path: logs/*.log +log_format: + type: web + fields: + src_ip: (\\d+\\.\\d+\\.\\d+\\.\\d+) + +rule_dir: rules +output_dir: output +""" + with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f: + f.write(config_content) + temp_path = f.name + + try: + manager = ConfigManager(temp_path) + manager.load_config() + # _set_defaults 应已生成 zhipu 段默认值 + zhipu = manager.get_config().get('zhipu', {}) + assert zhipu.get('model') == 'glm-4.6' + assert zhipu.get('coding_plan') is False + assert zhipu.get('timeout') == 120 # GLM-5 推理模型较慢，默认更长超时 + assert zhipu.get('max_tokens') == 4096 + # 模拟用户填入真实密钥后脱敏 + manager._config.setdefault('zhipu', {}) + manager._config['zhipu']['api_key'] = 'sk-zhipu-1234567890abcdef' + + safe = manager.get_safe_config() + masked = safe['zhipu']['api_key'] + assert masked != 'sk-zhipu-1234567890abcdef' # 已脱敏 + assert masked.startswith('sk-z') and masked.endswith('cdef') + assert '*' in masked + finally: + os.unlink(temp_path) + + class TestConfigValidationMethods: """配置验证方法测试"""