PortHunter MCP

import os import sys import time import json import asyncio from pathlib import Path from typing import Any, Dict from mcp import StdioServerParameters, types from mcp.client.stdio import stdio_client from mcp.client.session import ClientSession ROOT_DIR = Path(__file__).resolve().parents[1] DEFAULT_PCAP_DIR = ROOT_DIR / "captures" TOKEN = os.getenv("PORT_HUNTER_TOKEN", "MiTOKENultraSecreto123") PCAP_DIR = Path(os.getenv("PORT_HUNTER_ALLOWED_DIR", str(DEFAULT_PCAP_DIR))).resolve() SERVER_ENV = { "PORT_HUNTER_TOKEN": TOKEN, "PORT_HUNTER_ALLOWED_DIR": str(PCAP_DIR), "PORT_HUNTER_ALLOW_PRIVATE": os.getenv("PORT_HUNTER_ALLOW_PRIVATE", "false"), "PORT_HUNTER_CACHE_DIR": os.getenv("PORT_HUNTER_CACHE_DIR", ".cache/porthunter"), "PORT_HUNTER_REQUIRE_TOKEN": os.getenv("PORT_HUNTER_REQUIRE_TOKEN", "true"), "PORT_HUNTER_MAX_PCAP_MB": os.getenv("PORT_HUNTER_MAX_PCAP_MB", "200"), "OTX_API_KEY": os.getenv("OTX_API_KEY", ""), "GREYNOISE_API_KEY": os.getenv("GREYNOISE_API_KEY", ""), "GEOLITE2_CITY_DB": os.getenv("GEOLITE2_CITY_DB") or os.getenv("GEOIP_DB_PATH", ""), } PORT_HUNTER = StdioServerParameters( command=sys.executable, args=["-m", "porthunter.server"], # <- lo mismo que usa tu CLI env=SERVER_ENV, ) async def _call_tool(name: str, arguments: Dict[str, Any], timeout_s: float = 30.0) -> Any: async with stdio_client(PORT_HUNTER) as (read, write): async with ClientSession(read, write) as session: await session.initialize() rsp = await asyncio.wait_for( session.call_tool(name=name, arguments=arguments), timeout=timeout_s, ) sc = getattr(rsp, "structuredContent", None) if isinstance(sc, dict): return sc if rsp.content: txt = "" for block in rsp.content: if isinstance(block, types.TextContent): txt += block.text if txt: try: return json.loads(txt) except Exception: return txt return None def _abs_pcap(path: str) -> str: p = Path(path) if not p.is_absolute(): p = (PCAP_DIR / p).resolve() return str(p) async def run_bench(pcap_name: str) -> Dict[str, Any]: results: Dict[str, Any] = {"pcap": pcap_name} t0 = time.perf_counter() info = await _call_tool("get_info", {"auth_token": TOKEN}) t1 = time.perf_counter() p = _abs_pcap(pcap_name) ov0 = time.perf_counter() overview = await _call_tool("scan_overview", {"path": p, "auth_token": TOKEN}) ov1 = time.perf_counter() fe0 = time.perf_counter() first = await _call_tool("first_scan_event", {"path": p, "auth_token": TOKEN}) fe1 = time.perf_counter() results["get_info_s"] = round(t1 - t0, 3) results["overview_s"] = round(ov1 - ov0, 3) results["first_event_s"] = round(fe1 - fe0, 3) try: ov = overview.get("overview", {}) results["total_pkts"] = ov.get("total_pkts") results["interval_s"] = ov.get("interval_s") results["suspected_patterns"] = ov.get("suspected_patterns", []) results["scanners_count"] = len(ov.get("scanners", [])) except Exception: pass results["ok"] = bool(overview and first) return results def main(): import argparse ap = argparse.ArgumentParser(description="Benchmark PortHunter (MCP stdio)") ap.add_argument("pcap", help="Archivo .pcap o .pcapng dentro de PORT_HUNTER_ALLOWED_DIR") args = ap.parse_args() loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) res = loop.run_until_complete(run_bench(args.pcap)) print(json.dumps(res, indent=2, ensure_ascii=False)) if __name__ == "__main__": main()