mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-06-04 13:28:03 +02:00
公明
228 lines
8.7 KiB
Python
228 lines
8.7 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Pent Claude Agent MCP Server - 渗透测试工程师 MCP 服务
|
||
|
||
通过 MCP 协议暴露 AI 渗透测试能力：CyberStrikeAI 可指挥 pent_claude_agent 执行渗透测试任务。
|
||
pent_claude_agent 内部使用 Claude Agent SDK，可独立配置 MCP、工具等，作为独立的渗透测试工程师运行。
|
||
|
||
依赖：pip install mcp claude-agent-sdk（或使用项目 venv）
|
||
运行：python mcp_pent_claude_agent.py [--config /path/to/config.yaml]
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import asyncio
|
||
import os
|
||
from typing import Any
|
||
|
||
import yaml
|
||
from mcp.server.fastmcp import FastMCP
|
||
|
||
# 延迟导入，避免未安装时影响 MCP 启动
|
||
_claude_sdk_available = False
|
||
try:
|
||
from claude_agent_sdk import ClaudeAgentOptions, query
|
||
|
||
_claude_sdk_available = True
|
||
except ImportError:
|
||
pass
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 路径与配置
|
||
# ---------------------------------------------------------------------------
|
||
|
||
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
PROJECT_ROOT = os.path.dirname(os.path.dirname(SCRIPT_DIR))
|
||
_DEFAULT_CONFIG_PATH = os.path.join(SCRIPT_DIR, "pent_claude_agent_config.yaml")
|
||
|
||
# Agent 运行状态（简单内存状态，用于 status）
|
||
_last_task: str | None = None
|
||
_last_result: str | None = None
|
||
_task_count: int = 0
|
||
|
||
|
||
def _load_config(config_path: str | None) -> dict[str, Any]:
|
||
"""加载 YAML 配置，合并默认值与用户配置。"""
|
||
defaults: dict[str, Any] = {
|
||
"cwd": PROJECT_ROOT,
|
||
"allowed_tools": ["Read", "Write", "Bash", "Grep", "Glob"],
|
||
"env": {
|
||
"CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1",
|
||
"DISABLE_TELEMETRY": "1",
|
||
"DISABLE_ERROR_REPORTING": "1",
|
||
"DISABLE_BUG_COMMAND": "1",
|
||
},
|
||
"mcp_servers": {},
|
||
"system_prompt": (
|
||
"你是一名专业的渗透测试工程师。根据用户给出的任务，进行安全测试、漏洞分析、信息收集等。"
|
||
"请按步骤执行，输出清晰、可复现的结果。仅在授权范围内进行测试。"
|
||
),
|
||
}
|
||
path = config_path or os.environ.get("PENT_CLAUDE_AGENT_CONFIG", _DEFAULT_CONFIG_PATH)
|
||
if not os.path.isfile(path):
|
||
return defaults
|
||
try:
|
||
with open(path, "r", encoding="utf-8") as f:
|
||
user = yaml.safe_load(f) or {}
|
||
# 深度合并
|
||
def merge(base: dict, override: dict) -> dict:
|
||
out = dict(base)
|
||
for k, v in override.items():
|
||
if k in out and isinstance(out[k], dict) and isinstance(v, dict):
|
||
out[k] = merge(out[k], v)
|
||
else:
|
||
out[k] = v
|
||
return out
|
||
|
||
return merge(defaults, user)
|
||
except Exception:
|
||
return defaults
|
||
|
||
|
||
def _resolve_path(s: str) -> str:
|
||
"""解析路径占位符。"""
|
||
return s.replace("${PROJECT_ROOT}", PROJECT_ROOT).replace("${SCRIPT_DIR}", SCRIPT_DIR)
|
||
|
||
|
||
def _build_agent_options(config: dict[str, Any], cwd_override: str | None = None) -> ClaudeAgentOptions:
|
||
"""从配置构建 ClaudeAgentOptions。"""
|
||
raw_cwd = cwd_override or config.get("cwd", PROJECT_ROOT)
|
||
cwd = _resolve_path(str(raw_cwd)) if isinstance(raw_cwd, str) else str(raw_cwd)
|
||
env = dict(os.environ)
|
||
env.update(config.get("env", {}))
|
||
mcp_servers = config.get("mcp_servers") or {}
|
||
# 解析路径占位符
|
||
for name, cfg in list(mcp_servers.items()):
|
||
if isinstance(cfg, dict):
|
||
args = cfg.get("args") or []
|
||
cfg = dict(cfg)
|
||
cfg["args"] = [_resolve_path(str(a)) for a in args]
|
||
mcp_servers[name] = cfg
|
||
|
||
return ClaudeAgentOptions(
|
||
cwd=cwd,
|
||
allowed_tools=config.get("allowed_tools", ["Read", "Write", "Bash", "Grep", "Glob"]),
|
||
disallowed_tools=config.get("disallowed_tools", []),
|
||
mcp_servers=mcp_servers,
|
||
env=env,
|
||
system_prompt=config.get("system_prompt"),
|
||
setting_sources=config.get("setting_sources", ["user", "project"]),
|
||
)
|
||
|
||
|
||
async def _run_claude_agent(prompt: str, config_path: str | None = None, cwd: str | None = None) -> str:
|
||
"""内部执行 Claude Agent，返回最后一轮文本结果。"""
|
||
global _last_task, _last_result, _task_count
|
||
_last_task = prompt
|
||
_task_count += 1
|
||
|
||
if not _claude_sdk_available:
|
||
_last_result = "错误：未安装 claude-agent-sdk，请执行 pip install claude-agent-sdk"
|
||
return _last_result
|
||
|
||
config = _load_config(config_path)
|
||
options = _build_agent_options(config, cwd_override=cwd)
|
||
|
||
messages: list[Any] = []
|
||
try:
|
||
async for message in query(prompt=prompt, options=options):
|
||
messages.append(message)
|
||
except Exception as e:
|
||
_last_result = f"Agent 执行异常: {e}"
|
||
return _last_result
|
||
|
||
if not messages:
|
||
_last_result = "(无输出)"
|
||
return _last_result
|
||
|
||
# 多轮迭代时，取最后一个 ResultMessage（最后一波结果）
|
||
result_msgs = [m for m in messages if hasattr(m, "result") and getattr(m, "result", None) is not None]
|
||
last = result_msgs[-1] if result_msgs else messages[-1]
|
||
# 提取文本内容，优先 ResultMessage.result，避免输出 metadata
|
||
if hasattr(last, "result") and last.result is not None:
|
||
text = last.result
|
||
elif hasattr(last, "content") and last.content:
|
||
parts = []
|
||
for block in last.content:
|
||
if hasattr(block, "text") and block.text:
|
||
parts.append(block.text)
|
||
text = "\n".join(parts) if parts else "(无输出)"
|
||
else:
|
||
text = "(无输出)"
|
||
_last_result = text
|
||
return _last_result
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# MCP 服务与工具
|
||
# ---------------------------------------------------------------------------
|
||
|
||
app = FastMCP(
|
||
name="pent-claude-agent",
|
||
instructions="渗透测试工程师 MCP：接收任务后，内部启动 Claude Agent 独立执行渗透测试、漏洞分析等，并返回结果。",
|
||
)
|
||
|
||
|
||
@app.tool(
|
||
description="执行渗透测试任务。下发任务描述后，pent_claude_agent 会作为独立的渗透测试工程师，使用 Claude Agent 执行任务并返回结果。支持：端口扫描、漏洞探测、Web 安全测试、信息收集等。",
|
||
)
|
||
async def pent_claude_run_pentest_task(task: str) -> str:
|
||
"""Run a penetration testing task. The agent executes independently and returns results."""
|
||
return await _run_claude_agent(task)
|
||
|
||
|
||
@app.tool(
|
||
description="分析漏洞信息。传入漏洞描述、PoC、影响范围等，由 Agent 进行专业分析并给出修复建议。",
|
||
)
|
||
async def pent_claude_analyze_vulnerability(vuln_info: str) -> str:
|
||
"""Analyze vulnerability information and provide remediation suggestions."""
|
||
prompt = f"请对以下漏洞信息进行专业分析，包括：风险等级、影响范围、利用方式、修复建议。\n\n{vuln_info}"
|
||
return await _run_claude_agent(prompt)
|
||
|
||
|
||
@app.tool(
|
||
description="执行指定任务。通用任务执行入口，Agent 会根据任务内容自动选择合适的工具和方法。",
|
||
)
|
||
async def pent_agent_execute(task: str) -> str:
|
||
"""Execute a task. The agent chooses appropriate tools and methods."""
|
||
return await _run_claude_agent(task)
|
||
|
||
|
||
@app.tool(
|
||
description="对目标进行安全诊断。可传入 URL、IP、域名等，Agent 会进行初步的安全评估和诊断。",
|
||
)
|
||
async def pent_agent_diagnose(target: str) -> str:
|
||
"""Diagnose a target (URL, IP, domain) for security assessment."""
|
||
prompt = f"请对以下目标进行安全诊断和初步评估：{target}\n\n包括：可达性、开放服务、常见漏洞面等。"
|
||
return await _run_claude_agent(prompt)
|
||
|
||
|
||
@app.tool(
|
||
description="获取 pent_claude_agent 的当前状态：最近任务、结果摘要、执行次数等。",
|
||
)
|
||
def pent_claude_status() -> str:
|
||
"""Get the current status of pent_claude_agent."""
|
||
global _last_task, _last_result, _task_count
|
||
lines = [
|
||
f"任务执行次数: {_task_count}",
|
||
f"最近任务: {_last_task or '-'}",
|
||
f"最近结果摘要: {(str(_last_result or '-')[:200] + '...') if _last_result and len(str(_last_result)) > 200 else (_last_result or '-')}",
|
||
f"Claude SDK 可用: {_claude_sdk_available}",
|
||
]
|
||
return "\n".join(lines)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser(description="Pent Claude Agent MCP Server")
|
||
parser.add_argument(
|
||
"--config",
|
||
default=None,
|
||
help="Path to pent_claude_agent config YAML (env: PENT_CLAUDE_AGENT_CONFIG)",
|
||
)
|
||
args, _ = parser.parse_known_args()
|
||
# 将 config 路径存入环境，供工具调用时使用
|
||
if args.config:
|
||
os.environ["PENT_CLAUDE_AGENT_CONFIG"] = args.config
|
||
app.run(transport="stdio")
|