lightnow-proxy 1.0.0__py3-none-any.whl

lightnow_proxy/__init__.py

@@ -0,0 +1 @@
+__version__ = "1.0.0"

lightnow_proxy/app.py

@@ -0,0 +1,278 @@
+from __future__ import annotations
+
+from contextlib import AsyncExitStack, asynccontextmanager
+import logging
+from typing import Any, AsyncIterator
+
+from mcp.server import Server
+from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
+from mcp import types
+from mcp.types import CallToolResult, ReadResourceResult, TextContent
+from starlette.applications import Starlette
+from starlette.datastructures import Headers
+from starlette.requests import Request
+from starlette.responses import JSONResponse, Response
+from starlette.routing import Route
+from starlette.types import Receive, Scope, Send
+import structlog
+
+from lightnow_proxy import __version__
+from lightnow_proxy.auth import AuthError, TokenVerifier, has_required_group
+from lightnow_proxy.config import ProxyConfig
+from lightnow_proxy.request_context import current_principal
+from lightnow_proxy.router import ToolRouter, ToolRoutingError
+
+logger = structlog.get_logger("lightnow_proxy")
+
+
+class ProfileMCPApp:
+    def __init__(
+        self,
+        profile_name: str,
+        config: ProxyConfig,
+        router: ToolRouter,
+        verifier: TokenVerifier,
+    ):
+        self.profile_name = profile_name
+        self.config = config
+        self.router = router
+        self.verifier = verifier
+        self.server = self._build_server()
+        self.session_manager = StreamableHTTPSessionManager(
+            app=self.server,
+            json_response=False,
+            stateless=True,
+        )
+
+    def _build_server(self) -> Server:
+        server = Server(f"lightnow-proxy-{self.profile_name}", version=__version__)
+
+        @server.list_tools()
+        async def list_tools() -> list:
+            return await self.router.list_tools(self.profile_name)
+
+        @server.call_tool(validate_input=False)
+        async def call_tool(
+            name: str,
+            arguments: dict[str, Any] | None = None,
+            request_meta: dict[str, Any] | None = None,
+        ) -> CallToolResult:
+            try:
+                return await self.router.call_tool(self.profile_name, name, arguments or {}, request_meta)
+            except ToolRoutingError as exc:
+                return CallToolResult(content=[TextContent(type="text", text=str(exc))], isError=True)
+
+        @server.list_resources()
+        async def list_resources() -> list:
+            return await self.router.list_resources(self.profile_name)
+
+        @server.list_resource_templates()
+        async def list_resource_templates() -> list:
+            return await self.router.list_resource_templates(self.profile_name)
+
+        @server.read_resource()
+        async def read_resource(uri: str) -> ReadResourceResult:
+            return await self.router.read_resource(self.profile_name, uri)
+
+        register_unvalidated_call_tool_handler(server, call_tool)
+        return server
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        profile = self.config.profiles[self.profile_name]
+        try:
+            principal = await self.verifier.verify_headers(Headers(scope=scope))
+            if not has_required_group(principal, profile.required_groups):
+                logger.warning(
+                    "mcp_profile_forbidden",
+                    profile=self.profile_name,
+                    subject=principal.subject,
+                    username=principal.username,
+                    required_groups=profile.required_groups,
+                )
+                await forbidden(scope, receive, send)
+                return
+            scope.setdefault("state", {})
+            scope["state"]["principal"] = principal
+        except AuthError as exc:
+            logger.warning("mcp_profile_unauthorized", profile=self.profile_name, reason=str(exc))
+            await unauthorized(scope, receive, send, str(exc))
+            return
+        except Exception as exc:
+            logger.warning("mcp_profile_auth_failed", profile=self.profile_name, reason=exc.__class__.__name__)
+            await unauthorized(scope, receive, send, "invalid token")
+            return
+
+        principal_token = current_principal.set(principal)
+        try:
+            await self.session_manager.handle_request(scope, receive, send)
+        finally:
+            current_principal.reset(principal_token)
+
+
+class LocalProxyMCPApp:
+    def __init__(
+        self,
+        config: ProxyConfig,
+        router: ToolRouter,
+    ):
+        self.config = config
+        self.router = router
+        self.profile_name = config.local_proxy.profile
+        self.server = self._build_server()
+        self.session_manager = StreamableHTTPSessionManager(
+            app=self.server,
+            json_response=False,
+            stateless=True,
+        )
+
+    def _build_server(self) -> Server:
+        server = Server("lightnow-local-proxy", version=__version__)
+
+        @server.list_tools()
+        async def list_tools() -> list:
+            return await self.router.list_tools(self.profile_name)
+
+        @server.call_tool(validate_input=False)
+        async def call_tool(
+            name: str,
+            arguments: dict[str, Any] | None = None,
+            request_meta: dict[str, Any] | None = None,
+        ) -> CallToolResult:
+            try:
+                return await self.router.call_tool(self.profile_name, name, arguments or {}, request_meta)
+            except ToolRoutingError as exc:
+                return CallToolResult(content=[TextContent(type="text", text=str(exc))], isError=True)
+
+        @server.list_resources()
+        async def list_resources() -> list:
+            return await self.router.list_resources(self.profile_name)
+
+        @server.list_resource_templates()
+        async def list_resource_templates() -> list:
+            return await self.router.list_resource_templates(self.profile_name)
+
+        @server.read_resource()
+        async def read_resource(uri: str) -> ReadResourceResult:
+            return await self.router.read_resource(self.profile_name, uri)
+
+        register_unvalidated_call_tool_handler(server, call_tool)
+        return server
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        await self.session_manager.handle_request(scope, receive, send)
+
+
+async def unauthorized(scope: Scope, receive: Receive, send: Send, detail: str) -> None:
+    response = JSONResponse(
+        {"error": "unauthorized", "detail": detail},
+        status_code=401,
+        headers={"WWW-Authenticate": 'Bearer error="invalid_token"'},
+    )
+    await response(scope, receive, send)
+
+
+async def forbidden(scope: Scope, receive: Receive, send: Send) -> None:
+    response = JSONResponse({"error": "forbidden"}, status_code=403)
+    await response(scope, receive, send)
+
+
+async def healthz(_request: Request) -> Response:
+    return JSONResponse({"status": "ok"})
+
+
+async def status(_request: Request, config: ProxyConfig) -> Response:
+    """Return non-secret Local Proxy runtime status for local health checks."""
+    mode = "local_proxy" if config.local_proxy.enabled else "profile_proxy"
+    return JSONResponse(
+        {
+            "status": "ok",
+            "mode": mode,
+            "runner": {
+                "name": config.local_proxy.runner_name,
+                "version": config.local_proxy.runner_version,
+            },
+            "local_proxy": {
+                "enabled": config.local_proxy.enabled,
+                "profile": config.local_proxy.profile if config.local_proxy.enabled else None,
+                "path": config.local_proxy.path if config.local_proxy.enabled else None,
+                "sync_from_lightnow": config.local_proxy.sync_from_lightnow,
+                "client_name": config.local_proxy.client_name,
+                "client_version": config.local_proxy.client_version,
+                "client_transport": config.local_proxy.client_transport,
+                "telemetry_enabled": config.local_proxy.telemetry_enabled,
+                "policy_mode": config.local_proxy.policy_mode,
+                "allow_unmanaged_client_servers": config.local_proxy.allow_unmanaged_client_servers,
+            },
+            "registry_api": {
+                "enabled": bool(config.registry_api and config.registry_api.enabled),
+                "use_cli_session": bool(config.registry_api and config.registry_api.use_cli_session),
+            },
+            "profiles": sorted(config.profiles.keys()),
+            "static_upstreams": sorted(config.upstreams.keys()),
+        }
+    )
+
+
+def create_app(config: ProxyConfig, router: ToolRouter | None = None, verifier: TokenVerifier | None = None) -> Starlette:
+    logging.basicConfig(level=logging.INFO)
+    structlog.configure(
+        wrapper_class=structlog.make_filtering_bound_logger(logging.INFO),
+        processors=[
+            structlog.processors.add_log_level,
+            structlog.processors.TimeStamper(fmt="iso", utc=True),
+            structlog.processors.JSONRenderer(),
+        ],
+    )
+
+    router = router or ToolRouter(config)
+    verifier = verifier or TokenVerifier(config.auth)
+    local_proxy_app = LocalProxyMCPApp(config, router) if config.local_proxy.enabled else None
+    profile_apps = {}
+    if local_proxy_app is None:
+        profile_apps = {
+            profile_name: ProfileMCPApp(profile_name, config, router, verifier)
+            for profile_name in sorted(config.profiles.keys())
+        }
+
+    async def status_endpoint(request: Request) -> Response:
+        return await status(request, config)
+
+    routes = [
+        Route("/healthz", healthz, methods=["GET"]),
+        Route("/status", status_endpoint, methods=["GET"]),
+    ]
+    if local_proxy_app is not None:
+        routes.append(Route(config.local_proxy.path, endpoint=local_proxy_app))
+    for profile_name, profile_app in profile_apps.items():
+        routes.append(Route(f"/profiles/{profile_name}/mcp", endpoint=profile_app))
+
+    @asynccontextmanager
+    async def lifespan(_app: Starlette) -> AsyncIterator[None]:
+        async with AsyncExitStack() as stack:
+            if local_proxy_app is not None:
+                await stack.enter_async_context(local_proxy_app.session_manager.run())
+                await emit_lifecycle_event(router, "runner_started")
+            for profile_app in profile_apps.values():
+                await stack.enter_async_context(profile_app.session_manager.run())
+            try:
+                yield
+            finally:
+                if local_proxy_app is not None:
+                    await emit_lifecycle_event(router, "runner_stopped")
+
+    return Starlette(routes=routes, lifespan=lifespan)
+
+
+async def emit_lifecycle_event(router: Any, event_type: str) -> None:
+    emitter = getattr(router, "emit_lifecycle_event", None)
+    if emitter is not None:
+        await emitter(event_type)
+
+
+def register_unvalidated_call_tool_handler(server: Server, handler_func) -> None:
+    async def handler(req: types.CallToolRequest):
+        request_meta = req.params.meta.model_dump(mode="json", by_alias=True, exclude_none=True) if req.params.meta else None
+        result = await handler_func(req.params.name, req.params.arguments or {}, request_meta)
+        return types.ServerResult(result)
+
+    server.request_handlers[types.CallToolRequest] = handler

lightnow_proxy/auth.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+import time
+from typing import Any
+
+import httpx
+import jwt
+from jwt import PyJWTError
+from jwt import PyJWKClient
+from starlette.datastructures import Headers
+
+from lightnow_proxy.config import AuthConfig
+
+
+class AuthError(Exception):
+    pass
+
+
+class Principal:
+    def __init__(self, subject: str, username: str | None, groups: set[str], claims: dict[str, Any]):
+        self.subject = subject
+        self.username = username
+        self.groups = groups
+        self.claims = claims
+
+
+class TokenVerifier:
+    def __init__(self, config: AuthConfig):
+        self.config = config
+        self._jwks_client: PyJWKClient | None = None
+        self._jwks_client_expires_at = 0.0
+
+    async def verify_headers(self, headers: Headers) -> Principal:
+        if not self.config.enabled:
+            return Principal(subject="anonymous", username="anonymous", groups={"*"}, claims={})
+
+        auth_header = headers.get("authorization")
+        if not auth_header or not auth_header.lower().startswith("bearer "):
+            raise AuthError("missing bearer token")
+
+        token = auth_header.split(" ", 1)[1].strip()
+        if token in self.config.dev_bearer_tokens:
+            return self._principal_from_claims(self.config.dev_bearer_tokens[token])
+
+        claims = await self._decode_jwt(token)
+        return self._principal_from_claims(claims)
+
+    async def _decode_jwt(self, token: str) -> dict[str, Any]:
+        client = await self._jwks()
+        signing_key = client.get_signing_key_from_jwt(token)
+        try:
+            return self._decode_with_audience_or_authorized_party(token, signing_key.key)
+        except PyJWTError as exc:
+            raise AuthError(f"{exc.__class__.__name__}: {exc}") from exc
+
+    def _decode_with_audience_or_authorized_party(self, token: str, signing_key: Any) -> dict[str, Any]:
+        audiences = configured_audiences(self.config.audience)
+        if not audiences:
+            return jwt.decode(
+                token,
+                signing_key,
+                algorithms=["RS256", "RS384", "RS512"],
+                issuer=self.config.issuer,
+                options={"verify_aud": False},
+            )
+
+        unverified_claims = jwt.decode(token, options={"verify_signature": False})
+        if unverified_claims.get("aud"):
+            try:
+                return jwt.decode(
+                    token,
+                    signing_key,
+                    algorithms=["RS256", "RS384", "RS512"],
+                    audience=audiences,
+                    issuer=self.config.issuer,
+                )
+            except jwt.InvalidAudienceError:
+                if unverified_claims.get("azp") not in audiences:
+                    raise
+        elif unverified_claims.get("azp") not in audiences:
+            raise jwt.MissingRequiredClaimError("aud")
+
+        claims = jwt.decode(
+            token,
+            signing_key,
+            algorithms=["RS256", "RS384", "RS512"],
+            issuer=self.config.issuer,
+            options={"verify_aud": False},
+        )
+        if claims.get("azp") not in audiences:
+            raise jwt.InvalidAudienceError("authorized party is not allowed")
+        return claims
+
+    async def _jwks(self) -> PyJWKClient:
+        now = time.time()
+        if self._jwks_client and now < self._jwks_client_expires_at:
+            return self._jwks_client
+
+        async with httpx.AsyncClient(timeout=10) as client:
+            response = await client.get(f"{self.config.issuer.rstrip('/')}/.well-known/openid-configuration")
+            response.raise_for_status()
+            metadata = response.json()
+
+        jwks_uri = metadata.get("jwks_uri")
+        if not isinstance(jwks_uri, str) or not jwks_uri:
+            raise AuthError("issuer metadata does not contain jwks_uri")
+
+        self._jwks_client = PyJWKClient(jwks_uri)
+        self._jwks_client_expires_at = now + self.config.jwks_cache_seconds
+        return self._jwks_client
+
+    def _principal_from_claims(self, claims: dict[str, Any]) -> Principal:
+        subject = str(claims.get("sub") or "")
+        if not subject:
+            raise AuthError("token does not contain sub")
+
+        raw_groups = claims.get(self.config.groups_claim, [])
+        groups = normalize_groups(raw_groups)
+        username = claims.get("preferred_username") or claims.get("email")
+        return Principal(subject=subject, username=str(username) if username else None, groups=groups, claims=claims)
+
+
+def normalize_groups(raw_groups: Any) -> set[str]:
+    if raw_groups is None:
+        return set()
+    if isinstance(raw_groups, str):
+        raw_iterable = [raw_groups]
+    elif isinstance(raw_groups, list):
+        raw_iterable = raw_groups
+    else:
+        return set()
+
+    groups: set[str] = set()
+    for item in raw_iterable:
+        if not isinstance(item, str):
+            continue
+        normalized = item.strip()
+        if not normalized:
+            continue
+        groups.add(normalized)
+        groups.add(normalized.strip("/"))
+        groups.add(normalized.rsplit("/", 1)[-1])
+    return groups
+
+
+def configured_audiences(raw_audience: str | list[str] | None) -> list[str]:
+    if raw_audience is None:
+        return []
+    if isinstance(raw_audience, str):
+        return [raw_audience]
+    return [audience for audience in raw_audience if audience]
+
+
+def has_required_group(principal: Principal, required_groups: list[str]) -> bool:
+    if "*" in principal.groups:
+        return True
+    if not required_groups:
+        return True
+    wanted = {group.strip("/").rsplit("/", 1)[-1] for group in required_groups}
+    return bool(principal.groups.intersection(wanted))

lightnow_proxy/capture.py

@@ -0,0 +1,200 @@
+from __future__ import annotations
+
+from datetime import UTC, datetime
+import json
+from pathlib import Path
+from typing import Any, Callable
+
+import anyio
+from mcp.shared.message import SessionMessage
+from mcp import types
+
+SENSITIVE_KEY_FRAGMENTS = (
+    "authorization",
+    "api_key",
+    "apikey",
+    "access_token",
+    "refresh_token",
+    "token",
+    "password",
+    "secret",
+    "client_secret",
+)
+
+
+def capture_enabled(path: str | None) -> bool:
+    return bool(path and path.strip())
+
+
+def initialize_capture_file(path: str, config_path: str) -> None:
+    capture_path = Path(path).expanduser()
+    capture_path.parent.mkdir(parents=True, exist_ok=True)
+    with capture_path.open("a", encoding="utf-8") as handle:
+        handle.write("LightNow Local Proxy MCP Protocol Capture\n")
+        handle.write("=========================================\n")
+        handle.write(f"started_at: {timestamp()}\n")
+        handle.write(f"config: {config_path}\n")
+        handle.write("redaction: secret-like keys are replaced with [REDACTED]\n\n")
+
+
+SessionMessageObserver = Callable[[SessionMessage], None]
+
+
+async def capture_read_stream(
+    read_stream,
+    capture_path: str | None = None,
+    on_session_message: SessionMessageObserver | None = None,
+):
+    sender, receiver = anyio.create_memory_object_stream(0)
+
+    async def forward_messages() -> None:
+        async with sender:
+            async for item in read_stream:
+                if isinstance(item, SessionMessage):
+                    if on_session_message is not None:
+                        on_session_message(item)
+                    if capture_path:
+                        await write_capture(capture_path, describe_session_message(item))
+                await sender.send(item)
+
+    return receiver, forward_messages
+
+
+async def write_capture(capture_path: str, block: str | None) -> None:
+    if not block:
+        return
+
+    def append() -> None:
+        with Path(capture_path).expanduser().open("a", encoding="utf-8") as handle:
+            handle.write(block)
+            handle.write("\n")
+
+    await anyio.to_thread.run_sync(append)
+
+
+def describe_session_message(session_message: SessionMessage) -> str | None:
+    message = session_message.message.root
+    if isinstance(message, types.JSONRPCRequest):
+        return describe_request(message)
+    if isinstance(message, types.JSONRPCNotification):
+        return describe_notification(message)
+    return None
+
+
+def initialize_client_context(session_message: SessionMessage) -> dict[str, Any] | None:
+    message = session_message.message.root
+    if not isinstance(message, types.JSONRPCRequest):
+        return None
+
+    payload = message.model_dump(mode="json", by_alias=True, exclude_none=True)
+    if payload.get("method") != "initialize":
+        return None
+
+    params = payload.get("params") if isinstance(payload.get("params"), dict) else {}
+    client_info = params.get("clientInfo") if isinstance(params.get("clientInfo"), dict) else {}
+    capabilities = params.get("capabilities") if isinstance(params.get("capabilities"), dict) else {}
+    return {
+        "name": client_info.get("name") if isinstance(client_info.get("name"), str) else None,
+        "version": client_info.get("version") if isinstance(client_info.get("version"), str) else None,
+        "capability_keys": sorted(str(key) for key in capabilities.keys()),
+    }
+
+
+def describe_request(message: types.JSONRPCRequest) -> str:
+    payload = message.model_dump(mode="json", by_alias=True, exclude_none=True)
+    method = str(payload.get("method") or "")
+    params = payload.get("params") if isinstance(payload.get("params"), dict) else {}
+    redacted_params = redact(params)
+    lines = [
+        f"[{timestamp()}] request {method}",
+        f"id: {payload.get('id')}",
+        f"params_bytes: {json_size(params)}",
+    ]
+
+    if method == "initialize":
+        lines.extend(describe_initialize_params(redacted_params))
+    elif method == "tools/call":
+        lines.extend(describe_tool_call_params(redacted_params, params))
+    elif method in {"resources/read", "resources/list", "resources/templates/list", "tools/list"}:
+        lines.append("params:")
+        lines.append(pretty(redacted_params))
+    else:
+        lines.append("params:")
+        lines.append(pretty(redacted_params))
+
+    meta = redacted_params.get("_meta") if isinstance(redacted_params, dict) else None
+    if isinstance(meta, dict):
+        lines.append(f"meta_keys: {', '.join(sorted(meta.keys())) or '-'}")
+
+    return "\n".join(lines) + "\n"
+
+
+def describe_notification(message: types.JSONRPCNotification) -> str:
+    payload = message.model_dump(mode="json", by_alias=True, exclude_none=True)
+    params = payload.get("params") if isinstance(payload.get("params"), dict) else {}
+    return "\n".join(
+        [
+            f"[{timestamp()}] notification {payload.get('method')}",
+            f"params_bytes: {json_size(params)}",
+            "params:",
+            pretty(redact(params)),
+            "",
+        ]
+    )
+
+
+def describe_initialize_params(params: dict[str, Any]) -> list[str]:
+    client_info = params.get("clientInfo") if isinstance(params.get("clientInfo"), dict) else {}
+    capabilities = params.get("capabilities") if isinstance(params.get("capabilities"), dict) else {}
+    return [
+        f"protocol_version: {params.get('protocolVersion') or '-'}",
+        f"client_name: {client_info.get('name') or '-'}",
+        f"client_version: {client_info.get('version') or '-'}",
+        f"capability_keys: {', '.join(sorted(capabilities.keys())) or '-'}",
+        "params:",
+        pretty(params),
+    ]
+
+
+def describe_tool_call_params(params: dict[str, Any], raw_params: dict[str, Any]) -> list[str]:
+    arguments = params.get("arguments") if isinstance(params.get("arguments"), dict) else {}
+    raw_arguments = raw_params.get("arguments") if isinstance(raw_params.get("arguments"), dict) else {}
+    return [
+        f"tool_name: {params.get('name') or '-'}",
+        f"argument_keys: {', '.join(sorted(arguments.keys())) or '-'}",
+        f"argument_bytes: {json_size(raw_arguments)}",
+        "params:",
+        pretty(params),
+    ]
+
+
+def redact(value: Any) -> Any:
+    if isinstance(value, dict):
+        redacted: dict[str, Any] = {}
+        for key, nested in value.items():
+            key_text = str(key)
+            if is_sensitive_key(key_text):
+                redacted[key_text] = "[REDACTED]"
+            else:
+                redacted[key_text] = redact(nested)
+        return redacted
+    if isinstance(value, list):
+        return [redact(item) for item in value]
+    return value
+
+
+def is_sensitive_key(key: str) -> bool:
+    normalized = key.lower().replace("-", "_")
+    return any(fragment in normalized for fragment in SENSITIVE_KEY_FRAGMENTS)
+
+
+def json_size(value: Any) -> int:
+    return len(json.dumps(value or {}, ensure_ascii=False, separators=(",", ":"), sort_keys=True).encode("utf-8"))
+
+
+def pretty(value: Any) -> str:
+    return json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True)
+
+
+def timestamp() -> str:
+    return datetime.now(UTC).isoformat()