feat: add query for labels

Author: mmilanta

src/mcp_scan/MCPScanner.py

@@ -9,7 +9,7 @@
 
 from .mcp_client import check_server_with_timeout, scan_mcp_config_file
 from .StorageFile import StorageFile
-from .verify_api import verify_scan_path
+from .verify_api import verify_scan_path_and_labels
 
 # Set up logger for this module
 logger = logging.getLogger(__name__)
@@ -190,7 +190,9 @@ async def scan_path(self, path: str, inspect_only: bool = False) -> ScanPathResu
             logger.debug("Scanning server %d/%d: %s", i + 1, len(path_result.servers), server.name)
             path_result.servers[i] = await self.scan_server(server, inspect_only)
         logger.debug("Verifying server path: %s", path)
-        path_result = await verify_scan_path(path_result, base_url=self.base_url, run_locally=self.local_only)
+        path_result = await verify_scan_path_and_labels(
+            path_result, base_url=self.base_url, run_locally=self.local_only
+        )
         await self.emit("path_scanned", path_result)
         return path_result
 

src/mcp_scan/models.py

@@ -3,7 +3,7 @@
 from itertools import chain
 from typing import Any, Literal, TypeAlias
 
-from mcp.types import InitializeResult, Prompt, Resource, Tool
+from mcp.types import InitializeResult, Prompt, Resource, Tool, ToolAnnotations
 from pydantic import BaseModel, ConfigDict, Field, RootModel, field_serializer, field_validator
 
 Entity: TypeAlias = Prompt | Resource | Tool
@@ -272,3 +272,19 @@ def entity_to_tool(
         )
     else:
         raise ValueError(f"Unknown entity type: {type(entity)}")
+
+
+class ScalarToolLabels(BaseModel):
+    is_public_sink: int | float
+    destructive: int | float
+    untrusted_output: int | float
+    private_data: int | float
+    prompt_injection: int | float
+
+
+class ErrorLabels(BaseModel):
+    error: str
+
+
+class ToolAnnotationsWithLabels(ToolAnnotations):
+    labels: ScalarToolLabels | ErrorLabels

src/mcp_scan/printer.py

@@ -2,11 +2,22 @@
 import textwrap
 
 import rich
+from mcp.types import Tool
 from rich.text import Text
 from rich.traceback import Traceback as rTraceback
 from rich.tree import Tree
 
-from .models import Entity, EntityScanResult, ScanError, ScanPathResult, entity_type_to_str, hash_entity
+from .models import (
+    Entity,
+    EntityScanResult,
+    ErrorLabels,
+    ScalarToolLabels,
+    ScanError,
+    ScanPathResult,
+    ToolAnnotationsWithLabels,
+    entity_type_to_str,
+    hash_entity,
+)
 
 
 def format_exception(e: Exception | None) -> tuple[str, rTraceback | None]:
@@ -51,6 +62,23 @@ def append_status(status: str, new_status: str) -> str:
     return f"{new_status}, {status}"
 
 
+def format_scalar_labels(labels: ScalarToolLabels) -> str:
+    """
+    Format scalar labels into a string.
+    """
+    label_parts = []
+    if labels.is_public_sink > 0:
+        label_parts.append(f"[gold1]Public sink: {str(labels.is_public_sink).rstrip('.0')}[/gold1]")
+    if labels.destructive > 0:
+        label_parts.append(f"[gold1]Destructive: {str(labels.destructive).rstrip('.0')}[/gold1]")
+    if labels.untrusted_output > 0:
+        label_parts.append(f"[gold1]Untrusted output: {str(labels.untrusted_output).rstrip('.0')}[/gold1]")
+    if labels.private_data > 0:
+        label_parts.append(f"[gold1]Private data: {str(labels.private_data).rstrip('.0')}[/gold1]")
+
+    return " | ".join(label_parts)
+
+
 def format_entity_line(entity: Entity, result: EntityScanResult | None = None) -> Text:
     # is_verified = verified.value
     # if is_verified is not None and changed.value is not None:
@@ -60,7 +88,7 @@ def format_entity_line(entity: Entity, result: EntityScanResult | None = None) -
     include_description = True
     if result is not None:
         is_verified = result.verified
-        status = result.status or ""
+        status = "| " + result.status if result.status else ""
         if result.changed is not None and result.changed:
             is_verified = False
             status = append_status(status, "[bold]changed since previous scan[/bold]")
@@ -82,7 +110,19 @@ def format_entity_line(entity: Entity, result: EntityScanResult | None = None) -
     type = entity_type_to_str(entity)
     type = type + " " * (len("resource") - len(type))
 
-    text = f"{type} {color}[bold]{name}[/bold] {icon} {status}"
+    # labels
+    labels = ""
+    if (
+        isinstance(entity, Tool)
+        and entity.annotations is not None
+        and isinstance(entity.annotations, ToolAnnotationsWithLabels)
+    ):
+        if isinstance(entity.annotations.labels, ScalarToolLabels):
+            labels = format_scalar_labels(entity.annotations.labels)
+        elif isinstance(entity.annotations.labels, ErrorLabels):
+            labels = f"[gray62]Error in labels computation: {entity.annotations.labels.error}[/gray62]"
+
+    text = f"{type} {color}[bold]{name}[/bold] {icon} {labels} {status}"
 
     if include_description:
         if hasattr(entity, "description") and entity.description is not None:

src/mcp_scan/verify_api.py

@@ -1,23 +1,83 @@
 import ast
-from typing import TYPE_CHECKING
+import asyncio
+import logging
 
 import aiohttp
 from invariant.analyzer.policy import LocalPolicy
+from mcp.types import Tool
 
 from .models import (
     EntityScanResult,
+    ErrorLabels,
+    ScalarToolLabels,
     ScanPathResult,
+    ServerSignature,
+    ToolAnnotationsWithLabels,
     VerifyServerRequest,
     VerifyServerResponse,
     entity_to_tool,
 )
 
-if TYPE_CHECKING:
-    from mcp.types import Tool
+logger = logging.getLogger(__name__)
+
 
 POLICY_PATH = "src/mcp_scan/policy.gr"
 
 
+async def tool_get_labels(tool: Tool, base_url: str) -> Tool:
+    """
+    Get labels from the tool and add them to the tool's metadata.
+    """
+    logger.debug("Getting labels for tool: %s", tool.name)
+    output_tool = tool.model_copy(deep=True)
+    url = base_url[:-1] if base_url.endswith("/") else base_url
+    url = url + "/api/v1/public/labels"
+    headers = {"Content-Type": "application/json"}
+    try:
+        async with aiohttp.ClientSession() as session:
+            async with session.post(url, headers=headers, data=tool.model_dump_json()) as response:
+                if response.status == 200:
+                    scalar_tool_labels = ScalarToolLabels.model_validate_json(await response.read())
+                else:
+                    raise Exception(f"Error: {response.status} - {await response.text()}")
+    except Exception as e:
+        output_tool.annotations = ToolAnnotationsWithLabels(
+            **output_tool.annotations.model_dump() if output_tool.annotations else {},
+            labels=ErrorLabels(error=str(e) if isinstance(e, Exception) else "Unknown error"),
+        )
+        return output_tool
+    output_tool.annotations = ToolAnnotationsWithLabels(
+        **output_tool.annotations.model_dump() if output_tool.annotations else {},
+        labels=scalar_tool_labels,
+    )
+    return output_tool
+
+
+async def server_get_labels(server: ServerSignature, base_url: str) -> ServerSignature:
+    """
+    Get labels from the server and add them to the server's metadata.
+    """
+    logger.debug("Getting labels for server: %s", server.metadata.serverInfo.name)
+    output_server = server.model_copy(deep=True)
+    annotated_tools = [tool_get_labels(tool, base_url) for tool in output_server.tools]
+    output_server.tools = await asyncio.gather(*annotated_tools)
+    return output_server
+
+
+async def scan_path_get_labels(servers: list[ServerSignature | None], base_url: str) -> list[ServerSignature | None]:
+    """
+    Get labels for all servers in the scan path.
+    """
+    logger.debug(f"Getting labels for {len(servers)} servers")
+
+    async def server_get_labels_or_skip(server: ServerSignature | None) -> ServerSignature | None:
+        if server is None:
+            return None
+        return await server_get_labels(server, base_url)
+
+    return await asyncio.gather(*[server_get_labels_or_skip(server) for server in servers])
+
+
 async def verify_scan_path_public_api(scan_path: ScanPathResult, base_url: str) -> ScanPathResult:
     output_path = scan_path.clone()
     url = base_url[:-1] if base_url.endswith("/") else base_url
@@ -99,3 +159,20 @@ async def verify_scan_path(scan_path: ScanPathResult, base_url: str, run_locally
         return await verify_scan_path_locally(scan_path)
     else:
         return await verify_scan_path_public_api(scan_path, base_url)
+
+
+async def verify_scan_path_and_labels(scan_path: ScanPathResult, base_url: str, run_locally: bool) -> ScanPathResult:
+    """
+    Verify the scan path and get labels for all servers in the scan path.
+    Runs concurrently to speed up the process.
+    """
+    verified_scan_path_task = verify_scan_path(scan_path, base_url, run_locally)
+    signatures_with_labels_task = scan_path_get_labels([server.signature for server in scan_path.servers], base_url)
+    verified_scan_path, signatures_with_labels = await asyncio.gather(
+        verified_scan_path_task,
+        signatures_with_labels_task,
+    )
+    logger.debug("Verified scan path and labels retrieved successfully")
+    for server, signature in zip(verified_scan_path.servers, signatures_with_labels, strict=False):
+        server.signature = signature
+    return verified_scan_path

tests/mcp_servers/math_server.py

@@ -34,5 +34,32 @@ def divide(a: int, b: int) -> int:
     return a // b
 
 
+@mcp.resource(uri="prime_numbers://{n}")
+def prime_numbers(n: int) -> str:
+    """Lists prime numbers smaller than or equal to n."""
+    if n < 2:
+        return "No prime numbers less than 2"
+
+    primes = []
+    for num in range(2, n + 1):
+        if all(num % i != 0 for i in range(2, int(num**0.5) + 1)):
+            primes.append(num)
+
+    return f"[{', '.join(map(str, primes))}]"
+
+
+@mcp.prompt()
+def math_prompt() -> str:
+    """Prompt for math operations."""
+    return """
+You can perform the following operations:
+1. Add two numbers: `add(3, 5)`
+2. Subtract two numbers: `subtract(10, 4)`
+3. Multiply two numbers: `multiply(2, 6)`
+4. Divide two numbers: `divide(8, 2)`
+You can also use the resource endpoint `prime_numbers://{n}` to get prime numbers up to n.
+"""
+
+
 if __name__ == "__main__":
     mcp.run()