ITWorx_MCP_Hub-Public/scripts/mcp_stdio_bridge.py

"""Run the optional local stdio-to-HTTPS MCP bridge without printing credentials."""

from __future__ import annotations

import argparse
import base64
import json
import sys
from collections.abc import Mapping
from pathlib import Path
from time import monotonic
from typing import Any

import httpx

# Codex can launch this bridge while another repository is the working directory.
# Resolve imports relative to this checkout instead of relying on the caller's cwd.
_ROOT_DIR = Path(__file__).resolve().parents[1]
if str(_ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(_ROOT_DIR))

from packages.hub_clients import StdioBridge  # noqa: E402


class HttpTransport:
    def __init__(self) -> None:
        self.client = httpx.Client(timeout=30.0, follow_redirects=False)
        self._session_id: str | None = None

    def send(
        self, request: Mapping[str, Any], *, token: str, endpoint: str
    ) -> Mapping[str, Any] | None:
        headers = {
            "Authorization": f"Bearer {token}",
            "Accept": "application/json, text/event-stream",
            "Content-Type": "application/json",
        }
        if self._session_id is not None:
            headers["Mcp-Session-Id"] = self._session_id
        response = self.client.post(endpoint, json=dict(request), headers=headers)
        response.raise_for_status()
        session_id = response.headers.get("mcp-session-id")
        if session_id is not None:
            if not session_id.isascii() or not 1 <= len(session_id) <= 1024:
                raise ValueError("upstream returned an invalid MCP session identifier")
            self._session_id = session_id
        if response.status_code == 202 or not response.content:
            return None
        if "text/event-stream" in response.headers.get("content-type", ""):
            payload = _last_sse_json(response.text)
        else:
            payload = response.json()
        return payload if isinstance(payload, dict) else {"result": payload}


def _last_sse_json(body: str) -> Mapping[str, Any]:
    messages: list[Mapping[str, Any]] = []
    for line in body.splitlines():
        if not line.startswith("data:"):
            continue
        value = json.loads(line[5:].strip())
        if isinstance(value, dict):
            messages.append(value)
    if not messages:
        raise ValueError("upstream returned no MCP event data")
    return messages[-1]


class ClientCredentialsTokenProvider:
    """Fetch and cache a short-lived Authentik M2M token without emitting credentials."""

    def __init__(
        self,
        *,
        token_endpoint: str,
        client_id: str,
        username: str,
        secret_file: Path,
        scope: str = "openid profile",
    ) -> None:
        if not token_endpoint.startswith("https://"):
            raise ValueError("OAuth token endpoint must use HTTPS")
        self.token_endpoint = token_endpoint
        self.client_id = client_id
        self.username = username
        self.secret_file = secret_file
        self.scope = scope
        self._token: str | None = None
        self._refresh_at = 0.0
        self._client = httpx.Client(timeout=15.0, follow_redirects=False)

    def __call__(self) -> str:
        now = monotonic()
        if self._token is not None and now < self._refresh_at:
            return self._token
        secret = self._read_secret()
        credential = base64.b64encode(f"{self.username}:{secret}".encode()).decode("ascii")
        try:
            response = self._client.post(
                self.token_endpoint,
                data={
                    "grant_type": "client_credentials",
                    "client_id": self.client_id,
                    "client_secret": credential,
                    "scope": self.scope,
                },
                headers={"Accept": "application/json"},
            )
            response.raise_for_status()
            payload = response.json()
            token = payload.get("access_token") if isinstance(payload, dict) else None
            expires_in = payload.get("expires_in", 300) if isinstance(payload, dict) else 300
            if not isinstance(token, str) or not token or len(token) > 65536:
                raise ValueError("OAuth response has no bounded access token")
            lifetime = int(expires_in)
            if not 30 <= lifetime <= 86400:
                raise ValueError("OAuth response has invalid token lifetime")
            self._token = token
            self._refresh_at = now + max(1, lifetime - min(60, lifetime // 5))
            return token
        finally:
            del secret
            del credential

    def _read_secret(self) -> str:
        if self.secret_file.is_symlink():
            raise ValueError("OAuth client secret file is invalid")
        path = self.secret_file.resolve(strict=True)
        if not path.is_file() or path.stat().st_size > 8192:
            raise ValueError("OAuth client secret file is invalid")
        value = path.read_text(encoding="utf-8").strip()
        if not value or "\n" in value or "\r" in value:
            raise ValueError("OAuth client secret file must contain one non-empty line")
        return value


def main() -> int:
    parser = argparse.ArgumentParser(description="ITWorx MCP stdio bridge")
    parser.add_argument("--endpoint", required=True)
    parser.add_argument("--token-env", default=None)
    parser.add_argument("--client-id", required=True)
    parser.add_argument("--token-endpoint")
    parser.add_argument("--oauth-client-id")
    parser.add_argument("--oauth-username")
    parser.add_argument("--client-secret-file", type=Path)
    parser.add_argument("--scope", default="openid profile")
    args = parser.parse_args()
    oauth_values = (
        args.token_endpoint,
        args.oauth_client_id,
        args.oauth_username,
        args.client_secret_file,
    )
    if args.token_env and any(oauth_values):
        parser.error("choose either --token-env or OAuth client-credentials arguments")
    if not args.token_env and not all(oauth_values):
        parser.error("OAuth mode requires token endpoint, client id, username and secret file")
    provider = (
        ClientCredentialsTokenProvider(
            token_endpoint=args.token_endpoint,
            client_id=args.oauth_client_id,
            username=args.oauth_username,
            secret_file=args.client_secret_file,
            scope=args.scope,
        )
        if all(oauth_values)
        else None
    )
    bridge = StdioBridge(
        endpoint=args.endpoint,
        token_env=args.token_env or "ITWORX_MCP_ACCESS_TOKEN",
    )
    for output in bridge.forward(sys.stdin, HttpTransport(), token_provider=provider):
        sys.stdout.write(output + "\n")
        sys.stdout.flush()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())