diff --git a/infra/mcp-server/.env.example b/infra/mcp-server/.env.example new file mode 100644 index 0000000..2803dc1 --- /dev/null +++ b/infra/mcp-server/.env.example @@ -0,0 +1,13 @@ +# Dedicated, read-only PostgreSQL login. Do not reuse the n8n or ingest role. +DB_HOST=ocsentinel-postgres +DB_PORT=5432 +DB_NAME=ocsentinel +DB_USER=ocsentinel_mcp +DB_PASSWORD=replace-with-a-long-random-password + +# A long random bearer token for trusted MCP clients. Keep this file private. +MCP_AUTH_TOKEN=replace-with-a-second-long-random-token + +# Validate browser origins and Host headers when they are present. +MCP_ALLOWED_ORIGINS=http://localhost:6274,http://127.0.0.1:6274 +MCP_ALLOWED_HOSTS=localhost:8091,127.0.0.1:8091 diff --git a/infra/mcp-server/.gitignore b/infra/mcp-server/.gitignore new file mode 100644 index 0000000..d50a09f --- /dev/null +++ b/infra/mcp-server/.gitignore @@ -0,0 +1,2 @@ +.env +__pycache__/ diff --git a/infra/mcp-server/Dockerfile b/infra/mcp-server/Dockerfile new file mode 100644 index 0000000..b4b72ae --- /dev/null +++ b/infra/mcp-server/Dockerfile @@ -0,0 +1,16 @@ +FROM python:3.13-alpine + +WORKDIR /app +ENV PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +COPY server.py . + +RUN addgroup -S ocsentinel && adduser -S ocsentinel -G ocsentinel +USER ocsentinel + +EXPOSE 8080 +CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8080", "--proxy-headers", "--no-access-log"] diff --git a/infra/mcp-server/README.md b/infra/mcp-server/README.md new file mode 100644 index 0000000..e139650 --- /dev/null +++ b/infra/mcp-server/README.md @@ -0,0 +1,70 @@ +# OfficeCom Sentinel MCP + +Dieser Container stellt sichere, **schreibgeschuetzte** Abfragen der OfficeCom-Sentinel-Daten fuer KI-Agenten bereit. Er nutzt das offizielle Python-MCP-SDK mit Streamable HTTP unter `/mcp`. + +## Sicherheitsmodell + +- Der Dienst wird im ersten Schritt nur auf `127.0.0.1:8091` des n8n-Hosts gebunden. Er wird nicht ueber `sentinel.officecom.biz` veroeffentlicht. +- Jeder MCP-Aufruf verlangt einen eigenen Bearer-Token, prueft `Host` sowie vorhandene `Origin`-Header und wird ohne Aufrufparameter protokolliert. +- Der PostgreSQL-Zugang ist ein dedizierter Login mit `default_transaction_read_only=on`, einem 5-Sekunden-Statement-Timeout und ausschliesslich `SELECT`-Rechten. +- Die Werkzeuge haben feste, parametrisierte Abfragen und feste Ergebnisgrenzen. Es gibt kein Werkzeug fuer SQL, Schreiboperationen, Rohbeweise, Befehlszeilen oder Zugangsdaten. +- Die Antwort auf `get_device_security` und `search_security_events` enthaelt standardmaessig keine Kontonamen. Konten werden nur auf ausdrueckliche Tool-Anforderung ergaenzt. + +## Verfuegbare Tools + +| Tool | Zweck | +| --- | --- | +| `security_overview` | Gesamtlage, Abdeckung und dringende Systeme | +| `get_organization_status` | Status eines NinjaOne-Organisations-IDs | +| `get_device_security` | Bereinigte Sicherheitslage eines Systems | +| `search_security_events` | Zeitlich und mengenmaessig begrenzte Ereigniszusammenfassungen | +| `get_network_paths` | Beobachtete Quell-IP-zu-System-Pfade | +| `get_weekly_report` | Letzte woechentliche Kennzahlen ohne Bericht-HTML | + +Zusaetzlich gibt es die Resource `ocsentinel://read-only-policy` und den Prompt `incident_triage`. + +## Einmalig: Datenbankrolle anlegen + +Auf dem PostgreSQL-Container als Datenbankadministrator ausfuehren. Das Passwort in diesem Befehl durch ein langes, zufaelliges Kennwort ersetzen und danach nur in der lokalen `.env` hinterlegen. + +```sql +CREATE ROLE ocsentinel_mcp LOGIN PASSWORD 'replace-with-a-long-random-password' + NOSUPERUSER NOCREATEDB NOCREATEROLE NOINHERIT; +GRANT CONNECT ON DATABASE ocsentinel TO ocsentinel_mcp; +GRANT USAGE ON SCHEMA ocsentinel TO ocsentinel_mcp; +GRANT SELECT ON ocsentinel.device, ocsentinel.scan_report, + ocsentinel.weekly_organization_report TO ocsentinel_mcp; +GRANT SELECT ON ocsentinel.current_device_status, + ocsentinel.organization_summary TO ocsentinel_mcp; +``` + +Pruefung: + +```sql +SET ROLE ocsentinel_mcp; +SELECT * FROM ocsentinel.organization_summary; +INSERT INTO ocsentinel.device (machine_name, machine_name_key) VALUES ('must-fail', 'must-fail'); +``` + +Die letzte Anweisung muss scheitern. + +## Dockge-Bereitstellung + +1. Den Ordner `infra/mcp-server` als neuen Dockge-Stack auf dem n8n-Host ablegen. +2. `.env.example` nach `.env` kopieren, Datenbankpasswort und einen zweiten langen Zufallstoken setzen. +3. In `MCP_ALLOWED_HOSTS` nur die echten, erlaubten Host-Header lassen. Fuer den SSH-Tunnel sind `localhost:8091` und `127.0.0.1:8091` korrekt. +4. Stack starten. Der Endpunkt ist lokal: `http://127.0.0.1:8091/mcp`. + +Der Container hat keinen veroeffentlichten Zugriff auf das Internet. Fuer einen Arbeitsplatz wird ein Tunnel genutzt: + +```powershell +ssh -L 8091:127.0.0.1:8091 oc@172.16.41.197 -p 1022 +``` + +Danach ist der lokale MCP-Endpunkt `http://localhost:8091/mcp`. Der MCP-Client muss den Header `Authorization: Bearer ` mitsenden. + +## Betrieb + +- Logs: `docker logs ocsentinel-mcp --tail 100`. +- Niemals den Bearer-Token in einem Git-Repository, Screenshot oder Prompt speichern. +- Fuer einen spaeteren externen Zugriff wird ein separater OAuth-geschuetzter Reverse Proxy benoetigt. Der aktuelle Token-Modus ist ausschliesslich fuer den privaten Tunnel und vertrauenswuerdige Agenten gedacht. diff --git a/infra/mcp-server/compose.yml b/infra/mcp-server/compose.yml new file mode 100644 index 0000000..64d9c8c --- /dev/null +++ b/infra/mcp-server/compose.yml @@ -0,0 +1,23 @@ +services: + ocsentinel-mcp: + build: . + container_name: ocsentinel-mcp + restart: unless-stopped + env_file: .env + # The first version is intentionally only reachable through an SSH tunnel. + ports: + - "127.0.0.1:8091:8080" + networks: + - ocsentinel-network + read_only: true + tmpfs: + - /tmp + security_opt: + - no-new-privileges:true + cap_drop: + - ALL + +networks: + ocsentinel-network: + external: true + name: n8n_n8n-network diff --git a/infra/mcp-server/requirements.txt b/infra/mcp-server/requirements.txt new file mode 100644 index 0000000..edb5232 --- /dev/null +++ b/infra/mcp-server/requirements.txt @@ -0,0 +1,3 @@ +mcp==1.26.0 +psycopg[binary]==3.2.9 +uvicorn==0.35.0 diff --git a/infra/mcp-server/server.py b/infra/mcp-server/server.py new file mode 100644 index 0000000..546db4e --- /dev/null +++ b/infra/mcp-server/server.py @@ -0,0 +1,478 @@ +"""Read-only MCP access to curated OfficeCom Sentinel security data.""" + +import functools +import json +import logging +import os +import re +import time +from datetime import datetime +from typing import Any + +import psycopg +from psycopg.rows import dict_row +from mcp.server.fastmcp import FastMCP + + +LOGGER = logging.getLogger("ocsentinel.mcp") +logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO").upper(), format="%(asctime)s %(levelname)s %(message)s") + +MAX_RESULT_LIMIT = 100 +MAX_LOOKBACK_HOURS = 24 * 90 +IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z0-9._:-]{1,128}$") + + +def required_setting(name: str) -> str: + value = os.getenv(name, "").strip() + if not value or value.startswith("replace-with-"): + raise RuntimeError(f"{name} must be configured before starting OCSentinel MCP.") + return value + + +AUTH_TOKEN = required_setting("MCP_AUTH_TOKEN") +ALLOWED_ORIGINS = {value.strip() for value in os.getenv("MCP_ALLOWED_ORIGINS", "").split(",") if value.strip()} +ALLOWED_HOSTS = {value.strip().lower() for value in os.getenv("MCP_ALLOWED_HOSTS", "").split(",") if value.strip()} + + +def db_connection() -> psycopg.Connection: + return psycopg.connect( + host=required_setting("DB_HOST"), + port=os.getenv("DB_PORT", "5432"), + dbname=required_setting("DB_NAME"), + user=required_setting("DB_USER"), + password=required_setting("DB_PASSWORD"), + connect_timeout=5, + row_factory=dict_row, + options="-c default_transaction_read_only=on -c statement_timeout=5000", + ) + + +def rows(sql: str, parameters: tuple[Any, ...] = ()) -> list[dict[str, Any]]: + with db_connection() as connection, connection.cursor() as cursor: + cursor.execute(sql, parameters) + return list(cursor.fetchall()) + + +def row(sql: str, parameters: tuple[Any, ...] = ()) -> dict[str, Any] | None: + results = rows(sql, parameters) + return results[0] if results else None + + +def value(record: dict[str, Any] | None, *names: str, default: Any = None) -> Any: + if not isinstance(record, dict): + return default + for name in names: + if name in record: + return record[name] + return default + + +def json_safe(data: Any) -> Any: + return json.loads(json.dumps(data, default=lambda entry: entry.isoformat() if isinstance(entry, datetime) else str(entry))) + + +def bounded_limit(limit: int) -> int: + if not isinstance(limit, int) or isinstance(limit, bool): + raise ValueError("limit must be a whole number.") + return max(1, min(limit, MAX_RESULT_LIMIT)) + + +def bounded_hours(hours: int) -> int: + if not isinstance(hours, int) or isinstance(hours, bool): + raise ValueError("hours must be a whole number.") + return max(1, min(hours, MAX_LOOKBACK_HOURS)) + + +def checked_identifier(identifier: str, field_name: str) -> str: + value_to_check = (identifier or "").strip() + if value_to_check == "unknown" or IDENTIFIER_PATTERN.fullmatch(value_to_check): + return value_to_check + raise ValueError(f"{field_name} contains unsupported characters.") + + +def checked_machine_name(machine_name: str) -> str: + machine = (machine_name or "").strip() + if not machine or len(machine) > 255 or any(character in machine for character in "\r\n\x00"): + raise ValueError("machine_name must be a single device name of at most 255 characters.") + return machine + + +def organization_context(payload: dict[str, Any] | None) -> tuple[str, str]: + ninja = value(payload, "NinjaOne", "ninjaOne", default={}) or {} + return ( + str(value(ninja, "OrganizationId", "organizationId", default="unknown") or "unknown"), + str(value(ninja, "OrganizationName", "organizationName", default="Organisation unbekannt") or "Organisation unbekannt"), + ) + + +def sanitized_ransomware(payload: dict[str, Any] | None) -> dict[str, Any]: + ransomware = value(payload, "RansomwareBeta", "ransomwareBeta", default={}) or {} + sensors = [] + for sensor in value(ransomware, "Sensors", "sensors", default=[]) or []: + sensors.append( + { + "name": value(sensor, "Name", "name", default="unknown"), + "enabled": bool(value(sensor, "Enabled", "enabled", default=False)), + "available": bool(value(sensor, "Available", "available", default=False)), + "state": value(sensor, "State", "state", default="unknown"), + "eventCount": value(sensor, "EventCount", "eventCount", default=0), + } + ) + + signals = [] + for signal in value(ransomware, "Signals", "signals", default=[]) or []: + # Evidence and raw command lines intentionally never leave the MCP boundary. + signals.append( + { + "timestamp": value(signal, "Timestamp", "timestamp", default=None), + "category": value(signal, "Category", "category", default="signal"), + "process": value(signal, "Process", "process", default="-"), + "source": value(signal, "Source", "source", default="-"), + "confidence": value(signal, "Confidence", "confidence", default="low"), + } + ) + + return { + "enabled": bool(value(ransomware, "Enabled", "enabled", default=False)), + "state": value(ransomware, "State", "state", default="disabled"), + "reason": value(ransomware, "Reason", "reason", default="Keine Ransomware-Beta-Daten verfuegbar."), + "sensors": sensors, + "signals": signals[:20], + } + + +def summarized_events(payload: dict[str, Any] | None, include_accounts: bool = False) -> list[dict[str, Any]]: + groups: dict[tuple[str, str, str], dict[str, Any]] = {} + for event in value(payload, "Events", "events", default=[]) or []: + event_type = str(value(event, "Target", "target", default="Sicherheitsereignis")) + source_ip = str(value(event, "SourceIp", "sourceIp", default="-")) + account = str(value(event, "Username", "username", default="-")) if include_accounts else "" + key = (event_type, source_ip, account) + group = groups.setdefault(key, {"type": event_type, "sourceIp": source_ip, "count": 0, "latest": None}) + group["count"] += 1 + timestamp = value(event, "Timestamp", "timestamp", default=None) + if timestamp and (not group["latest"] or str(timestamp) > str(group["latest"])): + group["latest"] = timestamp + if include_accounts: + group["account"] = account + return sorted(groups.values(), key=lambda entry: (str(entry["latest"]), entry["count"]), reverse=True)[:25] + + +def audited(tool_name: str): + def decorator(function): + @functools.wraps(function) + def wrapped(*args, **kwargs): + started = time.monotonic() + try: + result = function(*args, **kwargs) + LOGGER.info("mcp_tool=%s outcome=ok duration_ms=%d", tool_name, (time.monotonic() - started) * 1000) + return result + except Exception: + LOGGER.exception("mcp_tool=%s outcome=error duration_ms=%d", tool_name, (time.monotonic() - started) * 1000) + raise + return wrapped + return decorator + + +mcp = FastMCP( + "OfficeCom Sentinel", + instructions=( + "Read-only security context from OfficeCom Sentinel. Use this data to investigate, summarize, and prioritize. " + "Do not treat it as authorization to alter devices, NinjaOne, PostgreSQL, or security controls." + ), + stateless_http=True, + json_response=True, + streamable_http_path="/mcp", +) + + +@mcp.resource("ocsentinel://read-only-policy") +def read_only_policy() -> str: + """Explain the data and safety boundary of this server.""" + return ( + "OfficeCom Sentinel MCP is read-only. It returns curated status, event summaries, ransomware sensor coverage, " + "network paths, and weekly-report metadata. Raw event evidence, command lines, SMB sessions, credentials, " + "and every write action are deliberately excluded." + ) + + +@mcp.prompt() +def incident_triage() -> str: + """Provide a safe, evidence-oriented workflow for analyzing Sentinel findings.""" + return ( + "Start with security_overview or get_organization_status. For a flagged device, use get_device_security and " + "search_security_events. Separate observed facts from hypotheses, identify the next reversible validation step, " + "and recommend escalation to the responsible OfficeCom technician for any containment action." + ) + + +@mcp.tool() +@audited("security_overview") +def security_overview() -> dict[str, Any]: + """Return the current cross-organization security posture and the most urgent devices.""" + summary = row("SELECT * FROM ocsentinel.organization_summary") or {} + coverage = row( + """ + SELECT count(*) AS known_devices, + count(*) FILTER (WHERE received_at >= now() - interval '36 hours') AS reporting_devices, + count(*) FILTER (WHERE received_at IS NULL OR received_at < now() - interval '36 hours') AS stale_devices + FROM ocsentinel.current_device_status + """ + ) or {} + urgent = rows( + """ + SELECT machine_name, alert_state, received_at, total_events, unique_ip_count, + payload #>> '{NinjaOne,OrganizationId}' AS organization_id, + payload #>> '{NinjaOne,OrganizationName}' AS organization_name + FROM ocsentinel.current_device_status + WHERE alert_state IN ('warning', 'critical') + ORDER BY CASE alert_state WHEN 'critical' THEN 0 ELSE 1 END, received_at DESC NULLS LAST + LIMIT 20 + """ + ) + return json_safe({"summary": summary, "coverage": coverage, "urgentDevices": urgent}) + + +@mcp.tool() +@audited("get_organization_status") +def get_organization_status(organization_id: str) -> dict[str, Any]: + """Return current coverage and alert state for one NinjaOne organization ID.""" + organization_id = checked_identifier(organization_id, "organization_id") + devices = rows( + """ + SELECT machine_name, received_at, alert_state, base_alert_state, total_events, unique_ip_count, + cve_total, cve_critical, payload + FROM ocsentinel.current_device_status + WHERE coalesce(payload #>> '{NinjaOne,OrganizationId}', 'unknown') = %s + ORDER BY CASE alert_state WHEN 'critical' THEN 0 WHEN 'warning' THEN 1 ELSE 2 END, machine_name + """, + (organization_id,), + ) + if not devices: + return {"organizationId": organization_id, "found": False, "devices": []} + _, organization_name = organization_context(devices[0]["payload"]) + status_counts = {state: sum(device["alert_state"] == state for device in devices) for state in ("ok", "warning", "critical", "unknown")} + return json_safe( + { + "organizationId": organization_id, + "organizationName": organization_name, + "found": True, + "deviceCount": len(devices), + "statusCounts": status_counts, + "lastReceivedAt": max((device["received_at"] for device in devices if device["received_at"]), default=None), + "devices": [ + { + "machineName": device["machine_name"], + "alertState": device["alert_state"], + "receivedAt": device["received_at"], + "totalEvents": device["total_events"], + "uniqueIpCount": device["unique_ip_count"], + "criticalCves": device["cve_critical"], + } + for device in devices + ], + } + ) + + +@mcp.tool() +@audited("get_device_security") +def get_device_security(machine_name: str, include_accounts: bool = False) -> dict[str, Any]: + """Return the latest sanitized security summary for one device. Accounts are omitted by default.""" + machine_name = checked_machine_name(machine_name) + device = row( + """ + SELECT machine_name, first_seen_at, last_seen_at, last_client_version, generated_at_utc, received_at, + alert_state, base_alert_state, total_events, unique_ip_count, cve_total, cve_critical, payload + FROM ocsentinel.current_device_status + WHERE machine_name = %s + """, + (machine_name,), + ) + if not device: + return {"machineName": machine_name, "found": False} + organization_id, organization_name = organization_context(device["payload"]) + return json_safe( + { + "found": True, + "machineName": device["machine_name"], + "organizationId": organization_id, + "organizationName": organization_name, + "receivedAt": device["received_at"], + "generatedAt": device["generated_at_utc"], + "alertState": device["alert_state"], + "baseAlertState": device["base_alert_state"], + "metrics": { + "totalEvents": device["total_events"], + "uniqueIpCount": device["unique_ip_count"], + "cveTotal": device["cve_total"], + "criticalCves": device["cve_critical"], + }, + "events": summarized_events(device["payload"], include_accounts), + "ransomwareBeta": sanitized_ransomware(device["payload"]), + } + ) + + +@mcp.tool() +@audited("search_security_events") +def search_security_events( + hours: int = 168, + organization_id: str | None = None, + alert_state: str | None = None, + limit: int = 25, + include_accounts: bool = False, +) -> dict[str, Any]: + """Find recent scanned reports with bounded filters. This returns event summaries, not raw evidence.""" + hours = bounded_hours(hours) + limit = bounded_limit(limit) + if organization_id is not None: + organization_id = checked_identifier(organization_id, "organization_id") + if alert_state is not None and alert_state not in {"ok", "warning", "critical", "unknown"}: + raise ValueError("alert_state must be ok, warning, critical, or unknown.") + + reports = rows( + """ + SELECT d.machine_name, r.received_at, r.generated_at_utc, r.alert_state, r.base_alert_state, + r.total_events, r.unique_ip_count, r.cve_total, r.cve_critical, r.payload + FROM ocsentinel.scan_report AS r + JOIN ocsentinel.device AS d ON d.id = r.device_id + WHERE r.received_at >= now() - (%s * interval '1 hour') + AND (%s::text IS NULL OR coalesce(r.payload #>> '{NinjaOne,OrganizationId}', 'unknown') = %s) + AND (%s::text IS NULL OR r.alert_state = %s) + ORDER BY r.received_at DESC + LIMIT %s + """, + (hours, organization_id, organization_id, alert_state, alert_state, limit), + ) + results = [] + for report in reports: + org_id, org_name = organization_context(report["payload"]) + results.append( + { + "machineName": report["machine_name"], + "organizationId": org_id, + "organizationName": org_name, + "receivedAt": report["received_at"], + "generatedAt": report["generated_at_utc"], + "alertState": report["alert_state"], + "baseAlertState": report["base_alert_state"], + "totalEvents": report["total_events"], + "uniqueIpCount": report["unique_ip_count"], + "criticalCves": report["cve_critical"], + "events": summarized_events(report["payload"], include_accounts), + "ransomwareBeta": sanitized_ransomware(report["payload"]), + } + ) + return json_safe({"lookbackHours": hours, "resultCount": len(results), "reports": results}) + + +@mcp.tool() +@audited("get_network_paths") +def get_network_paths(days: int = 14, organization_id: str | None = None, limit: int = 50) -> dict[str, Any]: + """Return recent observed source-IP to device paths from summarized security events.""" + days = max(1, min(bounded_hours(days * 24) // 24, 90)) + limit = bounded_limit(limit) + if organization_id is not None: + organization_id = checked_identifier(organization_id, "organization_id") + + reports = rows( + """ + SELECT machine_name, received_at, payload + FROM ocsentinel.current_device_status + WHERE received_at >= now() - (%s * interval '1 day') + AND (%s::text IS NULL OR coalesce(payload #>> '{NinjaOne,OrganizationId}', 'unknown') = %s) + """, + (days, organization_id, organization_id), + ) + flows: dict[tuple[str, str, str, str], dict[str, Any]] = {} + for report in reports: + org_id, org_name = organization_context(report["payload"]) + for event in value(report["payload"], "Events", "events", default=[]) or []: + source_ip = str(value(event, "SourceIp", "sourceIp", default="")) + if not source_ip or source_ip in {"-", "127.0.0.1", "::1"}: + continue + event_type = str(value(event, "Target", "target", default="Sicherheitsereignis")) + key = (source_ip, report["machine_name"], event_type, org_id) + flow = flows.setdefault( + key, + { + "sourceIp": source_ip, + "machineName": report["machine_name"], + "eventType": event_type, + "organizationId": org_id, + "organizationName": org_name, + "count": 0, + "lastSeen": report["received_at"], + }, + ) + flow["count"] += 1 + timestamp = value(event, "Timestamp", "timestamp", default=None) + if timestamp and (not flow["lastSeen"] or str(timestamp) > str(flow["lastSeen"])): + flow["lastSeen"] = timestamp + paths = sorted(flows.values(), key=lambda entry: (entry["count"], str(entry["lastSeen"])), reverse=True)[:limit] + return json_safe({"days": days, "pathCount": len(paths), "paths": paths}) + + +@mcp.tool() +@audited("get_weekly_report") +def get_weekly_report(organization_id: str) -> dict[str, Any]: + """Return the latest weekly report metadata and structured summary for one organization, without report HTML.""" + organization_id = checked_identifier(organization_id, "organization_id") + report = row( + """ + SELECT organization_id, organization_name, period_start_utc, period_end_utc, generated_at, + device_count, warning_count, critical_count, total_events, unique_ips, cve_total, + cve_critical, summary + FROM ocsentinel.weekly_organization_report + WHERE organization_id = %s + ORDER BY period_end_utc DESC, generated_at DESC + LIMIT 1 + """, + (organization_id,), + ) + if not report: + return {"organizationId": organization_id, "found": False} + report["found"] = True + return json_safe(report) + + +class GuardedMcpApp: + """Small ASGI guard without BaseHTTPMiddleware, which can disrupt MCP streaming.""" + + def __init__(self, wrapped_app): + self.wrapped_app = wrapped_app + + async def __call__(self, scope, receive, send): + if scope["type"] != "http": + await self.wrapped_app(scope, receive, send) + return + + headers = {key.decode("latin-1").lower(): value.decode("latin-1") for key, value in scope.get("headers", [])} + host = headers.get("host", "").lower() + origin = headers.get("origin") + authorization = headers.get("authorization", "") + + if ALLOWED_HOSTS and host not in ALLOWED_HOSTS: + await self.reject(send, 421, "Untrusted Host header.") + return + if origin and (not ALLOWED_ORIGINS or origin not in ALLOWED_ORIGINS): + await self.reject(send, 403, "Untrusted Origin header.") + return + if authorization != f"Bearer {AUTH_TOKEN}": + await self.reject(send, 401, "Bearer token required.", {b"www-authenticate": b"Bearer"}) + return + + await self.wrapped_app(scope, receive, send) + + @staticmethod + async def reject(send, status: int, message: str, extra_headers: dict[bytes, bytes] | None = None): + body = json.dumps({"error": message}).encode("utf-8") + headers = [(b"content-type", b"application/json"), (b"content-length", str(len(body)).encode("ascii"))] + if extra_headers: + headers.extend(extra_headers.items()) + await send({"type": "http.response.start", "status": status, "headers": headers}) + await send({"type": "http.response.body", "body": body}) + + +app = GuardedMcpApp(mcp.streamable_http_app())