MCP server: heaptrm as Model Context Protocol tool

4 tools for any MCP-compatible LLM agent:
heap_start — launch binary with instrumentation
heap_send — send data, get structured heap state
heap_check — force corruption validation
heap_observe — get current state

Works with Claude Code, SWE-agent, EnIGMA, or any MCP client.
Returns structured JSON with addresses, bins, risk score, corruptions.

heaptrm/mcp_config.json

@@ -0,0 +1,9 @@
+{
+  "mcpServers": {
+    "heaptrm": {
+      "command": "python3",
+      "args": ["heaptrm/mcp_server.py"],
+      "description": "Heap exploitation observability — structured heap state for exploit development"
+    }
+  }
+}

heaptrm/mcp_server.py

@@ -0,0 +1,166 @@
+#!/usr/bin/env python3
+"""
+heaptrm MCP server — structured heap observation for LLM agents.
+
+Provides 4 tools via Model Context Protocol:
+  - heap_start:   Launch a binary with heap instrumentation
+  - heap_send:    Send data to binary, get heap state back
+  - heap_check:   Force corruption validation (catches UAF/overflow)
+  - heap_observe:  Get current heap state without sending data
+
+Usage with Claude Code:
+  claude --mcp-server "python heaptrm/mcp_server.py"
+
+Usage with any MCP client:
+  Run as stdio server: python heaptrm/mcp_server.py
+"""
+
+import json
+import os
+import subprocess
+import signal
+import time
+from pathlib import Path
+from mcp.server.fastmcp import FastMCP
+
+mcp = FastMCP("heaptrm", instructions="""
+You have access to heaptrm, a heap exploitation observability tool.
+It instruments a binary's heap allocator and gives you structured
+observations: chunk addresses, sizes, states, freelist contents,
+corruption events, and exploit-readiness risk scores.
+
+Typical workflow:
+1. heap_start("./vulnerable_binary") — launch with instrumentation
+2. heap_send("1 0 64\\n") — send menu commands, get heap state
+3. heap_observe() — check current state
+4. heap_check() — validate for corruption after writes
+5. Use the addresses and bin info to compute your exploit
+
+The heap state includes safe-linking XOR keys derivable from chunk
+addresses (key = address >> 12). Use this for tcache poisoning.
+""")
+
+# Global state for the heaptrm subprocess
+_proc = None
+_last_state = None
+
+
+def _find_heaptrm_binary():
+    """Locate the compiled heaptrm Rust binary."""
+    candidates = [
+        Path(__file__).parent.parent / "heaptrm-cli" / "target" / "release" / "heaptrm",
+        Path("heaptrm-cli/target/release/heaptrm"),
+        Path("heaptrm"),
+    ]
+    for c in candidates:
+        if c.exists():
+            return str(c.resolve())
+    return None
+
+
+def _cmd(action: str, data: str = "") -> dict:
+    """Send a command to the heaptrm subprocess."""
+    global _proc, _last_state
+    if _proc is None:
+        return {"error": "No binary running. Call heap_start first."}
+
+    try:
+        _proc.stdin.write(json.dumps({"action": action, "data": data}) + "\n")
+        _proc.stdin.flush()
+        line = _proc.stdout.readline()
+        if not line:
+            return {"error": "heaptrm process closed"}
+        result = json.loads(line)
+        if result.get("heap"):
+            _last_state = result
+        return result
+    except Exception as e:
+        return {"error": str(e)}
+
+
+@mcp.tool()
+def heap_start(binary: str, args: str = "") -> str:
+    """Launch a binary with heap instrumentation.
+
+    Args:
+        binary: Path to the target binary
+        args: Optional space-separated arguments
+
+    Returns heap state after binary starts, or error message.
+    """
+    global _proc, _last_state
+
+    # Kill existing process
+    if _proc is not None:
+        try:
+            _proc.kill()
+            _proc.wait(timeout=2)
+        except:
+            pass
+        _proc = None
+        _last_state = None
+
+    heaptrm_bin = _find_heaptrm_binary()
+    if not heaptrm_bin:
+        return json.dumps({"error": "Cannot find heaptrm binary. Build with: cd heaptrm-cli && cargo build --release"})
+
+    cmd = [heaptrm_bin, binary]
+    if args:
+        cmd.extend(args.split())
+
+    try:
+        _proc = subprocess.Popen(
+            cmd,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.DEVNULL,
+            text=True,
+        )
+        time.sleep(0.1)
+        return json.dumps({"status": "started", "binary": binary, "pid": _proc.pid})
+    except Exception as e:
+        return json.dumps({"error": f"Failed to start: {e}"})
+
+
+@mcp.tool()
+def heap_send(data: str) -> str:
+    """Send data to the binary's stdin and get heap state back.
+
+    Args:
+        data: Data to send (e.g., "1 0 64\\n" for alloc). Must end with \\n.
+
+    Returns structured heap observation with chunks, bins, risk score,
+    addresses, and any corruption events.
+    """
+    if not data.endswith("\n"):
+        data += "\n"
+    result = _cmd("send", data)
+    return json.dumps(result, indent=2)
+
+
+@mcp.tool()
+def heap_check() -> str:
+    """Force heap validation to detect corruption from writes.
+
+    Call this after sending data that might corrupt the heap (edit/write
+    commands). Detects UAF writes, metadata corruption, and overflow.
+
+    Returns heap state with any newly detected corruptions.
+    """
+    result = _cmd("check")
+    return json.dumps(result, indent=2)
+
+
+@mcp.tool()
+def heap_observe() -> str:
+    """Get the current heap state without sending any data.
+
+    Returns the last observed heap state including chunks, bins,
+    risk score, and any detected corruption events.
+    """
+    result = _cmd("observe")
+    return json.dumps(result, indent=2)
+
+
+if __name__ == "__main__":
+    mcp.run(transport="stdio")