diff --git a/Detection/context_providers/source_code_analyzer_server.py b/Detection/context_providers/source_code_analyzer_server.py index 5abc52f..72df392 100644 --- a/Detection/context_providers/source_code_analyzer_server.py +++ b/Detection/context_providers/source_code_analyzer_server.py @@ -5,11 +5,19 @@ Provides MCP server source code for reasoning-based threat analysis. No pre-analysis or cheating metadata - just clean source code and basic info. """ -import yaml +import ast import logging +import os +from importlib.machinery import ( + BYTECODE_SUFFIXES, + EXTENSION_SUFFIXES, + BuiltinImporter, + FrozenImporter, +) from pathlib import Path -from typing import Dict, List, Any +from typing import Any, Dict, List, Optional, Tuple +import yaml from mcp.server.fastmcp import FastMCP # Configure logging @@ -39,6 +47,235 @@ def load_source_registry() -> Dict[str, Any]: source_registry = load_source_registry() + +MAX_LOCAL_SOURCE_FILES = 24 +MAX_LOCAL_SOURCE_CHARS = 240_000 + + +def _is_within(path: Path, root: Path) -> bool: + """Return whether a resolved path remains within the trusted server root.""" + + try: + path.resolve(strict=True).relative_to(root.resolve(strict=True)) + return True + except (FileNotFoundError, OSError, RuntimeError, ValueError): + return False + + +def _resolved_source_file(path: Path, root: Path) -> Optional[Path]: + """Return a regular in-root source file without following symlinked input.""" + + if not _is_within(path, root): + return None + try: + resolved = path.resolve(strict=True) + except (FileNotFoundError, OSError, RuntimeError): + return None + if resolved != Path(os.path.abspath(os.fspath(path))): + return None + return resolved if resolved.is_file() else None + + +def _has_file_with_suffixes(path: Path, suffixes: List[str]) -> bool: + """Return whether a competing loader file exists for a module path.""" + + return any(path.with_name(path.name + suffix).is_file() for suffix in suffixes) + + +def _runtime_import_precedes_local_source(module_name: str) -> bool: + """Return whether Python resolves a built-in or frozen module first.""" + + top_level = module_name.partition(".")[0] + if not top_level: + return False + return ( + BuiltinImporter.find_spec(top_level) is not None + or FrozenImporter.find_spec(top_level) is not None + ) + + +def _resolve_module_parts( + base: Path, + parts: List[str], + root: Path, +) -> Tuple[List[Path], Optional[Path]]: + """Resolve local source using Python's package-before-module precedence. + + The returned package directory is used only to resolve ``from package + import submodule``. Native-extension and bytecode-only modules are not + represented as Python source and block same-name ``.py`` candidates when + Python would select them first. + """ + + files: List[Path] = [] + if not parts: + initializer = _resolved_source_file(base / "__init__.py", root) + if initializer is not None: + files.append(initializer) + return files, base if _is_within(base, root) and base.is_dir() else None + + current = base + for index, part in enumerate(parts): + final = index == len(parts) - 1 + package_dir = current / part + module_stem = current / part + + # FileFinder checks a package directory before same-named module files. + if package_dir.is_dir() and _is_within(package_dir, root): + initializer_stem = package_dir / "__init__" + has_extension_initializer = _has_file_with_suffixes( + initializer_stem, EXTENSION_SUFFIXES + ) + has_source_initializer = (package_dir / "__init__.py").is_file() + has_bytecode_initializer = _has_file_with_suffixes( + initializer_stem, BYTECODE_SUFFIXES + ) + if has_extension_initializer: + initializer = None + else: + initializer = _resolved_source_file( + package_dir / "__init__.py", root + ) + if ( + has_extension_initializer + or has_source_initializer + or has_bytecode_initializer + ): + if has_source_initializer and initializer is None: + return files, None + if initializer is not None: + files.append(initializer) + if final: + return files, package_dir + current = package_dir + continue + + # Native extensions precede source modules. Source precedes bytecode. + if _has_file_with_suffixes(module_stem, EXTENSION_SUFFIXES): + return files, None + module = _resolved_source_file(module_stem.with_suffix(".py"), root) + if module is not None: + files.append(module) + return files, None + if _has_file_with_suffixes(module_stem, BYTECODE_SUFFIXES): + return files, None + + # A directory without an initializer becomes a namespace package only + # if no same-name module loader matched above. + if package_dir.is_dir() and _is_within(package_dir, root): + if final: + return files, package_dir + current = package_dir + continue + return files, None + + return files, None + + +def _local_import_candidates(current: Path, node: ast.AST, root: Path) -> List[Path]: + """Resolve Python imports that refer to files inside one registered server.""" + + candidates: List[Path] = [] + candidate_set = set() + + def add(paths: List[Path]) -> None: + for path in paths: + if path not in candidate_set: + candidates.append(path) + candidate_set.add(path) + + if isinstance(node, ast.Import): + for alias in node.names: + if _runtime_import_precedes_local_source(alias.name): + continue + files, _ = _resolve_module_parts( + root, + [part for part in alias.name.split(".") if part], + root, + ) + add(files) + return candidates + + if isinstance(node, ast.ImportFrom): + module = node.module or "" + level = int(node.level or 0) + if level: + base = current.parent + for _ in range(level - 1): + base = base.parent + else: + base = root + + if not _is_within(base, root): + return candidates + if not level and _runtime_import_precedes_local_source(module): + return candidates + + files, package_dir = _resolve_module_parts( + base, + [part for part in module.split(".") if part], + root, + ) + add(files) + if package_dir is not None: + for alias in node.names: + if alias.name == "*": + continue + alias_files, _ = _resolve_module_parts( + package_dir, + [part for part in alias.name.split(".") if part], + root, + ) + add(alias_files) + return candidates + + +def collect_local_source_files(entrypoint: Path) -> List[Dict[str, Any]]: + """Collect a bounded, dependency-aware source bundle for one MCP server. + + Registry paths are trusted, but import resolution is constrained to the + entrypoint's server directory. Returned paths are relative so dataset + directory names do not become classifier hints. + """ + + root = entrypoint.parent.resolve() + entrypoint = _resolved_source_file(entrypoint, root) + if entrypoint is None: + return [] + pending = [entrypoint] + seen = set() + files: List[Dict[str, str]] = [] + total_chars = 0 + while pending and len(files) < MAX_LOCAL_SOURCE_FILES: + path = pending.pop(0) + path = _resolved_source_file(path, root) + if path is None or path in seen: + continue + seen.add(path) + source = path.read_text(encoding="utf-8") + remaining = MAX_LOCAL_SOURCE_CHARS - total_chars + if remaining <= 0: + break + truncated = len(source) > remaining + included = source[:remaining] + files.append({ + "path": str(path.relative_to(root)), + "source_code": included, + "truncated": truncated, + }) + total_chars += len(included) + if truncated: + break + try: + tree = ast.parse(source, filename=str(path)) + except (SyntaxError, ValueError): + continue + for node in ast.walk(tree): + for dependency in _local_import_candidates(path, node, root): + if dependency not in seen and dependency not in pending: + pending.append(dependency) + return files + @mcp.tool() def get_source_code(server_names: List[str]) -> Dict[str, Any]: """Get the source code of MCP servers for analysis""" @@ -61,12 +298,15 @@ def get_source_code(server_names: List[str]) -> Dict[str, Any]: # Read source code server_path = server_info.get("path", "") - full_path = Path(__file__).parent / server_path + provider_root = Path(__file__).parent.resolve() + full_path = provider_root / server_path try: - if full_path.exists(): - with open(full_path, 'r') as f: + resolved_path = _resolved_source_file(full_path, provider_root) + if resolved_path is not None: + with open(resolved_path, 'r') as f: source_code = f.read() + source_files = collect_local_source_files(resolved_path) # Provide clean metadata without cheating indicators source_codes.append({ @@ -77,7 +317,14 @@ def get_source_code(server_names: List[str]) -> Dict[str, Any]: "description": server_info.get("description"), "capabilities": server_info.get("capabilities", []) }, - "source_code": source_code + # Keep the original field for compatibility and add the + # dependency-aware bundle for complete source reasoning. + "source_code": source_code, + "entrypoint": resolved_path.name, + "source_files": source_files, + "source_bundle_complete": not any( + item["truncated"] for item in source_files + ), }) else: source_codes.append({ diff --git a/Detection/tests/test_source_code_analyzer_server.py b/Detection/tests/test_source_code_analyzer_server.py new file mode 100644 index 0000000..0532409 --- /dev/null +++ b/Detection/tests/test_source_code_analyzer_server.py @@ -0,0 +1,197 @@ +"""Tests for dependency-aware source retrieval.""" + +import inspect +from importlib.machinery import EXTENSION_SUFFIXES + +from context_providers import source_code_analyzer_server as analyzer +from context_providers.source_code_analyzer_server import collect_local_source_files + + +def test_collects_local_imported_implementation_without_unrelated_files(tmp_path): + function_dir = tmp_path / "function" + function_dir.mkdir() + entrypoint = tmp_path / "server.py" + entrypoint.write_text( + "from function.core import run\n\ndef tool():\n return run()\n", + encoding="utf-8", + ) + (function_dir / "core.py").write_text( + "def run():\n return 'hidden implementation'\n", + encoding="utf-8", + ) + (tmp_path / "unrelated.py").write_text("SECRET = True\n", encoding="utf-8") + + files = collect_local_source_files(entrypoint) + by_path = {item["path"]: item for item in files} + assert set(by_path) == {"server.py", "function/core.py"} + assert "hidden implementation" in by_path["function/core.py"]["source_code"] + assert all(str(tmp_path) not in item["path"] for item in files) + + +def test_relative_imports_remain_inside_server_root(tmp_path): + package = tmp_path / "package" + package.mkdir() + entrypoint = package / "entry.py" + entrypoint.write_text("from .helper import value\n", encoding="utf-8") + (package / "helper.py").write_text("value = 1\n", encoding="utf-8") + + files = collect_local_source_files(entrypoint) + assert {item["path"] for item in files} == {"entry.py", "helper.py"} + + +def test_regular_package_precedes_same_named_module(tmp_path): + entrypoint = tmp_path / "server.py" + entrypoint.write_text("import dependency\n", encoding="utf-8") + (tmp_path / "dependency.py").write_text("KIND = 'module'\n", encoding="utf-8") + package = tmp_path / "dependency" + package.mkdir() + (package / "__init__.py").write_text("KIND = 'package'\n", encoding="utf-8") + + files = collect_local_source_files(entrypoint) + by_path = {item["path"]: item for item in files} + + assert set(by_path) == {"server.py", "dependency/__init__.py"} + assert "KIND = 'package'" in by_path["dependency/__init__.py"]["source_code"] + + +def test_namespace_package_does_not_shadow_same_named_module(tmp_path): + entrypoint = tmp_path / "server.py" + entrypoint.write_text("import dependency\n", encoding="utf-8") + (tmp_path / "dependency.py").write_text("KIND = 'module'\n", encoding="utf-8") + (tmp_path / "dependency").mkdir() + + files = collect_local_source_files(entrypoint) + + assert {item["path"] for item in files} == {"server.py", "dependency.py"} + + +def test_from_package_import_follows_initializer_and_submodule(tmp_path): + entrypoint = tmp_path / "server.py" + entrypoint.write_text("from dependency import implementation\n", encoding="utf-8") + package = tmp_path / "dependency" + package.mkdir() + (package / "__init__.py").write_text("PACKAGE = True\n", encoding="utf-8") + (package / "implementation.py").write_text("VALUE = 1\n", encoding="utf-8") + + files = collect_local_source_files(entrypoint) + + assert {item["path"] for item in files} == { + "server.py", + "dependency/__init__.py", + "dependency/implementation.py", + } + + +def test_builtin_module_precedes_same_named_local_source(tmp_path): + entrypoint = tmp_path / "server.py" + entrypoint.write_text("import sys\n", encoding="utf-8") + (tmp_path / "sys.py").write_text("LOCAL = True\n", encoding="utf-8") + + files = collect_local_source_files(entrypoint) + + assert {item["path"] for item in files} == {"server.py"} + + +def test_native_module_precedes_same_named_local_source(tmp_path): + entrypoint = tmp_path / "server.py" + entrypoint.write_text("import dependency\n", encoding="utf-8") + (tmp_path / "dependency.py").write_text("LOCAL = True\n", encoding="utf-8") + (tmp_path / f"dependency{EXTENSION_SUFFIXES[0]}").write_bytes(b"") + + files = collect_local_source_files(entrypoint) + + assert {item["path"] for item in files} == {"server.py"} + + +def test_source_discovery_never_executes_imported_module(tmp_path): + entrypoint = tmp_path / "server.py" + entrypoint.write_text("import dependency\n", encoding="utf-8") + marker = tmp_path / "executed" + (tmp_path / "dependency.py").write_text( + f"from pathlib import Path\nPath({str(marker)!r}).touch()\n", + encoding="utf-8", + ) + + files = collect_local_source_files(entrypoint) + + assert {item["path"] for item in files} == {"server.py", "dependency.py"} + assert not marker.exists() + + +def test_symlinked_import_is_not_returned_as_local_source(tmp_path): + entrypoint = tmp_path / "server.py" + entrypoint.write_text("import linked\n", encoding="utf-8") + implementation = tmp_path / "implementation.py" + implementation.write_text("VALUE = 1\n", encoding="utf-8") + (tmp_path / "linked.py").symlink_to(implementation) + + files = collect_local_source_files(entrypoint) + + assert {item["path"] for item in files} == {"server.py"} + + +def test_registry_entrypoint_cannot_escape_provider_root(tmp_path, monkeypatch): + provider_root = tmp_path / "provider" + provider_root.mkdir() + secret = tmp_path / "secret.py" + secret.write_text("SECRET = 'must not be returned'\n", encoding="utf-8") + monkeypatch.setattr(analyzer, "__file__", str(provider_root / "provider.py")) + monkeypatch.setattr( + analyzer, + "source_registry", + {"mcp_servers": [{"name": "escaped", "path": "../secret.py"}]}, + ) + + result = analyzer.get_source_code(["escaped"]) + + assert result == { + "source_codes": [{"server_name": "escaped", "status": "file_not_found"}], + "total_retrieved": 0, + } + + +def test_benchmark_public_contract_is_unchanged(tmp_path, monkeypatch): + provider_root = tmp_path / "provider" + provider_root.mkdir() + entrypoint = provider_root / "server.py" + entrypoint.write_text("import dependency\n", encoding="utf-8") + (provider_root / "dependency.py").write_text("VALUE = 1\n", encoding="utf-8") + monkeypatch.setattr(analyzer, "__file__", str(provider_root / "provider.py")) + monkeypatch.setattr( + analyzer, + "source_registry", + { + "mcp_servers": [ + { + "name": "example", + "path": "server.py", + "category": "test", + "description": "example server", + "capabilities": ["read"], + } + ] + }, + ) + + result = analyzer.get_source_code(["example"]) + row = result["source_codes"][0] + + assert set(result) == {"source_codes", "total_retrieved"} + assert set(row) == { + "server_name", + "status", + "metadata", + "source_code", + "entrypoint", + "source_files", + "source_bundle_complete", + } + assert all( + set(source_file) == {"path", "source_code", "truncated"} + for source_file in row["source_files"] + ) + assert list(inspect.signature(collect_local_source_files).parameters) == [ + "entrypoint" + ] + assert analyzer.MAX_LOCAL_SOURCE_FILES == 24 + assert analyzer.MAX_LOCAL_SOURCE_CHARS == 240_000