fix: add flows

Author: mmilanta

src/mcp_scan/printer.py

@@ -14,11 +14,15 @@
     ScalarToolLabels,
     ScanError,
     ScanPathResult,
+    ServerScanResult,
     ToolAnnotationsWithLabels,
     entity_type_to_str,
     hash_entity,
 )
 
+MAX_ENTITY_NAME_LENGTH = 25
+MAX_ENTITY_NAME_TOXIC_FLOW_LENGTH = 30
+
 
 def format_exception(e: Exception | None) -> tuple[str, rTraceback | None]:
     if e is None:
@@ -68,15 +72,15 @@ def format_scalar_labels(labels: ScalarToolLabels) -> str:
     """
     label_parts = []
     if labels.is_public_sink > 0:
-        label_parts.append(f"[gold1]Public sink: {str(labels.is_public_sink).rstrip('.0')}[/gold1]")
+        label_parts.append("Public sink")
     if labels.destructive > 0:
-        label_parts.append(f"[gold1]Destructive: {str(labels.destructive).rstrip('.0')}[/gold1]")
+        label_parts.append("Destructive")
     if labels.untrusted_output > 0:
-        label_parts.append(f"[gold1]Untrusted output: {str(labels.untrusted_output).rstrip('.0')}[/gold1]")
+        label_parts.append("Untrusted output")
     if labels.private_data > 0:
-        label_parts.append(f"[gold1]Private data: {str(labels.private_data).rstrip('.0')}[/gold1]")
+        label_parts.append("Private data")
 
-    return " | ".join(label_parts)
+    return "[gray62]" + " | ".join(label_parts) + "[/gray62]"
 
 
 def format_entity_line(entity: Entity, result: EntityScanResult | None = None) -> Text:
@@ -102,9 +106,9 @@ def format_entity_line(entity: Entity, result: EntityScanResult | None = None) -
 
     # right-pad & truncate name
     name = entity.name
-    if len(name) > 25:
-        name = name[:22] + "..."
-    name = name + " " * (25 - len(name))
+    if len(name) > MAX_ENTITY_NAME_LENGTH:
+        name = name[: (MAX_ENTITY_NAME_LENGTH - 3)] + "..."
+    name = name + " " * (MAX_ENTITY_NAME_LENGTH - len(name))
 
     # right-pad type
     type = entity_type_to_str(entity)
@@ -116,13 +120,14 @@ def format_entity_line(entity: Entity, result: EntityScanResult | None = None) -
         isinstance(entity, Tool)
         and entity.annotations is not None
         and isinstance(entity.annotations, ToolAnnotationsWithLabels)
+        and is_verified is not False
     ):
         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}"
+    text = f"{type} {color}[bold]{name}[/bold] {icon} {status} {labels}"
 
     if include_description:
         if hasattr(entity, "description") and entity.description is not None:
@@ -148,6 +153,86 @@ def format_entity_line(entity: Entity, result: EntityScanResult | None = None) -
     return formatted_text
 
 
+def format_tool_flow(tool_name: str, server_name: str, value: float) -> Text:
+    text = "{tool_name} {risk}"
+    tool_name = f"{server_name}/{tool_name}"
+    if len(tool_name) > MAX_ENTITY_NAME_TOXIC_FLOW_LENGTH:
+        tool_name = tool_name[: (MAX_ENTITY_NAME_TOXIC_FLOW_LENGTH - 3)] + "..."
+    tool_name = tool_name + " " * (MAX_ENTITY_NAME_TOXIC_FLOW_LENGTH - len(tool_name))
+
+    risk = "[gold1]Mild[/gold1]" if value <= 1.5 else "[red]High[/red]"
+    return Text.from_markup(text.format(tool_name=tool_name, risk=risk))
+
+
+def format_toxic_flows(servers: list[ServerScanResult]) -> list[Tree]:
+    """
+    Format toxic flows from the scan results into a tree structure.
+    """
+    untrusted_output_tools: list[tuple[str, str, float]] = []
+    destructive_tools: list[tuple[str, str, float]] = []
+    private_data_tools: list[tuple[str, str, float]] = []
+    is_public_sink_tools: list[tuple[str, str, float]] = []
+
+    for server in servers:
+        if server.signature is None:
+            continue
+        for tool in server.signature.tools:
+            if (
+                tool.annotations is not None
+                and isinstance(tool.annotations, ToolAnnotationsWithLabels)
+                and isinstance(tool.annotations.labels, ScalarToolLabels)
+            ):
+                if tool.annotations.labels.untrusted_output > 0:
+                    untrusted_output_tools.append(
+                        (tool.name, server.name or "", tool.annotations.labels.untrusted_output)
+                    )
+                if tool.annotations.labels.destructive > 0:
+                    destructive_tools.append((tool.name, server.name or "", tool.annotations.labels.destructive))
+                if tool.annotations.labels.private_data > 0:
+                    private_data_tools.append((tool.name, server.name or "", tool.annotations.labels.private_data))
+                if tool.annotations.labels.is_public_sink > 0:
+                    is_public_sink_tools.append((tool.name, server.name or "", tool.annotations.labels.is_public_sink))
+
+    untrusted_output_tools.sort(key=lambda x: x[2], reverse=True)
+    destructive_tools.sort(key=lambda x: x[2], reverse=True)
+    private_data_tools.sort(key=lambda x: x[2], reverse=True)
+    is_public_sink_tools.sort(key=lambda x: x[2], reverse=True)
+
+    toxic_flows: list[Tree] = []
+
+    # Flow 1: Untrusted output -> Private data -> Public sink
+    leak_data_flow = Tree("[bold]Leak data flow[/bold]")
+    untrusted_output_tree = Tree("[bold]Untrusted output[/bold]")
+    private_data_tree = Tree("[bold]Private data[/bold]")
+    public_sink_tree = Tree("[bold]Public sink[/bold]")
+    for tool_name, server_name, value in untrusted_output_tools:
+        untrusted_output_tree.add(format_tool_flow(tool_name, server_name, value))
+    for tool_name, server_name, value in private_data_tools:
+        private_data_tree.add(format_tool_flow(tool_name, server_name, value))
+    for tool_name, server_name, value in is_public_sink_tools:
+        public_sink_tree.add(format_tool_flow(tool_name, server_name, value))
+    if len(untrusted_output_tools) > 0 and len(private_data_tools) > 0 and len(is_public_sink_tools) > 0:
+        leak_data_flow.add(untrusted_output_tree)
+        leak_data_flow.add(private_data_tree)
+        leak_data_flow.add(public_sink_tree)
+        toxic_flows.append(leak_data_flow)
+
+    # Flow 2: Untrusted output -> Destructive
+    destructive_flow = Tree("[bold]Harm flow[/bold]")
+    untrusted_output_tree = Tree("[bold]Untrusted output[/bold]")
+    destructive_tree = Tree("[bold]Destructive[/bold]")
+    for tool_name, server_name, value in untrusted_output_tools:
+        untrusted_output_tree.add(format_tool_flow(tool_name, server_name, value))
+    for tool_name, server_name, value in destructive_tools:
+        destructive_tree.add(format_tool_flow(tool_name, server_name, value))
+    if len(untrusted_output_tools) > 0 and len(destructive_tools) > 0:
+        destructive_flow.add(untrusted_output_tree)
+        destructive_flow.add(destructive_tree)
+        toxic_flows.append(destructive_flow)
+
+    return toxic_flows
+
+
 def print_scan_path_result(result: ScanPathResult, print_errors: bool = False) -> None:
     if result.error is not None:
         err_status, traceback = format_error(result.error)
@@ -175,6 +260,14 @@ def print_scan_path_result(result: ScanPathResult, print_errors: bool = False) -
     if len(result.servers) > 0:
         rich.print(path_print_tree)
 
+    toxic_flows = format_toxic_flows(result.servers)
+    if toxic_flows:
+        toxic_flows_tree = Tree("● [bold][gold1]Toxic flows found:[/bold][/gold1]")
+        for flow in toxic_flows:
+            toxic_flows_tree.add(flow)
+        rich.print()
+        rich.print(toxic_flows_tree)
+
     if print_errors and len(server_tracebacks) > 0:
         console = rich.console.Console()
         for server, traceback in server_tracebacks: