安全 Agent：自动化漏洞扫描与报告生成架构

from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated
import operator


class SecurityScanState(TypedDict):
    target: dict
    scan_plan: list[dict]
    raw_results: Annotated[list[dict], operator.add]
    analyzed_findings: Annotated[list[dict], operator.add]
    report: dict | None
    iteration_count: int
    max_iterations: int


workflow = StateGraph(SecurityScanState)

workflow.add_node("parse_target", parse_target_node)
workflow.add_node("plan_scan", plan_scan_node)
workflow.add_node("execute_scan", execute_scan_node)
workflow.add_node("analyze_results", analyze_results_node)
workflow.add_node("generate_report", generate_report_node)
workflow.add_node("notify_user", notify_user_node)

workflow.set_entry_point("parse_target")
workflow.add_edge("parse_target", "plan_scan")
workflow.add_edge("plan_scan", "execute_scan")
workflow.add_edge("execute_scan", "analyze_results")
workflow.add_conditional_edges(
    "analyze_results",
    should_continue_scanning,
    {
        "continue": "plan_scan",
        "complete": "generate_report",
    },
)
workflow.add_edge("generate_report", "notify_user")
workflow.add_edge("notify_user", END)

graph = workflow.compile()

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum


class ToolCategory(Enum):
    SCANNER = "scanner"
    QUERY = "query"
    REPORT = "report"
    NOTIFICATION = "notification"


@dataclass
class ToolResult:
    success: bool
    data: dict
    metadata: dict = field(default_factory=dict)
    error: str | None = None


class SecurityTool(ABC):
    name: str
    category: ToolCategory
    description: str
    required_permissions: list[str]

    @abstractmethod
    async def execute(self, params: dict, context: dict) -> ToolResult:
        ...

    @abstractmethod
    def validate_params(self, params: dict) -> bool:
        ...

    def get_schema(self) -> dict:
        return {
            "name": self.name,
            "category": self.category.value,
            "description": self.description,
            "required_permissions": self.required_permissions,
        }

import asyncio
import json
from pathlib import Path


class NucleiScanner(SecurityTool):
    name = "nuclei_scanner"
    category = ToolCategory.SCANNER
    description = "基于 Nuclei 模板的漏洞扫描器"
    required_permissions = ["network:scan"]

    def __init__(self, nuclei_bin: str = "nuclei"):
        self.nuclei_bin = nuclei_bin

    def validate_params(self, params: dict) -> bool:
        required = ["target"]
        return all(k in params for k in required)

    async def execute(self, params: dict, context: dict) -> ToolResult:
        target = params["target"]
        templates = params.get("templates", "")
        severity = params.get("severity", "low,medium,high,critical")
        rate_limit = params.get("rate_limit", 100)

        cmd = [
            self.nuclei_bin,
            "-target", target,
            "-severity", severity,
            "-rate-limit", str(rate_limit),
            "-json",
            "-silent",
        ]
        if templates:
            cmd.extend(["-t", templates])

        try:
            proc = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )
            stdout, stderr = await asyncio.wait_for(
                proc.communicate(), timeout=3600
            )

            findings = []
            for line in stdout.decode().strip().split("\n"):
                if line:
                    findings.append(json.loads(line))

            return ToolResult(
                success=proc.returncode == 0,
                data={"findings": findings, "count": len(findings)},
                metadata={"target": target, "templates_used": templates},
            )
        except asyncio.TimeoutError:
            return ToolResult(
                success=False, data={},
                error="扫描超时（超过 1 小时限制）",
            )

import httpx
from datetime import datetime


class CVEDatabaseQuery(SecurityTool):
    name = "cve_query"
    category = ToolCategory.QUERY
    description = "查询 CVE 漏洞详情、CVSS 评分和修复建议"
    required_permissions = ["cve:read"]

    def __init__(self, nvd_api_key: str | None = None):
        self.base_url = "https://services.nvd.nist.gov/rest/json/cves/2.0"
        self.api_key = nvd_api_key
        self.cache: dict[str, dict] = {}

    def validate_params(self, params: dict) -> bool:
        return "cve_id" in params

    async def execute(self, params: dict, context: dict) -> ToolResult:
        cve_id = params["cve_id"]

        if cve_id in self.cache:
            return ToolResult(success=True, data=self.cache[cve_id])

        headers = {}
        if self.api_key:
            headers["apiKey"] = self.api_key

        async with httpx.AsyncClient() as client:
            resp = await client.get(
                f"{self.base_url}",
                params={"cveId": cve_id},
                headers=headers,
                timeout=30,
            )

            if resp.status_code != 200:
                return ToolResult(
                    success=False, data={},
                    error=f"NVD API 返回 {resp.status_code}",
                )

            data = resp.json()
            if not data.get("vulnerabilities"):
                return ToolResult(
                    success=False, data={},
                    error=f"未找到 CVE: {cve_id}",
                )

            vuln = data["vulnerabilities"][0]["cve"]
            result = {
                "cve_id": cve_id,
                "description": vuln.get("descriptions", [{}])[0].get("value", ""),
                "published": vuln.get("published", ""),
                "last_modified": vuln.get("lastModified", ""),
                "severity": self._extract_severity(vuln),
                "references": [
                    ref.get("url", "")
                    for ref in vuln.get("references", [])
                ],
            }
            self.cache[cve_id] = result
            return ToolResult(success=True, data=result)

    def _extract_severity(self, vuln: dict) -> str:
        metrics = vuln.get("metrics", {})
        for key in ["cvssMetricV31", "cvssMetricV30", "cvssMetricV2"]:
            if key in metrics and metrics[key]:
                return metrics[key][0].get("cvssData", {}).get("baseSeverity", "UNKNOWN")
        return "UNKNOWN"

from jinja2 import Template


class ReportGenerator(SecurityTool):
    name = "report_generator"
    category = ToolCategory.REPORT
    description = "生成结构化的安全扫描报告"
    required_permissions = ["report:write"]

    REPORT_TEMPLATE = Template("""# {{ title }}

## 执行摘要

| 项目 | 内容 |
|------|------|
| 扫描目标 | {{ target }} |
| 扫描时间 | {{ scan_time }} |
| 发现漏洞总数 | {{ total_findings }} |
| 高危漏洞 | {{ critical_count }} |
| 中危漏洞 | {{ high_count }} |
| 低危漏洞 | {{ medium_count }} |

## 风险评估

{{ risk_summary }}

## 漏洞详情

{% for finding in findings %}
### {{ finding.severity }}: {{ finding.title }}

- **CVE 编号**: {{ finding.cve_id }}
- **CVSS 评分**: {{ finding.cvss_score }}
- **影响组件**: {{ finding.affected_component }}
- **描述**: {{ finding.description }}

**修复建议**: {{ finding.remediation }}

---
{% endfor %}

## 修复优先级建议

{{ remediation_priorities }}
""")

    def validate_params(self, params: dict) -> bool:
        return all(k in params for k in ["findings", "target", "scan_time"])

    async def execute(self, params: dict, context: dict) -> ToolResult:
        findings = params["findings"]
        severity_order = {"CRITICAL": 0, "HIGH": 1, "MEDIUM": 2, "LOW": 3}
        sorted_findings = sorted(
            findings,
            key=lambda f: severity_order.get(f.get("severity", "LOW"), 4),
        )

        counts = {"CRITICAL": 0, "HIGH": 0, "MEDIUM": 0, "LOW": 0}
        for f in findings:
            sev = f.get("severity", "LOW").upper()
            counts[sev] = counts.get(sev, 0) + 1

        rendered = self.REPORT_TEMPLATE.render(
            title=f"安全扫描报告 - {params['target']}",
            target=params["target"],
            scan_time=params["scan_time"],
            total_findings=len(findings),
            critical_count=counts.get("CRITICAL", 0),
            high_count=counts.get("HIGH", 0),
            medium_count=counts.get("MEDIUM", 0),
            findings=sorted_findings,
            risk_summary=self._generate_risk_summary(counts),
            remediation_priorities=self._prioritize_remediation(findings),
        )

        return ToolResult(
            success=True,
            data={"report_markdown": rendered, "finding_counts": counts},
        )

import httpx


class AlertNotifier(SecurityTool):
    name = "alert_notifier"
    category = ToolCategory.NOTIFICATION
    description = "发送扫描结果告警通知"
    required_permissions = ["notify:send"]

    def __init__(self):
        self.channels: dict[str, str] = {}

    def register_channel(self, name: str, webhook_url: str):
        self.channels[name] = webhook_url

    def validate_params(self, params: dict) -> bool:
        return all(k in params for k in ["channel", "message", "severity"])

    async def execute(self, params: dict, context: dict) -> ToolResult:
        channel = params["channel"]
        if channel not in self.channels:
            return ToolResult(
                success=False, data={},
                error=f"未注册的通知渠道: {channel}",
            )

        payload = self._build_payload(
            params["message"],
            params["severity"],
            params.get("findings_summary", {}),
        )

        async with httpx.AsyncClient() as client:
            resp = await client.post(
                self.channels[channel],
                json=payload,
                timeout=10,
            )
            return ToolResult(
                success=resp.status_code == 200,
                data={"status_code": resp.status_code},
            )

    def _build_payload(self, message: str, severity: str, summary: dict) -> dict:
        color_map = {
            "CRITICAL": "red",
            "HIGH": "orange",
            "MEDIUM": "yellow",
            "LOW": "green",
        }
        return {
            "msg_type": "interactive",
            "card": {
                "header": {
                    "title": {"tag": "plain_text", "content": f"安全扫描告警 [{severity}]"},
                    "template": color_map.get(severity, "blue"),
                },
                "elements": [
                    {"tag": "markdown", "content": message},
                    {"tag": "markdown", "content": f"**漏洞统计**: {summary}"},
                ],
            },
        }

[用户输入] 扫描 https://api.example.com 的安全性

[Step 1 - 目标分析]
  Thought: 目标是一个 HTTPS API 端点，我需要先确认其技术栈和开放端口
  Action: 调用目标解析工具，识别技术指纹
  Observation: Nginx 1.24, Python 3.11, REST API, 开放端口 443

[Step 2 - 扫描计划]
  Thought: 基于技术栈信息，需要使用 Web 漏洞模板、API 安全模板和 SSL 配置模板
  Action: 生成扫描计划
    - Phase 1: SSL/TLS 配置检查 (nuclei-templates/ssl/)
    - Phase 2: Web 漏洞扫描 (nuclei-templates/http/)
    - Phase 3: API 安全评估 (nuclei-templates/http/api/)

[Step 3 - 扫描执行]
  Action: 执行 Nuclei 扫描，使用选定模板
  Observation: 发现 12 个漏洞
    - CRITICAL: 1 (SQL Injection in /api/users)
    - HIGH: 3 (CORS misconfig, Missing Rate Limiting, Verbose Error)
    - MEDIUM: 4 (X-Frame-Options missing, etc.)
    - LOW: 4 (Information disclosure, etc.)

[Step 4 - 结果分析]
  Thought: SQL Injection 是最严重的发现，需要深入验证其可利用性
  Action: 对 SQL Injection 进行深度分析，查询 CVE 关联
  Observation: 该漏洞与 CVE-2024-XXXX 关联，CVSS 9.8，已有公开 PoC

[Step 5 - 报告生成]
  Thought: 所有分析完成，生成结构化报告
  Action: 调用报告生成工具，整合所有发现
  Observation: 报告生成成功，包含执行摘要、详细发现和修复建议

[Step 6 - 通知推送]
  Action: 发送告警通知到安全团队飞书群
  Observation: 通知发送成功

SECURITY_AGENT_SYSTEM_PROMPT = """你是一个专业的安全评估 Agent，负责自动化漏洞扫描和报告生成。

## 核心原则
1. 始终将扫描范围限制在用户明确授权的目标内
2. 不执行任何破坏性操作（如 DoS、数据删除）
3. 扫描频率必须遵守 rate_limit 约束
4. 对于高危漏洞，必须进行二次验证以排除误报

## 推理框架
每次决策前，按以下步骤思考：
1. 当前目标系统的特征是什么？
2. 哪些扫描策略最匹配当前场景？
3. 上一步的结果揭示了什么新信息？
4. 是否需要调整后续扫描计划？

## 输出规范
- 漏洞描述必须包含：影响范围、利用条件、修复建议
- 报告需区分技术详情和管理层摘要
- 所有判断必须有工具返回数据支撑，禁止臆测
"""

class ScanScopeGuard:
    def __init__(self, allowed_cidrs: list[str], allowed_domains: list[str]):
        self.allowed_networks = [ipaddress.ip_network(c) for c in allowed_cidrs]
        self.allowed_domains = allowed_domains

    def validate_target(self, target: str) -> bool:
        parsed = urlparse(target)
        hostname = parsed.hostname or ""

        if hostname in self.allowed_domains:
            return True

        try:
            ip = ipaddress.ip_address(hostname)
            return any(ip in net for net in self.allowed_networks)
        except ValueError:
            resolved = socket.getaddrinfo(hostname, None)
            for _, _, _, _, addr in resolved:
                try:
                    ip = ipaddress.ip_address(addr[0])
                    if any(ip in net for net in self.allowed_networks):
                        return True
                except ValueError:
                    continue

        return False

class ResultSanitizer:
    SENSITIVE_PATTERNS = [
        (r'password["\s:=]+\S+', 'password=***'),
        (r'api[_-]?key["\s:=]+\S+', 'api_key=***'),
        (r'token["\s:=]+\S+', 'token=***'),
        (r'\b(?:10|172\.(?:1[6-9]|2\d|3[01])|192\.168)\.\d{1,3}\.\d{1,3}\b',
         '[INTERNAL_IP]'),
    ]

    def sanitize(self, data: dict) -> dict:
        import re
        import json

        text = json.dumps(data)
        for pattern, replacement in self.SENSITIVE_PATTERNS:
            text = re.sub(pattern, replacement, text, flags=re.IGNORECASE)
        return json.loads(text)

from dataclasses import dataclass
from datetime import datetime


@dataclass
class EvaluationResult:
    coverage: float
    false_positive_rate: float
    false_negative_rate: float
    report_completeness: float
    remediation_accuracy: float
    end_to_end_time: float
    overall_score: float


class SecurityAgentEvaluator:
    REQUIRED_REPORT_FIELDS = [
        "executive_summary", "scope", "methodology",
        "findings", "risk_matrix", "remediation_plan",
    ]

    def evaluate(
        self,
        scan_results: list[dict],
        ground_truth: list[dict],
        report: dict,
        start_time: datetime,
        end_time: datetime,
    ) -> EvaluationResult:
        detected_ids = {f.get("vuln_id") for f in scan_results}
        truth_ids = {f.get("vuln_id") for f in ground_truth}

        true_positives = detected_ids & truth_ids
        false_positives = detected_ids - truth_ids
        false_negatives = truth_ids - detected_ids

        total_reported = len(detected_ids)
        coverage = len(true_positives) / len(truth_ids) if truth_ids else 1.0
        fp_rate = len(false_positives) / total_reported if total_reported else 0.0
        fn_rate = len(false_negatives) / len(truth_ids) if truth_ids else 0.0

        completeness = sum(
            1 for field in self.REQUIRED_REPORT_FIELDS if field in report
        ) / len(self.REQUIRED_REPORT_FIELDS)

        e2e_time = (end_time - start_time).total_seconds()

        overall = (
            coverage * 0.3
            + (1 - fp_rate) * 0.25
            + (1 - fn_rate) * 0.25
            + completeness * 0.1
            + (1 - min(e2e_time / 3600, 1.0)) * 0.1
        )

        return EvaluationResult(
            coverage=coverage,
            false_positive_rate=fp_rate,
            false_negative_rate=fn_rate,
            report_completeness=completeness,
            remediation_accuracy=0.0,
            end_to_end_time=e2e_time,
            overall_score=overall,
        )

services:
  api:
    build: ./api
    ports:
      - "8000:8000"
    environment:
      - DATABASE_URL=postgresql://user:pass@db:5432/security_agent
      - REDIS_URL=redis://redis:6379/0
      - OPENAI_API_KEY=${OPENAI_API_KEY}
    depends_on:
      - db
      - redis

  worker:
    build: ./worker
    environment:
      - DATABASE_URL=postgresql://user:pass@db:5432/security_agent
      - REDIS_URL=redis://redis:6379/0
    depends_on:
      - db
      - redis

  db:
    image: pgvector/pgvector:pg16
    volumes:
      - pgdata:/var/lib/postgresql/data
    environment:
      - POSTGRES_DB=security_agent
      - POSTGRES_USER=user
      - POSTGRES_PASSWORD=pass

  redis:
    image: redis:7-alpine

volumes:
  pgdata: