23 | 24 |
25 | 26 | _mcp-scan in proxy mode._ 27 | 28 |
23 | Welcome to the Invariant MCP-scan Server!
56 |Use the API to interact with the server.
57 |Check the documentation for more information.
58 |Documentation: https://explorer.invariantlabs.ai/docs/mcp-scan
59 | """, 60 | media_type="text/html", 61 | status_code=200, 62 | ) 63 | 64 | async def on_startup(self): 65 | """Startup event for the FastAPI app.""" 66 | rich.print("[bold green]MCP-scan server started (http://localhost:" + str(self.port) + ")[/bold green]") 67 | 68 | # setup activity logger 69 | setup_activity_logger(self.app, pretty=self.pretty) 70 | 71 | from .routes.policies import load_guardrails_config_file 72 | 73 | await load_guardrails_config_file(self.config_file_path) 74 | 75 | async def life_span(self, app: FastAPI): 76 | """Lifespan event for the FastAPI app.""" 77 | await self.on_startup() 78 | 79 | yield 80 | 81 | if callable(self.on_exit): 82 | if inspect.iscoroutinefunction(self.on_exit): 83 | await self.on_exit() 84 | else: 85 | self.on_exit() 86 | 87 | def run(self): 88 | """Run the MCP scan server.""" 89 | uvicorn.run(self.app, host="0.0.0.0", port=self.port, log_level=self.log_level) 90 | -------------------------------------------------------------------------------- /src/mcp_scan_server/session_store.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass 2 | from datetime import datetime 3 | from enum import Enum 4 | from typing import Any 5 | 6 | 7 | class MergeNodeTypes(Enum): 8 | SELF = "self" 9 | OTHER = "other" 10 | SELF_TO = "self_to" 11 | OTHER_TO = "other_to" 12 | 13 | 14 | @dataclass(frozen=True) 15 | class MergeInstruction: 16 | node_type: MergeNodeTypes 17 | index: int 18 | 19 | 20 | @dataclass(frozen=True) 21 | class SessionNode: 22 | """ 23 | Represents a single event in a session. 24 | """ 25 | 26 | timestamp: datetime 27 | message: dict[str, Any] 28 | session_id: str 29 | server_name: str 30 | original_session_index: int 31 | 32 | def __hash__(self) -> int: 33 | """Assume uniqueness by session_id, index in session and time of event.""" 34 | return hash((self.session_id, self.original_session_index, self.timestamp)) 35 | 36 | def __lt__(self, other: "SessionNode") -> bool: 37 | """Sort by timestamp.""" 38 | return self.timestamp < other.timestamp 39 | 40 | 41 | class Session: 42 | """ 43 | Represents a sequence of SessionNodes, sorted by timestamp. 44 | """ 45 | 46 | def __init__( 47 | self, 48 | nodes: list[SessionNode] | None = None, 49 | ): 50 | self.nodes: list[SessionNode] = nodes or [] 51 | self.last_analysis_index: int = -1 52 | 53 | def _build_stack(self, other: "Session") -> list[MergeInstruction]: 54 | """ 55 | Construct a build stack of the nodes that make up the merged session. 56 | 57 | The build stack is a list of instructions for how to construct the merged session. 58 | 59 | We iterate over the nodes of the two sessions in reverse order, essentially performing a 60 | heap merge. From this, we construct a set of instructions on how to construct the merged session. 61 | We could have built the merged sessions directly, but then we couldn't iterate in reverse order 62 | and thus not be able to exit early when we have found an already inserted node. 63 | """ 64 | build_stack: list[MergeInstruction] = [] 65 | ptr_self, ptr_other = len(self.nodes) - 1, len(other.nodes) - 1 66 | early_exit = False 67 | 68 | while ptr_self >= 0 and ptr_other >= 0: 69 | # Exit early if we have found an already inserted node. 70 | if self.nodes[ptr_self] == other.nodes[ptr_other]: 71 | build_stack.append(MergeInstruction(MergeNodeTypes.SELF_TO, ptr_self)) 72 | early_exit = True 73 | break 74 | 75 | # Insert other node if it comes after the self node. 76 | elif self.nodes[ptr_self] < other.nodes[ptr_other]: 77 | build_stack.append(MergeInstruction(MergeNodeTypes.OTHER, ptr_other)) 78 | ptr_other -= 1 79 | 80 | # Insert self node if it comes after the other node. 81 | else: 82 | build_stack.append(MergeInstruction(MergeNodeTypes.SELF, ptr_self)) 83 | ptr_self -= 1 84 | 85 | # Handle remaining nodes in either self or other. 86 | # If we do not exit early, we should have some nodes 87 | # left in either self or other but not both. 88 | if not early_exit: 89 | if ptr_self >= 0: 90 | build_stack.append(MergeInstruction(MergeNodeTypes.SELF_TO, ptr_self)) 91 | elif ptr_other >= 0: 92 | build_stack.append(MergeInstruction(MergeNodeTypes.OTHER_TO, ptr_other)) 93 | 94 | return build_stack 95 | 96 | def _build_merged_nodes(self, build_stack: list[MergeInstruction], other: "Session") -> list[SessionNode]: 97 | """ 98 | Build the merged nodes from the build stack. 99 | 100 | The build stack is a stack of instructions for how to construct the merged session. 101 | The node_type is either "self", "other", "self_to" or "other_to". 102 | The index is the index of the node in the session. 103 | The "self_to" and "other_to" tuples are used to indicate that all nodes up to and including the index should be inserted from the respective session. 104 | The "self" and "other" tuples are used to indicate that the node at the index should be inserted from the respective session. 105 | """ 106 | merged_nodes = [] 107 | for merged_index, instruction in enumerate(reversed(build_stack)): 108 | if instruction.node_type == MergeNodeTypes.SELF: 109 | merged_nodes.append(self.nodes[instruction.index]) 110 | elif instruction.node_type == MergeNodeTypes.OTHER: 111 | merged_nodes.append(other.nodes[instruction.index]) 112 | # Update the last analysis index to the index of the last node from other. 113 | self.last_analysis_index = min(self.last_analysis_index, merged_index) 114 | elif instruction.node_type == MergeNodeTypes.SELF_TO: 115 | merged_nodes.extend(self.nodes[: instruction.index + 1]) 116 | elif instruction.node_type == MergeNodeTypes.OTHER_TO: 117 | merged_nodes.extend(other.nodes[: instruction.index + 1]) 118 | # Reset the index because we have inserted nodes from other that 119 | # were before the nodes from self. 120 | self.last_analysis_index = -1 121 | return merged_nodes 122 | 123 | def merge(self, other: "Session") -> None: 124 | """ 125 | Merge two session objects into a joint session. 126 | This assumes the precondition that both sessions are sorted and that duplicate nodes cannot exist 127 | (refer to the __hash__ method for session nodes). 128 | The postcondition is that the merged session is sorted, has no duplicates, and is the union of the two sessions. 129 | 130 | The algorithm proceeds in two steps: 131 | 1. Construct a build stack of the nodes that make up the merged session. 132 | 2. Iterate over the build stack in reverse order and build the merged nodes. 133 | 134 | When constructing the build stack, we can exit early if we have found an already inserted node 135 | using the precondition, since it implies that part of this trace has already been inserted -- 136 | specifically the part before the equal nodes. 137 | """ 138 | build_stack = self._build_stack(other) 139 | merged_nodes = self._build_merged_nodes(build_stack, other) 140 | self.nodes = merged_nodes 141 | 142 | def get_sorted_nodes(self) -> list[SessionNode]: 143 | return self.nodes 144 | 145 | def __repr__(self): 146 | return f"Session(nodes={self.get_sorted_nodes()})" 147 | 148 | 149 | class SessionStore: 150 | """ 151 | Stores sessions by client_name. 152 | """ 153 | 154 | def __init__(self): 155 | self.sessions: dict[str, Session] = {} 156 | 157 | def _default_session(self) -> Session: 158 | return Session() 159 | 160 | def __str__(self): 161 | return f"SessionStore(sessions={self.sessions})" 162 | 163 | def __getitem__(self, client_name: str) -> Session: 164 | if client_name not in self.sessions: 165 | self.sessions[client_name] = self._default_session() 166 | return self.sessions[client_name] 167 | 168 | def __setitem__(self, client_name: str, session: Session) -> None: 169 | self.sessions[client_name] = session 170 | 171 | def __repr__(self): 172 | return self.__str__() 173 | 174 | def fetch_and_merge(self, client_name: str, other: Session) -> Session: 175 | """ 176 | Fetch the session for the given client_name and merge it with the other session, returning the merged session. 177 | """ 178 | session = self[client_name] 179 | session.merge(other) 180 | return session 181 | 182 | 183 | async def to_session(messages: list[dict[str, Any]], server_name: str, session_id: str) -> Session: 184 | """ 185 | Convert a list of messages to a session. 186 | """ 187 | session_nodes: list[SessionNode] = [] 188 | for i, message in enumerate(messages): 189 | timestamp = datetime.fromisoformat(message["timestamp"]) 190 | session_nodes.append( 191 | SessionNode( 192 | server_name=server_name, 193 | message=message, 194 | original_session_index=i, 195 | session_id=session_id, 196 | timestamp=timestamp, 197 | ) 198 | ) 199 | 200 | return Session(nodes=session_nodes) 201 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | """Test package for mcp-scan.""" 2 | -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- 1 | """Global pytest fixtures for mcp-scan tests.""" 2 | 3 | import pytest 4 | 5 | from mcp_scan.utils import TempFile 6 | 7 | 8 | @pytest.fixture 9 | def claudestyle_config(): 10 | """Sample Claude-style MCP config.""" 11 | return """{ 12 | "mcpServers": { 13 | "claude": { 14 | "command": "mcp", 15 | "args": ["--server", "http://localhost:8000"], 16 | } 17 | } 18 | }""" 19 | 20 | 21 | @pytest.fixture 22 | def claudestyle_config_file(claudestyle_config): 23 | with TempFile(mode="w") as temp_file: 24 | temp_file.write(claudestyle_config) 25 | temp_file.flush() 26 | yield temp_file.name 27 | 28 | 29 | @pytest.fixture 30 | def vscode_mcp_config(): 31 | """Sample VSCode MCP config with inputs.""" 32 | return """{ 33 | // Inputs are prompted on first server start, then stored securely by VS Code. 34 | "inputs": [ 35 | { 36 | "type": "promptString", 37 | "id": "perplexity-key", 38 | "description": "Perplexity API Key", 39 | "password": true 40 | } 41 | ], 42 | "servers": { 43 | // https://github.com/ppl-ai/modelcontextprotocol/ 44 | "Perplexity": { 45 | "type": "stdio", 46 | "command": "npx", 47 | "args": ["-y", "@modelcontextprotocol/server-perplexity-ask"], 48 | "env": { 49 | "PERPLEXITY_API_KEY": "ASDF" 50 | } 51 | } 52 | } 53 | } 54 | """ 55 | 56 | 57 | @pytest.fixture 58 | def vscode_mcp_config_file(vscode_mcp_config): 59 | with TempFile(mode="w") as temp_file: 60 | temp_file.write(vscode_mcp_config) 61 | temp_file.flush() 62 | yield temp_file.name 63 | 64 | 65 | @pytest.fixture 66 | def vscode_config(): 67 | """Sample VSCode settings.json with MCP config.""" 68 | return """// settings.json 69 | { 70 | "mcp": { 71 | "servers": { 72 | "my-mcp-server": { 73 | "type": "stdio", 74 | "command": "my-command", 75 | "args": [] 76 | } 77 | } 78 | } 79 | }""" 80 | 81 | 82 | @pytest.fixture 83 | def vscode_config_file(vscode_config): 84 | with TempFile(mode="w") as temp_file: 85 | temp_file.write(vscode_config) 86 | temp_file.flush() 87 | yield temp_file.name 88 | 89 | 90 | @pytest.fixture 91 | def toy_server_add(): 92 | """Example toy server from the mcp docs.""" 93 | return """ 94 | from mcp.server.fastmcp import FastMCP 95 | 96 | # Create an MCP server 97 | mcp = FastMCP("Demo") 98 | 99 | # Add an addition tool 100 | @mcp.tool() 101 | def add(a: int, b: int) -> int: 102 | return a + b 103 | """ 104 | 105 | 106 | @pytest.fixture 107 | def toy_server_add_file(toy_server_add): 108 | with TempFile(mode="w", suffix=".py") as temp_file: 109 | temp_file.write(toy_server_add) 110 | temp_file.flush() 111 | yield temp_file.name.replace("\\", "/") 112 | 113 | # filename = "tmp_toy_server_" + str(uuid.uuid4()) + ".py" 114 | # # create the file 115 | # with open(filename, "w") as temp_file: 116 | # temp_file.write(toy_server_add) 117 | # temp_file.flush() 118 | # temp_file.seek(0) 119 | 120 | # # run tests 121 | # yield filename.replace("\\", "/") 122 | # # cleanup 123 | # import os 124 | 125 | # os.remove(filename) 126 | 127 | 128 | @pytest.fixture 129 | def toy_server_add_config(toy_server_add_file): 130 | return f""" 131 | {{ 132 | "mcpServers": {{ 133 | "toy": {{ 134 | "command": "mcp", 135 | "args": ["run", "{toy_server_add_file}"] 136 | }} 137 | }} 138 | }} 139 | """ 140 | 141 | 142 | @pytest.fixture 143 | def toy_server_add_config_file(toy_server_add_config): 144 | with TempFile(mode="w", suffix=".json") as temp_file: 145 | temp_file.write(toy_server_add_config) 146 | temp_file.flush() 147 | yield temp_file.name.replace("\\", "/") 148 | 149 | # filename = "tmp_config_" + str(uuid.uuid4()) + ".json" 150 | 151 | # # create the file 152 | # with open(filename, "w") as temp_file: 153 | # temp_file.write(toy_server_add_config) 154 | # temp_file.flush() 155 | # temp_file.seek(0) 156 | 157 | # # run tests 158 | # yield filename.replace("\\", "/") 159 | 160 | # # cleanup 161 | # import os 162 | 163 | # os.remove(filename) 164 | 165 | 166 | @pytest.fixture 167 | def math_server_config_path(): 168 | return "tests/mcp_servers/mcp_config.json" 169 | -------------------------------------------------------------------------------- /tests/e2e/__init__.py: -------------------------------------------------------------------------------- 1 | """End-to-end tests package for mcp-scan.""" 2 | -------------------------------------------------------------------------------- /tests/e2e/test_full_proxy_flow.py: -------------------------------------------------------------------------------- 1 | """End-to-end tests for complete MCP scanning workflow.""" 2 | 3 | import asyncio 4 | import os 5 | import subprocess 6 | import time 7 | 8 | import dotenv 9 | import pytest 10 | from mcp import ClientSession 11 | 12 | from mcp_scan.mcp_client import get_client, scan_mcp_config_file 13 | 14 | 15 | # Helper function to safely decode subprocess output 16 | def safe_decode(bytes_output, encoding="utf-8", errors="replace"): 17 | """Safely decode subprocess output, handling potential Unicode errors""" 18 | if bytes_output is None: 19 | return "" 20 | try: 21 | return bytes_output.decode(encoding) 22 | except UnicodeDecodeError: 23 | # Fall back to a more lenient error handler 24 | return bytes_output.decode(encoding, errors=errors) 25 | 26 | 27 | async def run_toy_server_client(config): 28 | async with get_client(config) as (read, write): 29 | async with ClientSession(read, write) as session: 30 | print("[Client] Initializing connection") 31 | await session.initialize() 32 | print("[Client] Listing tools") 33 | tools = await session.list_tools() 34 | print("[Client] Tools: ", tools.tools) 35 | 36 | print("[Client] Calling tool add") 37 | result = await session.call_tool("add", arguments={"a": 1, "b": 2}) 38 | result = result.content[0].text 39 | print("[Client] Result: ", result) 40 | 41 | return { 42 | "result": result, 43 | "tools": tools.tools, 44 | } 45 | return result 46 | 47 | 48 | async def ensure_config_file_contains_gateway(config_file, timeout=3): 49 | s = time.time() 50 | content = "" 51 | 52 | while True: 53 | with open(config_file) as f: 54 | content = f.read() 55 | if "invariant-gateway" in content: 56 | return True 57 | await asyncio.sleep(0.1) 58 | if time.time() - s > timeout: 59 | return False 60 | 61 | 62 | class TestFullProxyFlow: 63 | """Test cases for end-to-end scanning workflows.""" 64 | 65 | PORT = 9129 66 | 67 | @pytest.mark.asyncio 68 | @pytest.mark.parametrize("pretty", ["oneline", "full", "compact"]) 69 | # skip on windows 70 | @pytest.mark.skipif( 71 | os.name == "nt", 72 | reason="Skipping test on Windows due to subprocess handling issues", 73 | ) 74 | async def test_basic(self, toy_server_add_config_file, pretty): 75 | # if available, check for 'lsof' and make sure the port is not in use 76 | try: 77 | subprocess.check_output(["lsof", "-i", f":{self.PORT}"]) 78 | print(f"Port {self.PORT} is in use") 79 | return 80 | except subprocess.CalledProcessError: 81 | pass 82 | except FileNotFoundError: 83 | print("lsof not found, skipping port check") 84 | 85 | args = dotenv.dotenv_values(".env") 86 | gateway_dir = args.get("INVARIANT_GATEWAY_DIR", None) 87 | command = [ 88 | "uv", 89 | "run", 90 | "-m", 91 | "src.mcp_scan.run", 92 | "proxy", 93 | # ensure we are using the right ports 94 | "--mcp-scan-server-port", 95 | str(self.PORT), 96 | "--port", 97 | str(self.PORT), 98 | "--pretty", 99 | pretty, 100 | ] 101 | if gateway_dir is not None: 102 | command.extend(["--gateway-dir", gateway_dir]) 103 | command.append(toy_server_add_config_file) 104 | 105 | # start process in background 106 | env = {**os.environ, "COLUMNS": "256"} 107 | # Ensure proper handling of Unicode on Windows 108 | if os.name == "nt": # Windows 109 | # Explicitly set encoding for console on Windows 110 | env["PYTHONIOENCODING"] = "utf-8" 111 | 112 | process = subprocess.Popen( 113 | command, 114 | stdout=subprocess.PIPE, 115 | stderr=subprocess.PIPE, 116 | env=env, 117 | universal_newlines=False, # Binary mode for better Unicode handling 118 | ) 119 | 120 | # wait for gateway to be installed 121 | if not (await ensure_config_file_contains_gateway(toy_server_add_config_file)): 122 | # if process is not running, raise an error 123 | if process.poll() is not None: 124 | # process has terminated 125 | stdout, stderr = process.communicate() 126 | print(safe_decode(stdout)) 127 | print(safe_decode(stderr)) 128 | raise AssertionError("process terminated before gateway was installed") 129 | 130 | # print out toy_server_add_config_file 131 | with open(toy_server_add_config_file) as f: 132 | # assert that 'invariant-gateway' is in the file 133 | content = f.read() 134 | 135 | if "invariant-gateway" not in content: 136 | # terminate the process and get output 137 | process.terminate() 138 | process.wait() 139 | 140 | # get output 141 | stdout, stderr = process.communicate() 142 | print(safe_decode(stdout)) 143 | print(safe_decode(stderr)) 144 | 145 | assert "invariant-gateway" in content, ( 146 | "invariant-gateway wrapper was not found in the config file: " 147 | + content 148 | + "\nProcess output: " 149 | + safe_decode(stdout) 150 | + "\nError output: " 151 | + safe_decode(stderr) 152 | ) 153 | 154 | with open(toy_server_add_config_file) as f: 155 | # assert that 'invariant-gateway' is in the file 156 | content = f.read() 157 | print(content) 158 | 159 | # start client 160 | config = await scan_mcp_config_file(toy_server_add_config_file) 161 | servers = list(config.mcpServers.values()) 162 | assert len(servers) == 1 163 | server = servers[0] 164 | client_program = run_toy_server_client(server) 165 | 166 | # wait for client to finish 167 | try: 168 | client_output = await asyncio.wait_for(client_program, timeout=20) 169 | except asyncio.TimeoutError as e: 170 | print("Client timed out") 171 | process.terminate() 172 | process.wait() 173 | stdout, stderr = process.communicate() 174 | print(safe_decode(stdout)) 175 | print(safe_decode(stderr)) 176 | raise AssertionError("timed out waiting for MCP server to respond") from e 177 | 178 | assert int(client_output["result"]) == 3 179 | 180 | # shut down server and collect output 181 | process.terminate() 182 | stdout, stderr = process.communicate() 183 | process.wait() 184 | 185 | # print full outputs 186 | stdout_text = safe_decode(stdout) 187 | stderr_text = safe_decode(stderr) 188 | print("stdout: ", stdout_text) 189 | print("stderr: ", stderr_text) 190 | 191 | # basic checks for the log 192 | assert "used toy to tools/list" in stdout_text, "basic activity log statement not found" 193 | assert "call_1" in stdout_text, "call_1 not found in log" 194 | 195 | assert "call_2" in stdout_text, "call_2 not found in log" 196 | assert "to add" in stdout_text, "call to 'add' not found in log" 197 | 198 | # assert there is no 'address is already in use' error 199 | assert "address already in use" not in stderr_text, ( 200 | "mcp-scan proxy failed to start because the testing port " 201 | + str(self.PORT) 202 | + " is already in use. Please make sure to stop any other mcp-scan proxy server running on this port." 203 | ) 204 | -------------------------------------------------------------------------------- /tests/e2e/test_full_scan_flow.py: -------------------------------------------------------------------------------- 1 | """End-to-end tests for complete MCP scanning workflow.""" 2 | 3 | import json 4 | import subprocess 5 | 6 | import pytest 7 | from pytest_lazy_fixtures import lf 8 | 9 | from mcp_scan.utils import TempFile 10 | 11 | 12 | class TestFullScanFlow: 13 | """Test cases for end-to-end scanning workflows.""" 14 | 15 | @pytest.mark.parametrize( 16 | "sample_config_file", [lf("claudestyle_config_file"), lf("vscode_mcp_config_file"), lf("vscode_config_file")] 17 | ) 18 | def test_basic(self, sample_config_file): 19 | """Test a basic complete scan workflow from CLI to results. This does not mean that the results are correct or the servers can be run.""" 20 | # Run mcp-scan with JSON output mode 21 | result = subprocess.run( 22 | ["uv", "run", "-m", "src.mcp_scan.run", "scan", "--json", sample_config_file], 23 | capture_output=True, 24 | text=True, 25 | ) 26 | 27 | # Check that the command executed successfully 28 | assert result.returncode == 0, f"Command failed with error: {result.stderr}" 29 | 30 | print(result.stdout) 31 | print(result.stderr) 32 | 33 | # Try to parse the output as JSON 34 | try: 35 | output = json.loads(result.stdout) 36 | assert sample_config_file in output 37 | except json.JSONDecodeError: 38 | print(result.stdout) 39 | pytest.fail("Failed to parse JSON output") 40 | 41 | @pytest.mark.parametrize( 42 | "path, server_names", 43 | [ 44 | ("tests/mcp_servers/configs_files/weather_config.json", ["Weather"]), 45 | ("tests/mcp_servers/configs_files/math_config.json", ["Math"]), 46 | ("tests/mcp_servers/configs_files/all_config.json", ["Weather", "Math"]), 47 | ], 48 | ) 49 | def test_scan(self, path, server_names): 50 | path = "tests/mcp_servers/configs_files/all_config.json" 51 | result = subprocess.run( 52 | ["uv", "run", "-m", "src.mcp_scan.run", "scan", "--json", path], 53 | capture_output=True, 54 | text=True, 55 | ) 56 | assert result.returncode == 0, f"Command failed with error: {result.stderr}" 57 | output = json.loads(result.stdout) 58 | results: dict[str, dict] = {} 59 | for server in output[path]["servers"]: 60 | results[server["name"]] = server["result"] 61 | server["signature"]["metadata"]["serverInfo"]["version"] = ( 62 | "mcp_version" # swap actual version with placeholder 63 | ) 64 | 65 | with open(f"tests/mcp_servers/signatures/{server['name'].lower()}_server_signature.json") as f: 66 | assert server["signature"] == json.load(f), f"Signature mismatch for {server['name']} server" 67 | 68 | expected_results = { 69 | "Weather": [ 70 | { 71 | "changed": None, 72 | "messages": [], 73 | "status": None, 74 | "verified": True, 75 | "whitelisted": None, 76 | } 77 | ], 78 | "Math": [ 79 | { 80 | "changed": None, 81 | "messages": [], 82 | "status": None, 83 | "verified": True, 84 | "whitelisted": None, 85 | } 86 | ] 87 | * 4, 88 | } 89 | for server_name in server_names: 90 | assert results[server_name] == expected_results[server_name], f"Results mismatch for {server_name} server" 91 | 92 | def test_inspect(self): 93 | path = "tests/mcp_servers/configs_files/all_config.json" 94 | result = subprocess.run( 95 | ["uv", "run", "-m", "src.mcp_scan.run", "inspect", "--json", path], 96 | capture_output=True, 97 | text=True, 98 | ) 99 | assert result.returncode == 0, f"Command failed with error: {result.stderr}" 100 | output = json.loads(result.stdout) 101 | 102 | assert path in output 103 | for server in output[path]["servers"]: 104 | server["signature"]["metadata"]["serverInfo"]["version"] = ( 105 | "mcp_version" # swap actual version with placeholder 106 | ) 107 | 108 | with open(f"tests/mcp_servers/signatures/{server['name'].lower()}_server_signature.json") as f: 109 | assert server["signature"] == json.load(f), f"Signature mismatch for {server['name']} server" 110 | 111 | @pytest.fixture 112 | def vscode_settings_no_mcp_file(self): 113 | settings = { 114 | "[javascript]": {}, 115 | "github.copilot.advanced": {}, 116 | "github.copilot.chat.agent.thinkingTool": {}, 117 | "github.copilot.chat.codesearch.enabled": {}, 118 | "github.copilot.chat.languageContext.typescript.enabled": {}, 119 | "github.copilot.chat.welcomeMessage": {}, 120 | "github.copilot.enable": {}, 121 | "github.copilot.preferredAccount": {}, 122 | "settingsSync.ignoredExtensions": {}, 123 | "tabnine.experimentalAutoImports": {}, 124 | "workbench.colorTheme": {}, 125 | "workbench.startupEditor": {}, 126 | } 127 | with TempFile(mode="w") as temp_file: 128 | json.dump(settings, temp_file) 129 | temp_file.flush() 130 | yield temp_file.name 131 | 132 | def test_vscode_settings_no_mcp(self, vscode_settings_no_mcp_file): 133 | """Test scanning VSCode settings with no MCP configurations.""" 134 | result = subprocess.run( 135 | ["uv", "run", "-m", "src.mcp_scan.run", "scan", "--json", vscode_settings_no_mcp_file], 136 | capture_output=True, 137 | text=True, 138 | ) 139 | 140 | # Check that the command executed successfully 141 | assert result.returncode == 0, f"Command failed with error: {result.stderr}" 142 | 143 | # Try to parse the output as JSON 144 | try: 145 | output = json.loads(result.stdout) 146 | assert vscode_settings_no_mcp_file in output 147 | except json.JSONDecodeError: 148 | pytest.fail("Failed to parse JSON output") 149 | -------------------------------------------------------------------------------- /tests/mcp_servers/configs_files/all_config.json: -------------------------------------------------------------------------------- 1 | { 2 | "mcpServers": { 3 | "Weather": { 4 | "command": "uv run python", 5 | "args": ["tests/mcp_servers/weather_server.py"] 6 | }, 7 | "Math": { 8 | "command": "uv run python", 9 | "args": ["tests/mcp_servers/math_server.py"] 10 | } 11 | } 12 | } 13 | -------------------------------------------------------------------------------- /tests/mcp_servers/configs_files/math_config.json: -------------------------------------------------------------------------------- 1 | { 2 | "mcpServers": { 3 | "Math": { 4 | "command": "uv run python", 5 | "args": ["tests/mcp_servers/math_server.py"] 6 | } 7 | } 8 | } 9 | -------------------------------------------------------------------------------- /tests/mcp_servers/configs_files/weather_config.json: -------------------------------------------------------------------------------- 1 | { 2 | "mcpServers": { 3 | "Weather": { 4 | "command": "uv run python", 5 | "args": ["tests/mcp_servers/weather_server.py"] 6 | } 7 | } 8 | } 9 | -------------------------------------------------------------------------------- /tests/mcp_servers/math_server.py: -------------------------------------------------------------------------------- 1 | from mcp.server.fastmcp import FastMCP 2 | 3 | # Create an MCP server 4 | mcp = FastMCP("Math") 5 | 6 | 7 | # Add an addition tool 8 | @mcp.tool() 9 | def add(a: int, b: int) -> int: 10 | """Add two numbers.""" 11 | return a + b 12 | 13 | 14 | # Add a subtraction tool 15 | @mcp.tool() 16 | def subtract(a: int, b: int) -> int: 17 | """Subtract two numbers.""" 18 | return a - b 19 | 20 | 21 | # Add a multiplication tool 22 | @mcp.tool() 23 | def multiply(a: int, b: int) -> int: 24 | """Multiply two numbers.""" 25 | return a * b 26 | 27 | 28 | # Add a division tool 29 | @mcp.tool() 30 | def divide(a: int, b: int) -> int: 31 | """Divide two numbers.""" 32 | if b == 0: 33 | raise ValueError("Cannot divide by zero") 34 | return a // b 35 | 36 | 37 | if __name__ == "__main__": 38 | mcp.run() 39 | -------------------------------------------------------------------------------- /tests/mcp_servers/signatures/math_server_signature.json: -------------------------------------------------------------------------------- 1 | { 2 | "metadata": { 3 | "meta": null, 4 | "protocolVersion": "2024-11-05", 5 | "capabilities": { 6 | "experimental": {}, 7 | "logging": null, 8 | "prompts": { 9 | "listChanged": false 10 | }, 11 | "resources": { 12 | "subscribe": false, 13 | "listChanged": false 14 | }, 15 | "tools": { 16 | "listChanged": false 17 | } 18 | }, 19 | "serverInfo": { 20 | "name": "Math", 21 | "version": "mcp_version" 22 | }, 23 | "instructions": null 24 | }, 25 | "prompts": [], 26 | "resources": [], 27 | "tools": [ 28 | { 29 | "name": "add", 30 | "description": "Add two numbers.", 31 | "inputSchema": { 32 | "properties": { 33 | "a": { 34 | "title": "A", 35 | "type": "integer" 36 | }, 37 | "b": { 38 | "title": "B", 39 | "type": "integer" 40 | } 41 | }, 42 | "required": [ 43 | "a", 44 | "b" 45 | ], 46 | "title": "addArguments", 47 | "type": "object" 48 | }, 49 | "annotations": null 50 | }, 51 | { 52 | "name": "subtract", 53 | "description": "Subtract two numbers.", 54 | "inputSchema": { 55 | "properties": { 56 | "a": { 57 | "title": "A", 58 | "type": "integer" 59 | }, 60 | "b": { 61 | "title": "B", 62 | "type": "integer" 63 | } 64 | }, 65 | "required": [ 66 | "a", 67 | "b" 68 | ], 69 | "title": "subtractArguments", 70 | "type": "object" 71 | }, 72 | "annotations": null 73 | }, 74 | { 75 | "name": "multiply", 76 | "description": "Multiply two numbers.", 77 | "inputSchema": { 78 | "properties": { 79 | "a": { 80 | "title": "A", 81 | "type": "integer" 82 | }, 83 | "b": { 84 | "title": "B", 85 | "type": "integer" 86 | } 87 | }, 88 | "required": [ 89 | "a", 90 | "b" 91 | ], 92 | "title": "multiplyArguments", 93 | "type": "object" 94 | }, 95 | "annotations": null 96 | }, 97 | { 98 | "name": "divide", 99 | "description": "Divide two numbers.", 100 | "inputSchema": { 101 | "properties": { 102 | "a": { 103 | "title": "A", 104 | "type": "integer" 105 | }, 106 | "b": { 107 | "title": "B", 108 | "type": "integer" 109 | } 110 | }, 111 | "required": [ 112 | "a", 113 | "b" 114 | ], 115 | "title": "divideArguments", 116 | "type": "object" 117 | }, 118 | "annotations": null 119 | } 120 | ] 121 | } 122 | -------------------------------------------------------------------------------- /tests/mcp_servers/signatures/weather_server_signature.json: -------------------------------------------------------------------------------- 1 | { 2 | "metadata": { 3 | "meta": null, 4 | "protocolVersion": "2024-11-05", 5 | "capabilities": { 6 | "experimental": {}, 7 | "logging": null, 8 | "prompts": { 9 | "listChanged": false 10 | }, 11 | "resources": { 12 | "subscribe": false, 13 | "listChanged": false 14 | }, 15 | "tools": { 16 | "listChanged": false 17 | } 18 | }, 19 | "serverInfo": { 20 | "name": "Weather", 21 | "version": "mcp_version" 22 | }, 23 | "instructions": null 24 | }, 25 | "prompts": [], 26 | "resources": [], 27 | "tools": [ 28 | { 29 | "name": "weather", 30 | "description": "Get current weather for a location.", 31 | "inputSchema": { 32 | "properties": { 33 | "location": { 34 | "title": "Location", 35 | "type": "string" 36 | } 37 | }, 38 | "required": [ 39 | "location" 40 | ], 41 | "title": "weatherArguments", 42 | "type": "object" 43 | }, 44 | "annotations": null 45 | } 46 | ] 47 | } 48 | -------------------------------------------------------------------------------- /tests/mcp_servers/weather_server.py: -------------------------------------------------------------------------------- 1 | import random 2 | 3 | from mcp.server.fastmcp import FastMCP 4 | 5 | # Create an MCP server 6 | mcp = FastMCP("Weather") 7 | 8 | 9 | @mcp.tool() 10 | def weather(location: str) -> str: 11 | """Get current weather for a location.""" 12 | return random.choice(["Sunny", "Rainy", "Cloudy", "Snowy", "Windy"]) 13 | 14 | 15 | if __name__ == "__main__": 16 | mcp.run() 17 | -------------------------------------------------------------------------------- /tests/test_configs.json: -------------------------------------------------------------------------------- 1 | { 2 | "mcpServers": { 3 | "Random Facts MCP Server": { 4 | "command": "uv", 5 | "args": [ 6 | "run", 7 | "--with", 8 | "mcp[cli]", 9 | "mcp", 10 | "run", 11 | "/Users/marcomilanta/Documents/invariant/mcp-injection-experiments/whatsapp-takeover.py" 12 | ], 13 | "type": "stdio", 14 | "env": {} 15 | }, 16 | "WhatsApp Server": { 17 | "command": "uv", 18 | "args": [ 19 | "run", 20 | "--with", 21 | "mcp[cli]", 22 | "--with", 23 | "requests", 24 | "mcp", 25 | "run", 26 | "/Users/marcomilanta/Documents/invariant/mcp-injection-experiments/whatsapp.py" 27 | ], 28 | "type": "stdio", 29 | "env": {} 30 | } 31 | } 32 | } 33 | -------------------------------------------------------------------------------- /tests/unit/__init__.py: -------------------------------------------------------------------------------- 1 | """Unit tests package for mcp-scan.""" 2 | -------------------------------------------------------------------------------- /tests/unit/test_gateway.py: -------------------------------------------------------------------------------- 1 | import json 2 | 3 | import pyjson5 4 | import pytest 5 | from pytest_lazy_fixtures import lf 6 | 7 | from mcp_scan.gateway import MCPGatewayConfig, MCPGatewayInstaller, is_invariant_installed 8 | from mcp_scan.mcp_client import scan_mcp_config_file 9 | from mcp_scan.models import StdioServer 10 | 11 | 12 | @pytest.mark.asyncio 13 | @pytest.mark.parametrize("sample_config_file", [lf("claudestyle_config_file")]) 14 | async def test_install_gateway(sample_config_file): 15 | with open(sample_config_file) as f: 16 | config_dict = pyjson5.load(f) 17 | installer = MCPGatewayInstaller(paths=[sample_config_file]) 18 | for server in (await scan_mcp_config_file(sample_config_file)).get_servers().values(): 19 | if isinstance(server, StdioServer): 20 | assert not is_invariant_installed(server), "Invariant should not be installed" 21 | await installer.install( 22 | gateway_config=MCPGatewayConfig(project_name="test", push_explorer=True, api_key="my-very-secret-api-key"), 23 | verbose=True, 24 | ) 25 | 26 | # try to load the config 27 | with open(sample_config_file) as f: 28 | pyjson5.load(f) 29 | 30 | for server in (await scan_mcp_config_file(sample_config_file)).get_servers().values(): 31 | if isinstance(server, StdioServer): 32 | assert is_invariant_installed(server), "Invariant should be installed" 33 | 34 | await installer.uninstall(verbose=True) 35 | 36 | for server in (await scan_mcp_config_file(sample_config_file)).get_servers().values(): 37 | if isinstance(server, StdioServer): 38 | assert not is_invariant_installed(server), "Invariant should be uninstalled" 39 | 40 | with open(sample_config_file) as f: 41 | config_dict_uninstalled = pyjson5.load(f) 42 | 43 | # check for mcpServers.