diff --git a/AGENTS.md b/AGENTS.md index 6d4dbb1..95973ab 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -94,10 +94,18 @@ utils.py Helpers: create_http_flow, replay_flows, save_flows, decode_body - `FlowStore` assigns monotonically increasing integer IDs (`mitmproxy_mcp_id`) to each `HTTPFlow`. - `ProxyManager.call()` is the only thread-safe way to invoke mitmproxy commands on the running event loop. - Replay and save use mitmproxy's native commands (`replay.client`, `save.file`) rather than reimplementing logic. -- `CaptureAddon` filters flows with `capture_filter` and a runtime-updatable list of `CaptureRule` objects (`include`/`exclude`). +- `CaptureAddon` filters flows with `capture_filter` and a runtime-updatable list of `CaptureRule` objects (`include`/`exclude`). It tags each captured flow with `metadata["mitmproxy_mcp_source_proxy"]` to record which proxy instance captured it. - `RulesAddon` runs inside the mitmproxy event loop; its rule list is protected by an `RLock` and can be updated from the MCP tool thread. - `CryptoAddon` runs inside the mitmproxy event loop and applies user-loaded `CryptoHandler` scripts to decrypt/encrypt HTTP/WebSocket traffic. +### Auxiliary proxy (dual-proxy) + +The server supports an optional auxiliary proxy (`aux_proxy_manager`) for chained proxy scenarios where encryption/decryption is split across two mitmproxy instances. Both proxies share the same `FlowStore` but have independent `CaptureAddon`, `CryptoAddon`, `RulesAddon` and `WebSocketRulesAddon` instances. + +Tools that operate on flows (`flow_action` replay/resume/kill, `websocket_ctl` inject) automatically route to the correct proxy based on `flow.metadata["mitmproxy_mcp_source_proxy"]`. Other tools accept a `proxy_id` parameter (`"main"` or `"aux"`, default `"main"`) for explicit routing. + +The auxiliary proxy is entirely optional — it is only used when the user calls `proxy_ctl(cmd="start", proxy_id="aux", ...)`. + ## Automatic rules The server supports automatic rules via `rule_ctl`. A rule consists of: @@ -336,13 +344,14 @@ uv pip install -e ".[dev]" | Tool | Commands | |------|----------| -| `proxy_ctl(cmd, ...)` | `start`, `stop`, `status`, `list_options`, `clear_all`, `wireguard_config` | +| `proxy_ctl(cmd, proxy_id, ...)` | `start`, `stop`, `status`, `list_options`, `clear_all`, `wireguard_config` | | `ca_ctl(cmd, ...)` | `status`, `export_ca`, `set_verify_upstream`, `set_upstream_ca`, `clear_upstream_ca`, `set_client_cert`, `clear_client_cert` | -| `websocket_ctl(cmd, ...)` | `list`, `get`, `inject`, `connect`, `list_rules`, `add_rule`, `delete_rule`, `clear_rules` | +| `websocket_ctl(cmd, proxy_id, ...)` | `list`, `get`, `inject`, `connect`, `list_rules`, `add_rule`, `delete_rule`, `clear_rules` | | `http_ctl(cmd, ...)` | `list`, `get`, `delete`, `clear`, `load`, `save`, `extract_json` | | `flow_action(action, ...)` | `replay`, `resume`, `kill`, `update`, `create`, `send` | -| `rule_ctl(cmd, ...)` | `list`, `add`, `delete`, `clear` | -| `capture_rule_ctl(cmd, ...)` | `list`, `add`, `delete`, `clear` | +| `crypt_ctl(cmd, proxy_id, ...)` | `list`, `load`, `unload`, `reload`, `status` | +| `rule_ctl(cmd, proxy_id, ...)` | `list`, `add`, `delete`, `clear` | +| `capture_rule_ctl(cmd, proxy_id, ...)` | `list`, `add`, `delete`, `clear` | | `mock_server_ctl(cmd, ...)` | `start`, `add`, `stop`, `status` | | `map_local_ctl(cmd, ...)` | `list`, `add`, `delete`, `clear` | | `map_remote_ctl(cmd, ...)` | `list`, `add`, `delete`, `clear` | diff --git a/README.md b/README.md index ad3fadf..29c90a3 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,7 @@ - **查看**: `http_ctl(cmd="list")`, `http_ctl(cmd="get")` - **重放**: `flow_action(action="replay")`, `flow_action(action="send")` —— 基于 mitmproxy 原生 `replay.client` - **修改**: `flow_action(action="update")`, `flow_action(action="create")` +- **辅助代理(可选)** 支持同时运行第二个 mitmproxy 实例,用于链式代理场景下的加解密分工。 - **基于 mitmproxy 自身引擎** 实现重放和保存,不重复造轮子。 - **stdio 传输**,开箱兼容 Claude Desktop。 - **SSE 传输**,可远程或网络客户端连接(Claude Code、Cursor 等)。 @@ -351,6 +352,63 @@ class MyHandler(CryptoHandler): - 在 `decrypt_request` 中查询历史 handshake 流量来推导会话密钥。 - 返回 `CryptoResult(error="...")` 在 `crypt_ctl status` 中向 LLM 报告原因。 +## 辅助代理(双代理链式加解密) + +支持同时运行第二个 mitmproxy 实例(辅助代理),适用于链式代理场景下的加解密分工。 + +### 典型场景 + +``` +客户端 → mitmA (端口 8080) → Burp/其他代理 → mitmB (端口 8082) → 服务器 +``` + +- **mitmA**(主代理):解密客户端请求,加密返回给客户端的响应 +- **mitmB**(辅助代理):加密发往服务器的请求,解密服务器响应 + +两个代理共享同一个 FlowStore,捕获的流量统一在 `http_ctl` 中查看。 + +### 使用方式 + +```python +# 1. 启动主代理 +proxy_ctl(cmd="start", port=8080) + +# 2. 启动辅助代理(不同端口) +proxy_ctl(cmd="start", proxy_id="aux", port=8082) + +# 3. 分别加载加解密脚本 +crypt_ctl(cmd="load", script_path="/path/to/decrypt_a.py") # 主代理 +crypt_ctl(cmd="load", proxy_id="aux", script_path="/path/to/encrypt_b.py") # 辅助代理 + +# 4. 查看状态(自动合并显示两个代理) +proxy_ctl(cmd="status") + +# 5. 停止辅助代理 +proxy_ctl(cmd="stop", proxy_id="aux") +``` + +### 支持 proxy_id 的工具 + +以下工具通过 `proxy_id`(`"main"` 或 `"aux"`,默认 `"main"`)路由到指定代理: + +| 工具 | 路由行为 | +|------|----------| +| `proxy_ctl` | `start`/`stop`/`clear_all`/`wireguard_config` 按 proxy_id 路由;`status` 自动合并显示 | +| `crypt_ctl` | 加解密脚本按 proxy_id 隔离,互不干扰 | +| `rule_ctl` | 自动规则按 proxy_id 隔离 | +| `capture_rule_ctl` | 捕获规则按 proxy_id 隔离 | +| `websocket_ctl` | `inject` 自动按 flow 来源路由;`connect`/规则操作按 proxy_id 路由 | + +以下工具不需要 `proxy_id`,始终操作共享数据: + +| 工具 | 说明 | +|------|------| +| `http_ctl` | 操作共享 FlowStore,两个代理的流量统一查看 | +| `flow_action` | `replay`/`resume`/`kill` 自动按 flow 来源路由到正确的代理 | +| `flow_action(update/create/send)` | 操作 FlowStore 数据,不涉及代理路由 | +| `mock_server_ctl` | 始终在主代理上运行 | +| `map_local_ctl` / `map_remote_ctl` | 始终在主代理上运行 | + ## MCP Resources 除 tools 外,服务器还暴露一组只读的 MCP resources,客户端可以像读取文件一样直接获取状态,减少 tool 调用次数: @@ -383,14 +441,14 @@ class MyHandler(CryptoHandler): | 工具 | 命令 / 说明 | |------|------------| -| `proxy_ctl(cmd, ...)` | `start`, `stop`, `status`, `list_options`, `clear_all`, `wireguard_config` | +| `proxy_ctl(cmd, proxy_id, ...)` | `start`, `stop`, `status`, `list_options`, `clear_all`, `wireguard_config` | | `ca_ctl(cmd, ...)` | `status`, `export_ca`, `set_verify_upstream`, `set_upstream_ca`, `clear_upstream_ca`, `set_client_cert`, `clear_client_cert` | -| `websocket_ctl(cmd, ...)` | `list`, `get`, `inject`, `connect`, `list_rules`, `add_rule`, `delete_rule`, `clear_rules` | +| `websocket_ctl(cmd, proxy_id, ...)` | `list`, `get`, `inject`, `connect`, `list_rules`, `add_rule`, `delete_rule`, `clear_rules` | | `http_ctl(cmd, ...)` | `list`, `get`, `delete`, `clear`, `load`, `save`, `extract_json`, `export_har`, `import_har` | | `flow_action(action, ...)` | `replay`, `resume`, `kill`, `update`, `create`, `send` | -| `crypt_ctl(cmd, ...)` | `list`, `load`, `unload`, `reload`, `status`(用户自定义加解密脚本) | -| `rule_ctl(cmd, ...)` | `list`, `add`, `delete`, `clear`(自动规则) | -| `capture_rule_ctl(cmd, ...)` | `list`, `add`, `delete`, `clear`(捕获 include/exclude 规则) | +| `crypt_ctl(cmd, proxy_id, ...)` | `list`, `load`, `unload`, `reload`, `status`(用户自定义加解密脚本) | +| `rule_ctl(cmd, proxy_id, ...)` | `list`, `add`, `delete`, `clear`(自动规则) | +| `capture_rule_ctl(cmd, proxy_id, ...)` | `list`, `add`, `delete`, `clear`(捕获 include/exclude 规则) | | `mock_server_ctl(cmd, ...)` | `start`, `add`, `stop`, `status` | | `map_local_ctl(cmd, ...)` | `list`, `add`, `delete`, `clear`(URL → 本地文件) | | `map_remote_ctl(cmd, ...)` | `list`, `add`, `delete`, `clear`(URL 重写) | diff --git a/README_EN.md b/README_EN.md index 6ffbc25..9d65794 100644 --- a/README_EN.md +++ b/README_EN.md @@ -11,6 +11,7 @@ A lightweight [Model Context Protocol (MCP)](https://modelcontextprotocol.io) se - **View**: `http_ctl(cmd="list")`, `http_ctl(cmd="get")` - **Replay**: `flow_action(action="replay")`, `flow_action(action="send")` — backed by mitmproxy's native `replay.client` - **Modify**: `flow_action(action="update")`, `flow_action(action="create")` +- **Auxiliary proxy (optional)** — run a second mitmproxy instance for split encryption/decryption in chained proxy setups. - **Built on mitmproxy's own engine** for replay and save, so we don't reinvent the wheel. - **stdio transport** for out-of-the-box Claude Desktop compatibility. - **SSE transport** for remote or network-based MCP clients (Claude Code, Cursor, etc.). @@ -351,6 +352,63 @@ See `examples/crypto_xor_example.py` (simple XOR) and `examples/crypto_dynamic_k - Look up a previous handshake flow in `decrypt_request` to derive a session key. - Return `CryptoResult(error="...")` so the reason surfaces in `crypt_ctl status`. +## Auxiliary proxy (dual-proxy chained encryption) + +Supports running a second mitmproxy instance (auxiliary proxy) for split encryption/decryption in chained proxy setups. + +### Typical scenario + +``` +Client → mitmA (port 8080) → Burp/other tool → mitmB (port 8082) → Server +``` + +- **mitmA** (main proxy): decrypts client requests, encrypts responses back to client +- **mitmB** (auxiliary proxy): encrypts requests to server, decrypts server responses + +Both proxies share the same FlowStore; all captured traffic is visible in `http_ctl`. + +### Usage + +```python +# 1. Start main proxy +proxy_ctl(cmd="start", port=8080) + +# 2. Start auxiliary proxy (different port) +proxy_ctl(cmd="start", proxy_id="aux", port=8082) + +# 3. Load encryption scripts separately +crypt_ctl(cmd="load", script_path="/path/to/decrypt_a.py") # main +crypt_ctl(cmd="load", proxy_id="aux", script_path="/path/to/encrypt_b.py") # aux + +# 4. View status (auto-merged for both proxies) +proxy_ctl(cmd="status") + +# 5. Stop auxiliary proxy +proxy_ctl(cmd="stop", proxy_id="aux") +``` + +### Tools that support proxy_id + +These tools route to the specified proxy via `proxy_id` (`"main"` or `"aux"`, default `"main"`): + +| Tool | Routing behavior | +|------|------------------| +| `proxy_ctl` | `start`/`stop`/`clear_all`/`wireguard_config` route by proxy_id; `status` auto-merges | +| `crypt_ctl` | Crypto scripts are isolated per proxy_id | +| `rule_ctl` | Automatic rules are isolated per proxy_id | +| `capture_rule_ctl` | Capture rules are isolated per proxy_id | +| `websocket_ctl` | `inject` auto-routes by flow source; `connect`/rules route by proxy_id | + +These tools operate on shared data and do not need proxy_id: + +| Tool | Notes | +|------|-------| +| `http_ctl` | Operates on shared FlowStore; traffic from both proxies is visible | +| `flow_action` | `replay`/`resume`/`kill` auto-route to the correct proxy based on flow source | +| `flow_action(update/create/send)` | Operates on FlowStore data, no proxy routing | +| `mock_server_ctl` | Always runs on the main proxy | +| `map_local_ctl` / `map_remote_ctl` | Always runs on the main proxy | + ## MCP Resources In addition to tools, the server exposes read-only MCP resources that clients can read like files, reducing the need for repeated tool calls: @@ -383,14 +441,14 @@ Read mitmproxy://ca/status to see the CA/certificate configuration | Tool | Commands / Description | |------|------------------------| -| `proxy_ctl(cmd, ...)` | `start`, `stop`, `status`, `list_options`, `clear_all`, `wireguard_config` | +| `proxy_ctl(cmd, proxy_id, ...)` | `start`, `stop`, `status`, `list_options`, `clear_all`, `wireguard_config` | | `ca_ctl(cmd, ...)` | `status`, `export_ca`, `set_verify_upstream`, `set_upstream_ca`, `clear_upstream_ca`, `set_client_cert`, `clear_client_cert` | -| `websocket_ctl(cmd, ...)` | `list`, `get`, `inject`, `connect`, `list_rules`, `add_rule`, `delete_rule`, `clear_rules` | +| `websocket_ctl(cmd, proxy_id, ...)` | `list`, `get`, `inject`, `connect`, `list_rules`, `add_rule`, `delete_rule`, `clear_rules` | | `http_ctl(cmd, ...)` | `list`, `get`, `delete`, `clear`, `load`, `save`, `extract_json`, `export_har`, `import_har` | | `flow_action(action, ...)` | `replay`, `resume`, `kill`, `update`, `create`, `send` | -| `crypt_ctl(cmd, ...)` | `list`, `load`, `unload`, `reload`, `status` (user-defined encryption/decryption scripts) | -| `rule_ctl(cmd, ...)` | `list`, `add`, `delete`, `clear` (automatic rules) | -| `capture_rule_ctl(cmd, ...)` | `list`, `add`, `delete`, `clear` (capture include/exclude rules) | +| `crypt_ctl(cmd, proxy_id, ...)` | `list`, `load`, `unload`, `reload`, `status` (user-defined encryption/decryption scripts) | +| `rule_ctl(cmd, proxy_id, ...)` | `list`, `add`, `delete`, `clear` (automatic rules) | +| `capture_rule_ctl(cmd, proxy_id, ...)` | `list`, `add`, `delete`, `clear` (capture include/exclude rules) | | `mock_server_ctl(cmd, ...)` | `start`, `add`, `stop`, `status` | | `map_local_ctl(cmd, ...)` | `list`, `add`, `delete`, `clear` (URL → local file) | | `map_remote_ctl(cmd, ...)` | `list`, `add`, `delete`, `clear` (URL rewrite) | diff --git a/src/mitmproxy_mcp/proxy.py b/src/mitmproxy_mcp/proxy.py index 281ae60..2a2bb73 100644 --- a/src/mitmproxy_mcp/proxy.py +++ b/src/mitmproxy_mcp/proxy.py @@ -72,6 +72,7 @@ def __init__( capture_filter: str | None = None, capture_rules: list[CaptureRule] | None = None, event_buffer: EventBuffer | None = None, + source_proxy: str = "main", ) -> None: self.store = store self.capture_filter = capture_filter @@ -79,6 +80,7 @@ def __init__( self._lock = threading.RLock() self._capture_rules: list[CaptureRule] = [] self._event_buffer = event_buffer + self._source_proxy = source_proxy self._compile_filter() if capture_rules: self.set_rules(capture_rules) @@ -189,12 +191,18 @@ def _emit_flow_captured(self, store_id: int, flow: http.HTTPFlow) -> None: def response(self, flow: http.HTTPFlow) -> None: if self.should_capture(flow): + if flow.metadata is None: + flow.metadata = {} + flow.metadata["mitmproxy_mcp_source_proxy"] = self._source_proxy store_id = self.store.add(flow) self._emit_flow_captured(store_id, flow) def error(self, flow: http.HTTPFlow) -> None: # Also capture failed flows so errors are visible. if self.should_capture(flow): + if flow.metadata is None: + flow.metadata = {} + flow.metadata["mitmproxy_mcp_source_proxy"] = self._source_proxy store_id = self.store.add(flow) self._emit_flow_captured(store_id, flow) @@ -207,6 +215,9 @@ def websocket_start(self, flow: http.HTTPFlow) -> None: # ensure the WebSocket flow is tracked in case filters behave # differently at upgrade time. if self.should_capture(flow): + if flow.metadata is None: + flow.metadata = {} + flow.metadata["mitmproxy_mcp_source_proxy"] = self._source_proxy store_id = self.store.add(flow) self._emit_flow_captured(store_id, flow) if self._event_buffer is not None: @@ -259,7 +270,7 @@ def to_options(self) -> dict[str, Any]: class ProxyManager: """Manages a mitmproxy DumpMaster or WebMaster running in a background thread.""" - def __init__(self, store: FlowStore) -> None: + def __init__(self, store: FlowStore, source_proxy: str = "main") -> None: self.store = store self._master: DumpMaster | WebMaster | None = None self._thread: threading.Thread | None = None @@ -270,7 +281,7 @@ def __init__(self, store: FlowStore) -> None: self._wireguard_config: str | None = None self._ca_config = CaConfig() self.event_buffer = EventBuffer() - self.capture_addon = CaptureAddon(self.store, event_buffer=self.event_buffer) + self.capture_addon = CaptureAddon(self.store, event_buffer=self.event_buffer, source_proxy=source_proxy) self.rules_addon = RulesAddon(event_buffer=self.event_buffer) self.websocket_rules_addon = WebSocketRulesAddon(event_buffer=self.event_buffer) self.crypto_addon = CryptoAddon(self.store, event_buffer=self.event_buffer) @@ -843,8 +854,13 @@ def inject_websocket( to_client: bool, message: str, binary: bool = False, + target: ProxyManager | None = None, ) -> dict[str, Any]: - """Inject a message into an existing WebSocket connection.""" + """Inject a message into an existing WebSocket connection. + + If *target* is provided, the inject command runs in that proxy's event + loop (useful when the flow was captured by a different proxy instance). + """ flow = self.store.get(store_id) if flow is None: return {"success": False, "error": f"Flow with id {store_id} not found"} @@ -852,8 +868,9 @@ def inject_websocket( return {"success": False, "error": "Flow is not a WebSocket connection"} msg_bytes = message.encode("utf-8") if not binary else base64.b64decode(message) + proxy = target or self try: - self.call("inject.websocket", flow, to_client, msg_bytes, not binary) + proxy.call("inject.websocket", flow, to_client, msg_bytes, not binary) except Exception as e: return {"success": False, "error": f"Failed to inject message: {e}"} return {"success": True, "injected": 1} diff --git a/src/mitmproxy_mcp/server.py b/src/mitmproxy_mcp/server.py index 96315a1..439b92d 100644 --- a/src/mitmproxy_mcp/server.py +++ b/src/mitmproxy_mcp/server.py @@ -54,7 +54,8 @@ mcp = FastMCP("mitmproxy-mcp") store = FlowStore() -proxy_manager = ProxyManager(store) +proxy_manager = ProxyManager(store, source_proxy="main") +aux_proxy_manager = ProxyManager(store, source_proxy="aux") # ============================================================================= @@ -83,7 +84,10 @@ def config_rules() -> dict[str, Any]: @mcp.resource(EVENTS_LATEST_URI, name="Latest Events", mime_type="application/json") def events_latest() -> list[dict[str, Any]]: """Recent internal events (proxy lifecycle, captured flows, rule matches, crypto errors).""" - return events_latest_resource(proxy_manager.event_buffer) + main_events = events_latest_resource(proxy_manager.event_buffer, limit=10) + aux_events = events_latest_resource(aux_proxy_manager.event_buffer, limit=10) + combined = sorted(main_events + aux_events, key=lambda e: e.get("timestamp", 0), reverse=True) + return combined[:10] @mcp.resource(CRYPTO_SCRIPTS_URI, name="Crypto Scripts", mime_type="application/json") @@ -113,6 +117,20 @@ def ca_status() -> dict[str, Any]: # ============================================================================= +def _get_proxy_by_id(proxy_id: str) -> ProxyManager: + """Return the proxy manager for the given id ('main' or 'aux').""" + if proxy_id == "aux": + return aux_proxy_manager + return proxy_manager + + +def _get_source_proxy_for_flow(flow: http.HTTPFlow) -> str: + """Return the proxy id that captured the given flow (defaults to 'main').""" + if flow.metadata: + return flow.metadata.get("mitmproxy_mcp_source_proxy", "main") + return "main" + + def _get_flow_or_raise(flow_id: int) -> http.HTTPFlow: flow = store.get(flow_id) if flow is None: @@ -131,6 +149,7 @@ def _get_flows_by_ids(flow_ids: list[int]) -> list[http.HTTPFlow]: def _proxy_start( + target_proxy: ProxyManager | None = None, host: str = "127.0.0.1", port: int = 8080, capture_filter: str | None = None, @@ -140,7 +159,8 @@ def _proxy_start( webui: bool = False, web_port: int = 8081, ) -> dict[str, Any]: - return proxy_manager.start( + target = target_proxy or proxy_manager + return target.start( host=host, port=port, capture_filter=capture_filter, @@ -404,33 +424,39 @@ def _request_send( def _flow_replay(flow_id: int, use_modified: bool = True) -> dict[str, Any]: flow = _get_flow_or_raise(flow_id) - if not proxy_manager.is_running: + source = _get_source_proxy_for_flow(flow) + target = _get_proxy_by_id(source) + if not target.is_running: return { "success": False, - "error": "Proxy is not running. Start it with proxy_start before replaying.", + "error": f"Proxy '{source}' is not running. Start it before replaying.", } - return replay_flows(proxy_manager.call, [flow], use_modified=use_modified) + return replay_flows(target.call, [flow], use_modified=use_modified) def _flow_resume(flow_id: int) -> dict[str, Any]: flow = _get_flow_or_raise(flow_id) - if not proxy_manager.is_running: + source = _get_source_proxy_for_flow(flow) + target = _get_proxy_by_id(source) + if not target.is_running: return { "success": False, - "error": "Proxy is not running. Start it with proxy_start before resuming.", + "error": f"Proxy '{source}' is not running. Start it before resuming.", } - proxy_manager.call("flow.resume", [flow]) + target.call("flow.resume", [flow]) return {"success": True} def _flow_kill(flow_id: int) -> dict[str, Any]: flow = _get_flow_or_raise(flow_id) - if not proxy_manager.is_running: + source = _get_source_proxy_for_flow(flow) + target = _get_proxy_by_id(source) + if not target.is_running: return { "success": False, - "error": "Proxy is not running. Start it with proxy_start before killing.", + "error": f"Proxy '{source}' is not running. Start it before killing.", } - proxy_manager.call("flow.kill", [flow]) + target.call("flow.kill", [flow]) return {"success": True} @@ -514,6 +540,7 @@ def proxy_ctl( cmd: Literal[ "start", "stop", "status", "list_options", "clear_all", "wireguard_config" ], + proxy_id: Literal["main", "aux"] = "main", host: str = "127.0.0.1", port: int = 8080, capture_filter: str | None = None, @@ -527,7 +554,9 @@ def proxy_ctl( """Control the proxy. Commands: start, stop, status, list_options, clear_all, wireguard_config. Use tool_info('proxy_ctl') for details.""" try: if cmd == "start": + target = _get_proxy_by_id(proxy_id) return _proxy_start( + target_proxy=target, host=host, port=port, capture_filter=capture_filter, @@ -538,15 +567,21 @@ def proxy_ctl( web_port=web_port, ) if cmd == "stop": - return proxy_manager.stop() + target = _get_proxy_by_id(proxy_id) + return target.stop() if cmd == "status": - return proxy_manager.status() + result = proxy_manager.status() + if aux_proxy_manager.is_running: + result["aux_proxy"] = aux_proxy_manager.status() + return result if cmd == "wireguard_config": - return proxy_manager.wireguard_config() + target = _get_proxy_by_id(proxy_id) + return target.wireguard_config() if cmd == "list_options": return _proxy_list_options() if cmd == "clear_all": - return proxy_manager.clear_all(stop_proxy=stop_proxy) + target = _get_proxy_by_id(proxy_id) + return target.clear_all(stop_proxy=stop_proxy) return {"success": False, "error": f"Unknown proxy command: {cmd}"} except Exception as e: return {"success": False, "error": str(e)} @@ -610,29 +645,31 @@ def _websocket_get(flow_id: int, include_content: bool, max_content_size: int | @mcp.tool() def crypt_ctl( cmd: Literal["list", "load", "unload", "reload", "status"], + proxy_id: Literal["main", "aux"] = "main", script_path: str | None = None, script_id: str | None = None, ) -> dict[str, Any]: """Load and manage user-written encryption/decryption scripts. Commands: list, load, unload, reload, status. Use tool_info('crypt_ctl') for details.""" try: + target = _get_proxy_by_id(proxy_id) if cmd == "list": - return {"success": True, "scripts": proxy_manager.list_crypto_scripts()} + return {"success": True, "scripts": target.list_crypto_scripts()} if cmd == "load": if script_path is None: return {"success": False, "error": "script_path is required"} - return proxy_manager.load_crypto_script(script_path) + return target.load_crypto_script(script_path) if cmd == "unload": if script_id is None: return {"success": False, "error": "script_id is required"} - return proxy_manager.unload_crypto_script(script_id) + return target.unload_crypto_script(script_id) if cmd == "reload": if script_id is None: return {"success": False, "error": "script_id is required"} - return proxy_manager.reload_crypto_script(script_id) + return target.reload_crypto_script(script_id) if cmd == "status": if script_id is None: return {"success": False, "error": "script_id is required"} - status = proxy_manager.get_crypto_script_status(script_id) + status = target.get_crypto_script_status(script_id) if status is None: return {"success": False, "error": f"Crypto script '{script_id}' not found"} return {"success": True, "script": status} @@ -647,6 +684,7 @@ def websocket_ctl( "list", "get", "inject", "connect", "list_rules", "add_rule", "delete_rule", "clear_rules", ], + proxy_id: Literal["main", "aux"] = "main", flow_id: int | None = None, to_client: bool = True, message: str = "", @@ -682,12 +720,18 @@ def websocket_ctl( if cmd == "inject": if flow_id is None: return {"success": False, "error": "flow_id is required"} - return proxy_manager.inject_websocket(flow_id, to_client, message, binary) + flow = store.get(flow_id) + if flow is None: + return {"success": False, "error": f"Flow with id {flow_id} not found"} + source = _get_source_proxy_for_flow(flow) + target = _get_proxy_by_id(source) + return target.inject_websocket(flow_id, to_client, message, binary) if cmd == "connect": if url is None: return {"success": False, "error": "url is required"} + target = _get_proxy_by_id(proxy_id) header_dict = {h.name: h.value for h in headers} if headers else None - return proxy_manager.connect_websocket( + return target.connect_websocket( url=url, headers=header_dict, subprotocols=subprotocols, @@ -696,22 +740,26 @@ def websocket_ctl( timeout=timeout, ) if cmd == "list_rules": - rules = proxy_manager.list_websocket_rules() + target = _get_proxy_by_id(proxy_id) + rules = target.list_websocket_rules() return {"success": True, "rules": [r.model_dump() for r in rules]} if cmd == "add_rule": if rule is None: return {"success": False, "error": "rule is required"} from mitmproxy_mcp.websocket_rules import WebSocketRule + target = _get_proxy_by_id(proxy_id) ws_rule = WebSocketRule(**rule) - proxy_manager.add_websocket_rule(ws_rule) + target.add_websocket_rule(ws_rule) return {"success": True, "rule": ws_rule.model_dump()} if cmd == "delete_rule": if rule_id is None: return {"success": False, "error": "rule_id is required"} - deleted = proxy_manager.delete_websocket_rule(rule_id) + target = _get_proxy_by_id(proxy_id) + deleted = target.delete_websocket_rule(rule_id) return {"success": deleted, "deleted": deleted} if cmd == "clear_rules": - count = proxy_manager.clear_websocket_rules() + target = _get_proxy_by_id(proxy_id) + count = target.clear_websocket_rules() return {"success": True, "cleared": count} return {"success": False, "error": f"Unknown websocket command: {cmd}"} except Exception as e: @@ -903,28 +951,30 @@ def flow_action( @mcp.tool() def rule_ctl( cmd: Literal["list", "add", "delete", "clear"], + proxy_id: Literal["main", "aux"] = "main", rule: dict[str, Any] | None = None, rule_id: str | None = None, ) -> dict[str, Any]: """Automatic rules. Commands: list, add, delete, clear. Use tool_info('rule_ctl') for details.""" try: + target = _get_proxy_by_id(proxy_id) if cmd == "list": - rules = proxy_manager.list_rules() + rules = target.list_rules() return {"success": True, "rules": [r.model_dump(exclude_none=True) for r in rules]} if cmd == "add": if rule is None: return {"success": False, "error": "rule is required"} rule_obj = Rule(**rule) - proxy_manager.add_rule(rule_obj) + target.add_rule(rule_obj) return {"success": True, "rule": rule_obj.model_dump(exclude_none=True)} if cmd == "delete": if rule_id is None: return {"success": False, "error": "rule_id is required"} - if proxy_manager.delete_rule(rule_id): + if target.delete_rule(rule_id): return {"success": True} return {"success": False, "error": f"Rule with id {rule_id} not found"} if cmd == "clear": - count = proxy_manager.clear_rules() + count = target.clear_rules() return {"success": True, "cleared": count} return {"success": False, "error": f"Unknown rule command: {cmd}"} except Exception as e: @@ -934,31 +984,33 @@ def rule_ctl( @mcp.tool() def capture_rule_ctl( cmd: Literal["list", "add", "delete", "clear"], + proxy_id: Literal["main", "aux"] = "main", rule: dict[str, Any] | None = None, rule_id: str | None = None, ) -> dict[str, Any]: """Capture rules. Commands: list, add, delete, clear. Use tool_info('capture_rule_ctl') for details.""" try: + target = _get_proxy_by_id(proxy_id) if cmd == "list": - rules = proxy_manager.list_capture_rules() + rules = target.list_capture_rules() return {"success": True, "rules": [r.model_dump(exclude_none=True) for r in rules]} if cmd == "add": if rule is None: return {"success": False, "error": "rule is required"} rule_obj = CaptureRule(**rule) - proxy_manager.add_capture_rule(rule_obj) + target.add_capture_rule(rule_obj) return {"success": True, "rule": rule_obj.model_dump(exclude_none=True)} if cmd == "delete": if rule_id is None: return {"success": False, "error": "rule_id is required"} - if proxy_manager.delete_capture_rule(rule_id): + if target.delete_capture_rule(rule_id): return {"success": True} return { "success": False, "error": f"Capture rule with id {rule_id} not found", } if cmd == "clear": - count = proxy_manager.clear_capture_rules() + count = target.clear_capture_rules() return {"success": True, "cleared": count} return {"success": False, "error": f"Unknown capture rule command: {cmd}"} except Exception as e: diff --git a/src/mitmproxy_mcp/tool_info.py b/src/mitmproxy_mcp/tool_info.py index 946ddc6..bbb6eff 100644 --- a/src/mitmproxy_mcp/tool_info.py +++ b/src/mitmproxy_mcp/tool_info.py @@ -11,12 +11,13 @@ TOOL_INFO: dict[str, ToolInfo] = { "proxy_ctl": { - "summary": "Start, stop and inspect the mitmproxy capture proxy.", + "summary": "Start, stop and inspect the mitmproxy capture proxy. Supports a main proxy and an optional auxiliary proxy via proxy_id.", "commands": { "start": { - "description": "Start the capture proxy in a background thread. Set webui=True to also start mitmproxy's built-in web interface.", + "description": "Start the capture proxy in a background thread. Set webui=True to also start mitmproxy's built-in web interface. Use proxy_id='aux' to start the auxiliary proxy.", "required": [], "optional": [ + "proxy_id ('main' or 'aux', default 'main')", "host (str, default 127.0.0.1)", "port (int, default 8080)", "capture_filter (str, optional mitmproxy flowfilter expression)", @@ -29,13 +30,13 @@ "example": {"cmd": "start", "port": 8080, "webui": True, "web_port": 8081}, }, "stop": { - "description": "Stop the running proxy.", + "description": "Stop the running proxy. Use proxy_id='aux' to stop the auxiliary proxy.", "required": [], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": {"cmd": "stop"}, }, "status": { - "description": "Return proxy running state, listen address and captured flow count.", + "description": "Return proxy running state, listen address and captured flow count. Automatically includes auxiliary proxy info if it is running.", "required": [], "optional": [], "example": {"cmd": "status"}, @@ -43,7 +44,7 @@ "wireguard_config": { "description": "Return the WireGuard client configuration when the proxy was started in WireGuard mode. The returned INI can be imported into iOS, Android, macOS or Windows WireGuard clients. Certificate trust is still required for HTTPS/HTTP3 decryption.", "required": [], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": {"cmd": "wireguard_config"}, }, "list_options": { @@ -53,9 +54,9 @@ "example": {"cmd": "list_options"}, }, "clear_all": { - "description": "Clear all flows, rules, capture rules and mappings. Optionally stop the proxy.", + "description": "Clear all flows, rules, capture rules and mappings. Optionally stop the proxy. Use proxy_id='aux' to target the auxiliary proxy.", "required": [], - "optional": ["stop_proxy (bool, default False)"], + "optional": ["proxy_id ('main' or 'aux', default 'main')", "stop_proxy (bool, default False)"], "example": {"cmd": "clear_all", "stop_proxy": False}, }, }, @@ -185,18 +186,18 @@ }, }, "rule_ctl": { - "summary": "Manage automatic request/response modification rules.", + "summary": "Manage automatic request/response modification rules. Use proxy_id to target the auxiliary proxy.", "commands": { "list": { "description": "List all automatic rules.", "required": [], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": {"cmd": "list"}, }, "add": { "description": "Add or replace an automatic rule. Existing rule with same id is overwritten.", "required": ["rule"], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": { "cmd": "add", "rule": { @@ -210,30 +211,30 @@ "delete": { "description": "Delete a rule by id.", "required": ["rule_id"], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": {"cmd": "delete", "rule_id": "mock"}, }, "clear": { "description": "Delete all automatic rules.", "required": [], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": {"cmd": "clear"}, }, }, }, "capture_rule_ctl": { - "summary": "Manage include/exclude capture rules.", + "summary": "Manage include/exclude capture rules. Use proxy_id to target the auxiliary proxy.", "commands": { "list": { "description": "List all capture rules.", "required": [], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": {"cmd": "list"}, }, "add": { "description": "Add or replace a capture rule. Existing rule with same id is overwritten.", "required": ["rule"], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": { "cmd": "add", "rule": {"id": "api", "filter": "~u api.example.com", "action": "include"}, @@ -242,13 +243,13 @@ "delete": { "description": "Delete a capture rule by id.", "required": ["rule_id"], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": {"cmd": "delete", "rule_id": "api"}, }, "clear": { "description": "Delete all capture rules.", "required": [], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": {"cmd": "clear"}, }, }, @@ -411,42 +412,42 @@ }, }, "crypt_ctl": { - "summary": "Load and manage user-written encryption/decryption scripts for transparent traffic transformation.", + "summary": "Load and manage user-written encryption/decryption scripts for transparent traffic transformation. Use proxy_id to target the auxiliary proxy.", "commands": { "list": { "description": "List all loaded crypto handler scripts with error counts.", "required": [], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": {"cmd": "list"}, }, "load": { "description": "Load a CryptoHandler script from a Python file. The script must define a CryptoHandler subclass.", "required": ["script_path"], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": {"cmd": "load", "script_path": "/path/to/crypto_script.py"}, }, "unload": { "description": "Unload a crypto script by id.", "required": ["script_id"], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": {"cmd": "unload", "script_id": "my-handler"}, }, "reload": { "description": "Reload an already loaded crypto script by id.", "required": ["script_id"], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": {"cmd": "reload", "script_id": "my-handler"}, }, "status": { "description": "Show detailed status for a loaded crypto script, including the last error.", "required": ["script_id"], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": {"cmd": "status", "script_id": "my-handler"}, }, }, }, "websocket_ctl": { - "summary": "Manage WebSocket connections: inspect, inject, connect and modify messages with rules.", + "summary": "Manage WebSocket connections: inspect, inject, connect and modify messages with rules. Use proxy_id to target the auxiliary proxy.", "commands": { "list": { "description": "List captured WebSocket flows.", @@ -461,7 +462,7 @@ "example": {"cmd": "get", "flow_id": 1, "max_content_size": 4096}, }, "inject": { - "description": "Inject a message into an existing WebSocket connection.", + "description": "Inject a message into an existing WebSocket connection. Automatically routes to the proxy that captured the flow.", "required": ["flow_id", "message"], "optional": ["to_client (bool, default True)", "binary (bool, default False)"], "example": {"cmd": "inject", "flow_id": 1, "message": "hello from mcp", "to_client": False}, @@ -470,6 +471,7 @@ "description": "Actively open a WebSocket connection through the running proxy and capture it.", "required": ["url"], "optional": [ + "proxy_id ('main' or 'aux', default 'main')", "headers (list[Header])", "subprotocols (list[str])", "messages (list[str])", @@ -481,13 +483,13 @@ "list_rules": { "description": "List WebSocket message modification rules.", "required": [], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": {"cmd": "list_rules"}, }, "add_rule": { "description": "Add a rule that modifies or drops WebSocket messages in real time.", "required": ["rule"], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": { "cmd": "add_rule", "rule": { @@ -503,13 +505,13 @@ "delete_rule": { "description": "Delete a WebSocket rule by id.", "required": ["rule_id"], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": {"cmd": "delete_rule", "rule_id": "replace-ping"}, }, "clear_rules": { "description": "Delete all WebSocket rules.", "required": [], - "optional": [], + "optional": ["proxy_id ('main' or 'aux', default 'main')"], "example": {"cmd": "clear_rules"}, }, },