mythos/mcp/exploit_generation_mcp.py

"""``exploit-generation-mcp`` — Pwntools + ROPGadget + heap + privesc."""

from __future__ import annotations

from ._mcp_runtime import MCPServer
from ..exploit.pwntools_synth import PwntoolsSynth
from ..exploit.rop_chain import ROPChainBuilder
from ..exploit.heap_exploit import HeapExploitKit
from ..exploit.privesc_kb import PrivEscKB

server = MCPServer("exploit-generation-mcp")
_pwn = PwntoolsSynth()
_rop = ROPChainBuilder()
_heap = HeapExploitKit()
_pe = PrivEscKB()


@server.tool("rop_chain", {"crash": "object"})
def rop_chain(crash: dict):
    return _rop.build(crash)


@server.tool("pwntools_assemble", {"crash": "object", "rop_chain": "array"})
def pwntools_assemble(crash: dict, rop_chain: list):
    return _pwn.assemble(crash, rop_chain)


@server.tool("heap_template", {"crash": "object"})
def heap_template(crash: dict):
    return _heap.spray_template(crash)


@server.tool("privesc_suggest", {"hypotheses": "array"})
def privesc_suggest(hypotheses: list):
    return _pe.suggest(hypotheses)


if __name__ == "__main__":  # pragma: no cover
    server.serve_stdio()