# PWN-MCP

面向 CTF Pwn 题的 AI 辅助调试平台。基于 MCP（Model Context Protocol）协议，让 AI 模型直接操作远程靶机上的 GDB、与 CTF 服务交互、分析堆状态，实现自动化堆利用。

## 架构

```
AI (MCP Client) ──stdio──▶ PWN-MCP Server ──SSH──▶ 远程靶机
                              │                        ├─ GDB + heap_logger.py
                              │                        ├─ CTF binary (socat)
                              │                        └─ /tmp/heap_*.json
                              └── 本地 IDA headless
```

## 安装

```bash
pip install -e .
```

依赖：`mcp`、`paramiko`

## 运行

```bash
# 以下三种方式均可
ctf-debugger
python -m src
python src/server.py
```

### MCP 客户端配置

```json
{
  "mcpServers": {
    "ctf-debugger": {
      "command": "python",
      "args": ["-m", "src"],
      "cwd": "/path/to/PWN-MCP"
    }
  }
}
```

## 模块总览

| 模块 | 文件 | 职责 |
|------|------|------|
| **MCP Server** | `src/server.py` | FastMCP 入口，注册全部工具组，stdio transport |
| **SSH Manager** | `src/ssh_manager.py` | 单例 SSH 连接管理（paramiko），提供 exec/upload/download |
| **GDB Controller** | `src/gdb_controller.py` | 单例 GDB 会话控制，通过 SSH PTY Channel 操控远程 GDB |
| **CTF Connector** | `src/ctf_connector.py` | 单例 CTF 服务连接，SSH direct-tcpip 端口转发，pwntools 风格接口 |
| **Heap Analyzer** | `src/heap_analyzer.py` | 解析 heap_logger 产生的 JSON（delta/snapshot），过滤与格式化 |
| **GDB Plugin** | `remote_scripts/heap_logger.py` | 上传到靶机的 GDB Python 插件，hook malloc/calloc/realloc/free |

## MCP Tools

### SSH 工具组（`src/tools/ssh_tools.py`）

| Tool | 说明 |
|------|------|
| `ssh_connect` | 建立 SSH 连接（密钥/密码认证） |
| `ssh_exec` | 远程执行 shell 命令 |
| `ssh_upload_file` | SFTP 上传文件 |
| `ssh_download_file` | SFTP 下载文件内容 |
| `ssh_disconnect` | 断开 SSH 连接 |

### GDB 工具组（`src/tools/gdb_tools.py`）

| Tool | 说明 |
|------|------|
| `gdb_start` | 启动 GDB 会话，可加载 heap_logger.py 插件 |
| `gdb_run` | 运行目标程序 |
| `gdb_continue` | 继续执行（前台） |
| `gdb_continue_background` | `continue&` 后台继续，用于请求粒度采集 |
| `gdb_interrupt` | Ctrl-C 中断程序 |
| `gdb_pause` | 暂停程序（优先 GDB `interrupt` 命令，兼容后台 continue 场景） |
| `gdb_restart` | kill + run 重启 |
| `gdb_stop` | 退出 GDB 会话 |
| `gdb_command` | 执行任意 GDB/pwndbg/gef 命令 |
| `gdb_memory` | 内存 dump（bytes + words 双视图，支持 GDB 表达式） |
| `gdb_chunk` | 查看 glibc malloc chunk header（prev_size/size + pwndbg malloc_chunk） |
| `gdb_source_window` | 围绕 source chunk 查看溢出窗口内存 |
| `gdb_read_output` | 读取 GDB 输出缓冲区（不等待 prompt） |

### Heap 分析工具组（`src/tools/heap_tools.py`）

**请求粒度追踪**

| Tool | 说明 |
|------|------|
| `heap_begin_action` | 标记堆操作阶段开始 |
| `heap_end_action` | 结束阶段，返回增量 delta（支持 caller/地址/类型过滤） |
| `heap_snapshot` | 获取完整堆状态快照（所有活跃 chunk） |
| `heap_get_state` | 快速查看堆统计概览 |
| `heap_query_delta` | 对最近一次 delta 做二次查询过滤 |
| `heap_query_log` | 查询完整文本日志 heaplog.txt |

**Bin 与 Chunk 追踪**

| Tool | 说明 |
|------|------|
| `heap_bins_json` | 导出 glibc bin 状态 JSON（fastbins/unsorted/smallbins/largebins） |
| `heap_find_chunk` | 查找 chunk 在哪个 bin 中 |
| `heap_chunk_bin` | 返回 chunk 所属 bin 描述 |
| `track_chunk` | 开始追踪 chunk 生命周期（header/bin/memory diff） |
| `untrack_chunk` | 停止追踪 chunk |
| `tracked_chunks_status` | 查看已追踪 chunk 状态 |

**请求级 Timeline**

| Tool | 说明 |
|------|------|
| `begin_request` | 开始网络请求粒度的 heap timeline |
| `end_request` | 结束请求 timeline |
| `request_timeline` | 读取请求级 heap timeline |
| `request_events` | 按 label 查询请求内 heap 事件 |

**高级能力**

| Tool | 说明 |
|------|------|
| `add_probe` / `remove_probe` / `list_probes` | 条件断点 + 自动采样 probe |
| `mem_snapshot` | 保存内存快照 |
| `mem_diff` | 比较两个内存快照（offset/old/new/u32/ascii） |
| `run_heap_sweep` | Payload sweep 自动化框架（批量测试配置，按停止条件终止） |

### CTF 交互工具组（`src/tools/ctf_tools.py`）

| Tool | 说明 |
|------|------|
| `ctf_connect` | 通过 SSH 端口转发连接 CTF 服务（自动接收 banner） |
| `ctf_send` | 发送数据（支持 hex 模式发送原始字节） |
| `ctf_sendline` | 发送一行文本（自动追加换行） |
| `ctf_recv` | 接收响应数据 |
| `ctf_recv_until` | 接收到指定 pattern |
| `ctf_close` | 关闭 CTF 连接 |
| `ctf_menu_action` | 组合工具：自动完成菜单式交互（选菜单 → 等 prompt → 填值） |

### IDA 静态分析工具组（`src/tools/ida_tools.py`）

| Tool | 说明 |
|------|------|
| `ida_headless_scan` | 调用本地 IDA headless 分析 ELF，输出 JSON 扫描结果 |
| `ida_find_functions` | 在扫描结果中按正则/子串搜索函数 |
| `ida_get_pseudocode` | 获取函数反编译伪代码 |
| `ida_get_xrefs` | 获取函数交叉引用（xrefs_to / calls_from） |

## 项目结构

```
PWN-MCP/
├── pyproject.toml              # 项目配置与依赖
├── README.md
├── src/
│   ├── __init__.py
│   ├── __main__.py             # python -m src 入口
│   ├── server.py               # MCP Server 入口，注册全部工具
│   ├── ssh_manager.py          # SSH 连接管理器（单例）
│   ├── gdb_controller.py       # GDB 会话控制器（单例，PTY Channel）
│   ├── ctf_connector.py        # CTF 服务连接器（单例，端口转发）
│   ├── heap_analyzer.py        # 堆日志 JSON 解析与格式化
│   └── tools/
│       ├── __init__.py
│       ├── ssh_tools.py        # SSH 工具组（5 tools）
│       ├── gdb_tools.py        # GDB 工具组（13 tools）
│       ├── heap_tools.py       # Heap 工具组（20 tools）
│       ├── ctf_tools.py        # CTF 工具组（7 tools）
│       └── ida_tools.py        # IDA 工具组（4 tools）
├── remote_scripts/
│   └── heap_logger.py          # GDB Python 插件（上传到靶机）
├── scripts/
│   └── verify_mcp_heap_debug.py
├── skills/
│   └── ctf-heap-spray.md       # AI Skill：堆喷射自动化工作流
└── docs/
    └── mcp_heap_debug_capabilities.md
```

## 典型工作流

```
ssh_connect → ssh_upload_file(heap_logger.py)
→ gdb_start(binary, plugin) → gdb_command("set disable-randomization on") → gdb_run
→ ctf_connect(127.0.0.1, port)
→ 迭代循环:
    heap_begin_action("spray")
    → gdb_continue → ctf_menu_action("1", [...]) → gdb_interrupt
    → heap_end_action()        # 分析 delta
    → heap_snapshot()          # 关键步骤验证布局
→ 发送 exploit payload
```

## GDB 插件（heap_logger.py）

上传到靶机后通过 `gdb -x heap_logger.py` 加载。特性：

- x86_64 / i386 ABI 感知的参数与返回值提取
- Hook malloc / calloc / realloc / free（raw + pending libc 断点）
- user pointer → chunk header 转换，chunk size 字段读取
- caller 解析到主程序调用点
- action-scoped delta JSON（`/tmp/heap_delta.json`）和全量 snapshot JSON（`/tmp/heap_snapshot.json`）
- 自定义 GDB 命令：`heap-begin`、`heap-end`、`heap-snapshot`、`heap-status`、`heap-bins-json`、`heap-find-chunk`、`track-chunk`、`add-probe-json` 等

## 技术栈

- **MCP SDK** — MCP 协议支持（FastMCP + stdio transport）
- **Paramiko** — SSH2 协议（连接、SFTP、PTY Channel、端口转发）
- **asyncio** — 异步运行时，通过 `asyncio.to_thread` 桥接同步调用
- **IDA Pro** — headless 模式静态分析（可选）