hexstrike_mcp_http.py

#!/usr/bin/env python3
"""HexStrike AI — remote MCP over streamable HTTP.

Upstream ships ``hexstrike_mcp.py`` as a FastMCP server that only speaks the
**stdio** transport (``mcp.run()`` with no transport), so it can only be driven
by a client running on the same machine.

This module reuses upstream's ``setup_mcp_server()`` — every ``@mcp.tool()``
definition, unchanged — but serves it over the MCP **streamable-HTTP** transport
so the platform can be driven remotely (e.g. from Claude Code over the network).

An optional bearer token (``HEXSTRIKE_MCP_TOKEN`` / ``--token``) gates the
endpoint. Because the streamable-HTTP transport streams responses over SSE, the
guard is implemented as a lightweight *pure-ASGI* middleware that only inspects
request headers and short-circuits with 401 — it never buffers the response body
(a Starlette ``BaseHTTPMiddleware`` would break the stream).
"""

import argparse
import hmac
import logging
import os
import shutil
import sys

import uvicorn

# Reuse the upstream stdio MCP module wholesale — all tool definitions live in
# setup_mcp_server(). This file must sit next to hexstrike_mcp.py on PYTHONPATH.
import hexstrike_mcp as base

# The MCP streamable-HTTP transport applies DNS-rebinding protection: it
# validates the request Host header and, by default, only accepts localhost.
# A remote client hitting the container's LAN IP/hostname is otherwise rejected
# with "Invalid Host header". We relax this below (the endpoint is already
# guarded by the bearer token).
try:
    from mcp.server.transport_security import TransportSecuritySettings
except Exception:  # pragma: no cover - older/newer SDK layouts
    TransportSecuritySettings = None

logger = logging.getLogger("hexstrike-mcp-http")

# ---------------------------------------------------------------------------
# Tool availability filter.
#
# Upstream registers ~150 MCP tools, but many wrap external CLI binaries that
# are not installed in a given image. Exposing all of them (a) misleads the
# model into calling tools that will fail and (b) overflows tool-count limits on
# some model providers (e.g. OpenAI caps function tools at ~128). Since the MCP
# process runs inside the same container as the tools, we can check each backing
# binary with shutil.which() and hide the ones that are missing.
#
# The map only needs entries for binary-backed tools. Anything NOT listed here
# is always kept — that covers the AI/orchestration/file/process/HTTP-framework
# tools and pip-module-backed ones (angr, pwntools) that are not PATH binaries.
# Values are a binary name or a tuple of acceptable alternatives (any present).
TOOL_REQUIRES = {
    # network / recon
    "amass_scan": "amass", "nmap_scan": "nmap", "nmap_advanced_scan": "nmap",
    "masscan_high_speed": "masscan", "rustscan_fast_scan": "rustscan",
    "autorecon_scan": "autorecon", "autorecon_comprehensive": "autorecon",
    "nbtscan_netbios": "nbtscan", "arp_scan_discovery": "arp-scan",
    "responder_credential_harvest": "responder", "netexec_scan": ("netexec", "nxc"),
    "enum4linux_scan": "enum4linux", "enum4linux_ng_advanced": "enum4linux-ng",
    "rpcclient_enumeration": "rpcclient", "fierce_scan": "fierce",
    "dnsenum_scan": "dnsenum", "subfinder_scan": "subfinder", "smbmap_scan": "smbmap",
    # web
    "gobuster_scan": "gobuster", "feroxbuster_scan": "feroxbuster", "ffuf_scan": "ffuf",
    "dirb_scan": "dirb", "dirsearch_scan": "dirsearch", "nikto_scan": "nikto",
    "sqlmap_scan": "sqlmap", "wpscan_analyze": "wpscan", "nuclei_scan": "nuclei",
    "httpx_probe": "httpx", "katana_crawl": "katana", "hakrawler_crawl": "hakrawler",
    "gau_discovery": "gau", "waybackurls_discovery": "waybackurls",
    "arjun_parameter_discovery": "arjun", "arjun_scan": "arjun",
    "paramspider_discovery": "paramspider", "paramspider_mining": "paramspider",
    "dalfox_xss_scan": "dalfox", "wafw00f_scan": "wafw00f",
    "x8_parameter_discovery": "x8", "jaeles_vulnerability_scan": "jaeles",
    "dotdotpwn_scan": "dotdotpwn", "xsser_scan": "xsser", "wfuzz_scan": "wfuzz",
    "graphql_scanner": "graphql-scanner", "jwt_analyzer": "jwt-analyzer",
    "burpsuite_scan": "burpsuite", "zap_scan": "zaproxy",
    "qsreplace_parameter_replacement": "qsreplace", "uro_url_filtering": "uro",
    "anew_data_processing": "anew",
    # password / crack
    "hydra_attack": "hydra", "john_crack": "john", "hashcat_crack": "hashcat",
    "hashpump_attack": "hashpump",
    # binary / RE
    "gdb_analyze": "gdb", "gdb_peda_debug": "gdb", "radare2_analyze": "radare2",
    "binwalk_analyze": "binwalk", "ropgadget_search": "ROPgadget",
    "ropper_gadget_search": "ropper", "one_gadget_search": "one_gadget",
    "checksec_analyze": "checksec", "objdump_analyze": "objdump",
    "strings_extract": "strings", "xxd_hexdump": "xxd", "ghidra_analysis": "ghidra",
    "pwninit_setup": "pwninit", "libc_database_lookup": "libc-database",
    "msfvenom_generate": "msfvenom", "metasploit_run": "msfconsole",
    "volatility_analyze": ("volatility", "vol"),
    "volatility3_analyze": ("volatility3", "vol"),
    # forensics
    "steghide_analysis": "steghide", "foremost_carving": "foremost",
    "exiftool_extract": "exiftool",
    # cloud / container
    "prowler_scan": "prowler", "scout_suite_assessment": ("scout-suite", "scout"),
    "trivy_scan": "trivy", "kube_hunter_scan": "kube-hunter",
    "kube_bench_cis": "kube-bench", "docker_bench_security_scan": "docker-bench-security",
    "checkov_iac_scan": "checkov", "terrascan_iac_scan": "terrascan",
    "falco_runtime_monitoring": "falco", "clair_vulnerability_scan": "clair",
    "cloudmapper_analysis": "cloudmapper", "pacu_exploitation": "pacu",
}


def _binary_present(req):
    """True if the required binary (str) or any alternative (tuple) is on PATH."""
    if isinstance(req, (tuple, list)):
        return any(shutil.which(b) for b in req)
    return shutil.which(req) is not None


def _csv_env(value):
    return {x.strip() for x in (value or "").split(",") if x.strip()}


def _remove_tool(mcp, name):
    """Remove a registered tool across FastMCP versions."""
    remover = getattr(mcp, "remove_tool", None)
    if callable(remover):
        try:
            remover(name)
            return True
        except Exception:
            pass
    tm = getattr(mcp, "_tool_manager", None)
    tools = getattr(tm, "_tools", None)
    if isinstance(tools, dict):
        return tools.pop(name, None) is not None
    return False


def apply_tool_filter(mcp, args):
    """Prune MCP tools whose backing binary is not installed.

    Controlled by (arg / env):
      --tool-filter / HEXSTRIKE_MCP_TOOL_FILTER   auto (default) | off
      --include-only / HEXSTRIKE_MCP_INCLUDE_ONLY expose ONLY these (whitelist)
      --exclude-tools / HEXSTRIKE_MCP_EXCLUDE_TOOLS  always drop these
      --keep-tools / HEXSTRIKE_MCP_KEEP_TOOLS     protect these from auto-drop
    """
    tm = getattr(mcp, "_tool_manager", None)
    tools = getattr(tm, "_tools", None)
    if not isinstance(tools, dict):
        logger.warning("Cannot access tool registry — skipping tool filter")
        return
    all_names = set(tools)
    include_only = _csv_env(getattr(args, "include_only", ""))
    exclude = _csv_env(getattr(args, "exclude_tools", ""))
    keep = _csv_env(getattr(args, "keep_tools", ""))
    mode = (getattr(args, "tool_filter", "auto") or "auto").lower()

    to_drop = set()
    if include_only:
        to_drop |= {n for n in all_names if n not in include_only}
    elif mode != "off":
        to_drop |= {n for n, req in TOOL_REQUIRES.items()
                    if n in all_names and not _binary_present(req)}
    to_drop -= keep          # explicit keep wins over auto-drop
    to_drop |= (exclude & all_names)   # explicit exclude always drops

    for name in sorted(to_drop):
        _remove_tool(mcp, name)

    remaining = len(tools)
    logger.info("Tool filter (%s): %d dropped, %d exposed%s",
                "off" if (mode == "off" and not include_only and not exclude) else mode,
                len(to_drop), remaining,
                " [whitelist]" if include_only else "")
    if to_drop:
        logger.info("Hidden tools: %s", ", ".join(sorted(to_drop)))
    return remaining


DEFAULT_MCP_HOST = "0.0.0.0"
DEFAULT_MCP_PORT = 8889
DEFAULT_MCP_PATH = "/mcp"
DEFAULT_TIMEOUT = getattr(base, "DEFAULT_REQUEST_TIMEOUT", 300)


class BearerTokenASGIMiddleware:
    """Reject requests to the MCP path lacking a valid ``Authorization: Bearer``.

    No-op when ``token`` is falsy, preserving upstream's open trust model for
    isolated lab/CTF use. Kept as raw ASGI so streaming (SSE) responses pass
    through untouched.
    """

    def __init__(self, app, token, protected_path):
        self.app = app
        self.token = token
        self.protected_path = protected_path
        self._expected = f"Bearer {token}" if token else None

    async def __call__(self, scope, receive, send):
        if self._expected and scope.get("type") == "http":
            if scope.get("path", "").startswith(self.protected_path):
                headers = dict(scope.get("headers") or [])
                auth = headers.get(b"authorization", b"").decode("latin-1")
                if not hmac.compare_digest(auth, self._expected):
                    await self._reject(send)
                    return
        await self.app(scope, receive, send)

    @staticmethod
    async def _reject(send):
        body = b'{"error":"unauthorized"}'
        await send({
            "type": "http.response.start",
            "status": 401,
            "headers": [
                (b"content-type", b"application/json"),
                (b"content-length", str(len(body)).encode()),
                (b"www-authenticate", b'Bearer realm="hexstrike-mcp"'),
            ],
        })
        await send({"type": "http.response.body", "body": body})


def parse_args(argv=None):
    parser = argparse.ArgumentParser(
        description="Run the HexStrike AI MCP server over streamable HTTP")
    parser.add_argument(
        "--server", type=str,
        default=os.environ.get("HEXSTRIKE_SERVER_URL", base.DEFAULT_HEXSTRIKE_SERVER),
        help="HexStrike API server URL (default: %(default)s)")
    parser.add_argument(
        "--host", type=str,
        default=os.environ.get("HEXSTRIKE_MCP_HOST", DEFAULT_MCP_HOST),
        help="Interface to bind the MCP HTTP server (default: %(default)s)")
    parser.add_argument(
        "--port", type=int,
        default=int(os.environ.get("HEXSTRIKE_MCP_PORT", DEFAULT_MCP_PORT)),
        help="Port to bind the MCP HTTP server (default: %(default)s)")
    parser.add_argument(
        "--path", type=str,
        default=os.environ.get("HEXSTRIKE_MCP_PATH", DEFAULT_MCP_PATH),
        help="URL path for the streamable-HTTP endpoint (default: %(default)s)")
    parser.add_argument(
        "--token", type=str,
        default=os.environ.get("HEXSTRIKE_MCP_TOKEN", ""),
        help="Optional bearer token; when set, clients must send "
             "'Authorization: Bearer <token>'")
    parser.add_argument(
        "--timeout", type=int,
        default=int(os.environ.get("HEXSTRIKE_TIMEOUT", DEFAULT_TIMEOUT)),
        help="HexStrike API request timeout in seconds (default: %(default)s)")
    parser.add_argument(
        "--allowed-hosts", type=str,
        default=os.environ.get("HEXSTRIKE_MCP_ALLOWED_HOSTS", ""),
        help="Comma-separated Host/Origin allowlist for DNS-rebinding protection "
             "(e.g. 'host.example:8889'). When empty (default), the Host check is "
             "disabled so the endpoint is reachable via any address/hostname.")
    parser.add_argument(
        "--tool-filter", type=str,
        default=os.environ.get("HEXSTRIKE_MCP_TOOL_FILTER", "auto"),
        help="'auto' (default) hides tools whose backing binary is not installed; "
             "'off' exposes every registered tool.")
    parser.add_argument(
        "--include-only", type=str,
        default=os.environ.get("HEXSTRIKE_MCP_INCLUDE_ONLY", ""),
        help="Comma-separated whitelist: expose ONLY these tools (overrides auto).")
    parser.add_argument(
        "--exclude-tools", type=str,
        default=os.environ.get("HEXSTRIKE_MCP_EXCLUDE_TOOLS", ""),
        help="Comma-separated tools to always hide.")
    parser.add_argument(
        "--keep-tools", type=str,
        default=os.environ.get("HEXSTRIKE_MCP_KEEP_TOOLS", ""),
        help="Comma-separated tools to always keep (protects them from auto-drop).")
    parser.add_argument("--debug", action="store_true", help="Enable debug logging")
    return parser.parse_args(argv)


def build_app(args):
    """Construct the streamable-HTTP ASGI app wrapping upstream's MCP tools."""
    client = base.HexStrikeClient(args.server, args.timeout)

    # Non-fatal health probe — the API server may still be starting up.
    try:
        health = client.check_health()
    except Exception as exc:  # pragma: no cover - defensive
        health = {"error": str(exc)}
    if "error" in health:
        logger.warning("HexStrike API not reachable yet at %s: %s",
                       args.server, health.get("error"))
    else:
        logger.info("Connected to HexStrike API at %s (status=%s, version=%s)",
                    args.server, health.get("status"), health.get("version"))

    mcp = base.setup_mcp_server(client)

    # Hide tools whose backing binary is not installed (keeps the exposed set
    # under provider tool-count limits and avoids surfacing tools that can only
    # fail). Configurable via --tool-filter / --include-only / etc.
    apply_tool_filter(mcp, args)

    # Steer FastMCP's settings so streamable_http_app() mounts where we expect
    # and (if run via mcp.run) binds the right interface. Guarded because
    # attribute names vary slightly across mcp SDK versions.
    for attr, value in (("host", args.host), ("port", args.port),
                        ("streamable_http_path", args.path)):
        try:
            setattr(mcp.settings, attr, value)
        except Exception:
            logger.debug("Could not set mcp.settings.%s", attr)

    # Configure DNS-rebinding / Host-header validation on the transport so the
    # endpoint is reachable remotely (by default upstream only allows localhost).
    if TransportSecuritySettings is not None:
        allowed = [h.strip() for h in (args.allowed_hosts or "").split(",") if h.strip()]
        try:
            if allowed:
                mcp.settings.transport_security = TransportSecuritySettings(
                    enable_dns_rebinding_protection=True,
                    allowed_hosts=allowed,
                    allowed_origins=allowed,
                )
                logger.info("Host-header allowlist enabled: %s", allowed)
            else:
                mcp.settings.transport_security = TransportSecuritySettings(
                    enable_dns_rebinding_protection=False,
                )
                logger.info("DNS-rebinding Host check disabled (any host accepted)")
        except Exception as exc:
            logger.warning("Could not configure transport_security: %s", exc)
    else:
        logger.warning("TransportSecuritySettings unavailable — remote Host headers "
                       "may be rejected with 'Invalid Host header'")

    app = mcp.streamable_http_app()

    token = (args.token or "").strip() or None
    if token:
        logger.info("Bearer-token auth ENABLED on %s", args.path)
    else:
        logger.warning("Bearer-token auth DISABLED — expose only on a trusted "
                       "network / VPN. Set HEXSTRIKE_MCP_TOKEN to require auth.")
    return BearerTokenASGIMiddleware(app, token, args.path)


def main(argv=None):
    args = parse_args(argv)
    logging.basicConfig(
        level=logging.DEBUG if args.debug else logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
    logger.info("Starting HexStrike AI MCP (streamable-HTTP) on %s:%s%s -> API %s",
                args.host, args.port, args.path, args.server)
    try:
        app = build_app(args)
    except Exception as exc:
        logger.error("Failed to build MCP app: %s", exc)
        import traceback
        logger.error(traceback.format_exc())
        sys.exit(1)
    uvicorn.run(app, host=args.host, port=args.port,
                log_level="debug" if args.debug else "info")


if __name__ == "__main__":
    main()