| 1 | import asyncio |
| 2 | from shutil import which |
| 3 | from typing import Any |
| 4 | from urllib.parse import urlparse |
| 5 | |
| 6 | from helpers.api import ApiHandler, Request, Response |
| 7 | from helpers.mcp_handler import MCPConfig, normalize_name |
| 8 | |
| 9 | |
| 10 | _PROMPT_INJECTION_MARKERS = ( |
| 11 | "ignore previous", |
| 12 | "ignore all previous", |
| 13 | "system prompt", |
| 14 | "developer message", |
| 15 | "hidden instruction", |
| 16 | "exfiltrate", |
| 17 | "leak secret", |
| 18 | "credential", |
| 19 | ) |
| 20 | |
| 21 | |
| 22 | class McpServerScan(ApiHandler): |
| 23 | async def process(self, input: dict[Any, Any], request: Request) -> dict[Any, Any] | Response: |
| 24 | server = dict(input.get("server") or {}) |
| 25 | allow_local_execution = bool(input.get("allow_local_execution", False)) |
| 26 | allow_remote_network = bool(input.get("allow_remote_network", False)) |
| 27 | inspect_runtime = input.get("inspect_runtime", True) is not False |
| 28 | |
| 29 | server = self._normalize_server(server) |
| 30 | warnings = self._static_warnings(server) |
| 31 | is_local = not (server.get("url") or server.get("serverUrl")) |
| 32 | has_static_errors = any(warning.get("level") == "error" for warning in warnings) |
| 33 | |
| 34 | runtime_status: list[dict[str, Any]] = [] |
| 35 | runtime_detail: dict[str, Any] = {} |
| 36 | runtime_error = "" |
| 37 | |
| 38 | should_inspect_runtime = ( |
| 39 | inspect_runtime |
| 40 | and not has_static_errors |
| 41 | and ((is_local and allow_local_execution) or (not is_local and allow_remote_network)) |
| 42 | ) |
| 43 | |
| 44 | if should_inspect_runtime: |
| 45 | try: |
| 46 | scan_config = await asyncio.to_thread( |
| 47 | lambda: MCPConfig(servers_list=[server], config_scope="scan") |
| 48 | ) |
| 49 | runtime_status = scan_config.get_servers_status() |
| 50 | runtime_detail = scan_config.get_server_detail(server.get("name", "")) |
| 51 | warnings.extend(self._tool_warnings(runtime_detail.get("tools", []))) |
| 52 | except Exception as exc: |
| 53 | runtime_error = str(exc) |
| 54 | warnings.append( |
| 55 | { |
| 56 | "level": "error", |
| 57 | "title": "Runtime inspection failed", |
| 58 | "message": runtime_error, |
| 59 | } |
| 60 | ) |
| 61 | elif is_local and inspect_runtime: |
| 62 | warnings.append( |
| 63 | { |
| 64 | "level": "warning", |
| 65 | "title": "Local command not executed", |
| 66 | "message": "Local stdio MCP inspection requires explicit trust because it runs the configured command.", |
| 67 | } |
| 68 | ) |
| 69 | elif not is_local and inspect_runtime and has_static_errors: |
| 70 | warnings.append( |
| 71 | { |
| 72 | "level": "info", |
| 73 | "title": "Runtime inspection skipped", |
| 74 | "message": "Fix static scan errors before attempting runtime MCP inspection.", |
| 75 | } |
| 76 | ) |
| 77 | elif not is_local and inspect_runtime: |
| 78 | warnings.append( |
| 79 | { |
| 80 | "level": "info", |
| 81 | "title": "Remote runtime inspection skipped", |
| 82 | "message": "Enable trusted remote inspection to contact the MCP URL and list exposed tools.", |
| 83 | } |
| 84 | ) |
| 85 | |
| 86 | return { |
| 87 | "success": True, |
| 88 | "server": self._redact_server(server), |
| 89 | "risk_level": self._risk_level(warnings), |
| 90 | "warnings": warnings, |
| 91 | "status": runtime_status, |
| 92 | "detail": runtime_detail, |
| 93 | "runtime_error": runtime_error, |
| 94 | } |
| 95 | |
| 96 | def _normalize_server(self, server: dict[str, Any]) -> dict[str, Any]: |
| 97 | name = str(server.get("name") or "").strip() |
| 98 | url = str(server.get("url") or server.get("serverUrl") or "").strip() |
| 99 | command = str(server.get("command") or "").strip() |
| 100 | |
| 101 | if not name: |
| 102 | name = self._derive_name(url, command) |
| 103 | server["name"] = normalize_name(name or "mcp_server") |
| 104 | |
| 105 | if url: |
| 106 | server["url"] = url |
| 107 | server.setdefault("type", "streamable-http") |
| 108 | elif command: |
| 109 | server["command"] = command |
| 110 | server["type"] = "stdio" |
| 111 | |
| 112 | return server |
| 113 | |
| 114 | def _derive_name(self, url: str, command: str) -> str: |
| 115 | if url: |
| 116 | parsed = urlparse(url) |
| 117 | parts = [part for part in parsed.path.split("/") if part] |
| 118 | return parts[-1] if parts else parsed.hostname or "remote_mcp" |
| 119 | if command: |
| 120 | return command.rsplit("/", 1)[-1] |
| 121 | return "mcp_server" |
| 122 | |
| 123 | def _static_warnings(self, server: dict[str, Any]) -> list[dict[str, str]]: |
| 124 | warnings: list[dict[str, str]] = [] |
| 125 | url = str(server.get("url") or "").strip() |
| 126 | command = str(server.get("command") or "").strip() |
| 127 | |
| 128 | if url: |
| 129 | parsed = urlparse(url) |
| 130 | if parsed.scheme not in {"http", "https"}: |
| 131 | warnings.append( |
| 132 | { |
| 133 | "level": "error", |
| 134 | "title": "Unsupported URL scheme", |
| 135 | "message": "Remote MCP URLs should use http or https.", |
| 136 | } |
| 137 | ) |
| 138 | elif parsed.scheme == "http" and parsed.hostname not in {"localhost", "127.0.0.1", "::1"}: |
| 139 | warnings.append( |
| 140 | { |
| 141 | "level": "warning", |
| 142 | "title": "Unencrypted remote URL", |
| 143 | "message": "Prefer HTTPS for remote MCP servers outside localhost.", |
| 144 | } |
| 145 | ) |
| 146 | if not parsed.netloc: |
| 147 | warnings.append( |
| 148 | { |
| 149 | "level": "error", |
| 150 | "title": "Invalid remote URL", |
| 151 | "message": "The remote MCP URL is missing a host.", |
| 152 | } |
| 153 | ) |
| 154 | elif command: |
| 155 | if which(command) is None: |
| 156 | warnings.append( |
| 157 | { |
| 158 | "level": "warning", |
| 159 | "title": "Command not found", |
| 160 | "message": f"'{command}' is not currently available on PATH.", |
| 161 | } |
| 162 | ) |
| 163 | if command in {"bash", "sh", "zsh", "fish", "python", "python3", "node"}: |
| 164 | warnings.append( |
| 165 | { |
| 166 | "level": "warning", |
| 167 | "title": "General-purpose interpreter", |
| 168 | "message": "Review the command and arguments carefully before running this local MCP server.", |
| 169 | } |
| 170 | ) |
| 171 | else: |
| 172 | warnings.append( |
| 173 | { |
| 174 | "level": "error", |
| 175 | "title": "Missing connection target", |
| 176 | "message": "Provide either a remote URL or a local command.", |
| 177 | } |
| 178 | ) |
| 179 | |
| 180 | if isinstance(server.get("headers"), dict) and server["headers"]: |
| 181 | warnings.append( |
| 182 | { |
| 183 | "level": "info", |
| 184 | "title": "Headers configured", |
| 185 | "message": "Header values are redacted in scan output. Keep tokens in trusted settings only.", |
| 186 | } |
| 187 | ) |
| 188 | |
| 189 | if isinstance(server.get("env"), dict) and server["env"]: |
| 190 | warnings.append( |
| 191 | { |
| 192 | "level": "info", |
| 193 | "title": "Environment configured", |
| 194 | "message": "Environment values are redacted in scan output. Avoid hardcoding secrets in MCP configs.", |
| 195 | } |
| 196 | ) |
| 197 | |
| 198 | return warnings |
| 199 | |
| 200 | def _tool_warnings(self, tools: Any) -> list[dict[str, str]]: |
| 201 | warnings: list[dict[str, str]] = [] |
| 202 | if not isinstance(tools, list): |
| 203 | return warnings |
| 204 | |
| 205 | for tool in tools: |
| 206 | if not isinstance(tool, dict): |
| 207 | continue |
| 208 | haystack = f"{tool.get('name', '')}\n{tool.get('description', '')}".lower() |
| 209 | if any(marker in haystack for marker in _PROMPT_INJECTION_MARKERS): |
| 210 | warnings.append( |
| 211 | { |
| 212 | "level": "warning", |
| 213 | "title": "Suspicious tool description", |
| 214 | "message": f"Review tool '{tool.get('name', 'unknown')}' for prompt-injection style language.", |
| 215 | } |
| 216 | ) |
| 217 | return warnings |
| 218 | |
| 219 | def _redact_server(self, server: dict[str, Any]) -> dict[str, Any]: |
| 220 | redacted = dict(server) |
| 221 | for key in ("headers", "env"): |
| 222 | if isinstance(redacted.get(key), dict): |
| 223 | redacted[key] = {name: "***" for name in redacted[key]} |
| 224 | return redacted |
| 225 | |
| 226 | def _risk_level(self, warnings: list[dict[str, str]]) -> str: |
| 227 | levels = {warning.get("level", "info") for warning in warnings} |
| 228 | if "error" in levels: |
| 229 | return "error" |
| 230 | if "warning" in levels: |
| 231 | return "warning" |
| 232 | return "ok" |