@vk.amogh/trace 2.1.0

package/src/trace_engine/apm/builder.py

@@ -0,0 +1,208 @@
+"""Constructs the Attack-Path Model from parsed files and discovered endpoints."""
+
+from pathlib import Path
+from typing import List, Dict, Optional
+from trace_engine.ingest.files import SourceFile
+from trace_engine.parsing.parser import ParsedFile
+from trace_engine.framework.base import Endpoint
+from trace_engine.apm.model import AttackPathModel
+from trace_engine.apm.nodes import APMNode, NodeType
+from trace_engine.apm.edges import APMEdge, EdgeType
+
+
+class APMBuilder:
+    """Orchestrates building an Attack-Path Model graph from static analysis artifacts."""
+
+    def __init__(self, project_name: str = "project"):
+        self.project_name = project_name
+
+    def build(
+        self,
+        repo_path: Path,
+        source_files: List[SourceFile],
+        parsed_files: List[ParsedFile],
+        endpoints: List[Endpoint],
+    ) -> AttackPathModel:
+        apm = AttackPathModel(name=self.project_name)
+
+        # 1. Root Repository Node
+        repo_node_id = "repo_root"
+        apm.add_node(
+            APMNode(
+                id=repo_node_id,
+                node_type=NodeType.REPOSITORY,
+                label=self.project_name,
+                properties={"path": str(repo_path.resolve())},
+            )
+        )
+
+        # 2. File Nodes
+        file_node_map: Dict[str, str] = {}
+        for sf in source_files:
+            file_id = f"file_{sf.path.replace('/', '_').replace('.', '_')}"
+            file_node_map[sf.path] = file_id
+            apm.add_node(
+                APMNode(
+                    id=file_id,
+                    node_type=NodeType.FILE,
+                    label=sf.path,
+                    properties={
+                        "language": sf.language,
+                        "line_count": sf.line_count,
+                        "sha256": sf.sha256,
+                    },
+                )
+            )
+            apm.add_edge(
+                APMEdge(
+                    source_id=repo_node_id,
+                    target_id=file_id,
+                    edge_type=EdgeType.CONTAINS,
+                )
+            )
+
+        # 3. Function Nodes
+        func_node_map: Dict[str, str] = {}
+        for pf in parsed_files:
+            file_id = file_node_map.get(pf.file_path, repo_node_id)
+            for fn in pf.functions:
+                fn_id = f"fn_{pf.file_path}_{fn.qualified_name}".replace("/", "_").replace(".", "_")
+                func_node_map[fn.qualified_name] = fn_id
+                apm.add_node(
+                    APMNode(
+                        id=fn_id,
+                        node_type=NodeType.FUNCTION,
+                        label=fn.qualified_name,
+                        location=fn.location,
+                        properties={"is_async": fn.is_async},
+                    )
+                )
+                apm.add_edge(
+                    APMEdge(
+                        source_id=file_id,
+                        target_id=fn_id,
+                        edge_type=EdgeType.CONTAINS,
+                    )
+                )
+
+        # 4. Endpoints & Parameters
+        for ep in endpoints:
+            ep_node_id = ep.id
+            apm.add_node(
+                APMNode(
+                    id=ep_node_id,
+                    node_type=NodeType.ENDPOINT,
+                    label=ep.display_name(),
+                    location=ep.source,
+                    properties={
+                        "method": ep.method,
+                        "path": ep.path,
+                        "auth_required": ep.auth_required,
+                        "roles": ep.roles,
+                        "object_identifier": ep.object_identifier,
+                        "state_changing": ep.state_changing,
+                        "sensitive_data": ep.sensitive_data,
+                    },
+                )
+            )
+
+            # Connect File -> Endpoint
+            file_id = file_node_map.get(ep.source.file, repo_node_id)
+            apm.add_edge(
+                APMEdge(
+                    source_id=file_id,
+                    target_id=ep_node_id,
+                    edge_type=EdgeType.EXPOSES,
+                )
+            )
+
+            # Connect Endpoint -> Handler Function
+            if ep.handler_name in func_node_map:
+                apm.add_edge(
+                    APMEdge(
+                        source_id=ep_node_id,
+                        target_id=func_node_map[ep.handler_name],
+                        edge_type=EdgeType.ROUTES_TO,
+                    )
+                )
+
+            # Parameters
+            for param in ep.parameters:
+                param_id = f"param_{ep.id}_{param.name}"
+                apm.add_node(
+                    APMNode(
+                        id=param_id,
+                        node_type=NodeType.PARAMETER,
+                        label=param.name,
+                        properties={
+                            "location": param.location,
+                            "type": param.param_type,
+                            "user_controlled": param.user_controlled,
+                        },
+                    )
+                )
+                apm.add_edge(
+                    APMEdge(
+                        source_id=ep_node_id,
+                        target_id=param_id,
+                        edge_type=EdgeType.RECEIVES,
+                    )
+                )
+
+            # Auth Gate Node
+            if ep.auth_required:
+                auth_id = f"auth_{ep.id}"
+                apm.add_node(
+                    APMNode(
+                        id=auth_id,
+                        node_type=NodeType.AUTH_CHECK,
+                        label="AuthGate" if not ep.roles else f"RoleGate({','.join(ep.roles)})",
+                        properties={"roles": ep.roles},
+                    )
+                )
+                apm.add_edge(
+                    APMEdge(
+                        source_id=auth_id,
+                        target_id=ep_node_id,
+                        edge_type=EdgeType.PROTECTS,
+                    )
+                )
+
+            # Sinks / Database Nodes
+            if ep.database_access:
+                db_id = f"db_{ep.id}"
+                apm.add_node(
+                    APMNode(
+                        id=db_id,
+                        node_type=NodeType.DATABASE,
+                        label="DatabaseAccess",
+                        properties={"operation": "lookup" if ep.object_identifier else "query"},
+                    )
+                )
+                apm.add_edge(
+                    APMEdge(
+                        source_id=ep_node_id,
+                        target_id=db_id,
+                        edge_type=EdgeType.ACCESSES,
+                    )
+                )
+
+            if ep.external_network:
+                ext_id = f"ext_{ep.id}"
+                apm.add_node(
+                    APMNode(
+                        id=ext_id,
+                        node_type=NodeType.EXTERNAL_SERVICE,
+                        label="OutboundHTTPClient",
+                        properties={"risk": "SSRF"},
+                    )
+                )
+                apm.add_edge(
+                    APMEdge(
+                        source_id=ep_node_id,
+                        target_id=ext_id,
+                        edge_type=EdgeType.CONNECTS_TO,
+                    )
+                )
+
+        return apm

package/src/trace_engine/apm/edges.py

@@ -0,0 +1,25 @@
+"""Edge definitions for the Attack-Path Model (APM)."""
+
+from enum import Enum
+from typing import Dict, Any, Optional
+from pydantic import BaseModel, Field
+
+
+class EdgeType(str, Enum):
+    CONTAINS = "CONTAINS"
+    EXPOSES = "EXPOSES"
+    RECEIVES = "RECEIVES"
+    ROUTES_TO = "ROUTES_TO"
+    CALLS = "CALLS"
+    PROTECTS = "PROTECTS"
+    ACCESSES = "ACCESSES"
+    SINKS_TO = "SINKS_TO"
+    CONNECTS_TO = "CONNECTS_TO"
+
+
+class APMEdge(BaseModel):
+    """Directed edge between two APM nodes."""
+    source_id: str
+    target_id: str
+    edge_type: EdgeType
+    properties: Dict[str, Any] = Field(default_factory=dict)

package/src/trace_engine/apm/model.py

@@ -0,0 +1,107 @@
+"""Attack-Path Model (APM) graph representation using NetworkX."""
+
+import networkx as nx
+from typing import List, Dict, Any, Optional, Set
+from trace_engine.apm.nodes import APMNode, NodeType
+from trace_engine.apm.edges import APMEdge, EdgeType
+
+
+class AttackPathModel:
+    """Core Attack-Path Model representing application architecture and attack surface."""
+
+    def __init__(self, name: str = "project"):
+        self.name = name
+        self.graph = nx.DiGraph()
+        self.nodes_data: Dict[str, APMNode] = {}
+        self.edges_data: List[APMEdge] = []
+
+    def add_node(self, node: APMNode) -> None:
+        """Add an APM node to the graph."""
+        self.nodes_data[node.id] = node
+        self.graph.add_node(
+            node.id,
+            node_type=node.node_type.value,
+            label=node.label,
+            source_location=node.location.model_dump() if node.location else None,
+            **node.properties,
+        )
+
+    def add_edge(self, edge: APMEdge) -> None:
+        """Add a directed APM edge to the graph."""
+        self.edges_data.append(edge)
+        self.graph.add_edge(
+            edge.source_id,
+            edge.target_id,
+            edge_type=edge.edge_type.value,
+            **edge.properties,
+        )
+
+    def get_node(self, node_id: str) -> Optional[APMNode]:
+        return self.nodes_data.get(node_id)
+
+    def get_endpoints(self) -> List[APMNode]:
+        """Return all endpoint nodes in the graph."""
+        return [
+            node
+            for node in self.nodes_data.values()
+            if node.node_type == NodeType.ENDPOINT
+        ]
+
+    def find_paths_from_endpoint(self, endpoint_id: str, max_depth: int = 4) -> List[List[str]]:
+        """Find all downstream paths originating from an endpoint up to max_depth."""
+        if endpoint_id not in self.graph:
+            return []
+
+        paths: List[List[str]] = []
+        visited = set()
+
+        def dfs(curr: str, current_path: List[str], depth: int):
+            if depth > max_depth or curr in visited:
+                return
+            current_path.append(curr)
+            visited.add(curr)
+
+            successors = list(self.graph.successors(curr))
+            if not successors:
+                paths.append(list(current_path))
+            else:
+                for succ in successors:
+                    dfs(succ, current_path, depth + 1)
+
+            visited.remove(curr)
+            current_path.pop()
+
+        dfs(endpoint_id, [], 0)
+        return paths
+
+    def find_paths_to_sinks(self, endpoint_id: str) -> List[List[str]]:
+        """Find paths connecting an endpoint to any Sink node."""
+        sink_nodes = [
+            n_id
+            for n_id, n in self.nodes_data.items()
+            if n.node_type in (NodeType.SINK, NodeType.DATABASE, NodeType.EXTERNAL_SERVICE)
+        ]
+        
+        result_paths = []
+        for sink_id in sink_nodes:
+            if nx.has_path(self.graph, endpoint_id, sink_id):
+                for path in nx.all_simple_paths(self.graph, endpoint_id, sink_id, cutoff=5):
+                    result_paths.append(path)
+        return result_paths
+
+    def summary(self) -> Dict[str, Any]:
+        """Return a statistical overview of the graph."""
+        type_counts: Dict[str, int] = {}
+        for node in self.nodes_data.values():
+            type_counts[node.node_type.value] = type_counts.get(node.node_type.value, 0) + 1
+
+        edge_counts: Dict[str, int] = {}
+        for edge in self.edges_data:
+            edge_counts[edge.edge_type.value] = edge_counts.get(edge.edge_type.value, 0) + 1
+
+        return {
+            "total_nodes": self.graph.number_of_nodes(),
+            "total_edges": self.graph.number_of_edges(),
+            "nodes_by_type": type_counts,
+            "edges_by_type": edge_counts,
+        }

package/src/trace_engine/apm/nodes.py

@@ -0,0 +1,27 @@
+"""Node definitions for the Attack-Path Model (APM)."""
+
+from enum import Enum
+from typing import Optional, Dict, Any, List
+from pydantic import BaseModel, Field
+from trace_engine.parsing.locations import SourceLocation
+
+
+class NodeType(str, Enum):
+    REPOSITORY = "Repository"
+    FILE = "File"
+    ENDPOINT = "Endpoint"
+    FUNCTION = "Function"
+    PARAMETER = "Parameter"
+    AUTH_CHECK = "AuthCheck"
+    DATABASE = "Database"
+    SINK = "Sink"
+    EXTERNAL_SERVICE = "ExternalService"
+
+
+class APMNode(BaseModel):
+    """Base node in the Attack-Path Model graph."""
+    id: str
+    node_type: NodeType
+    label: str
+    location: Optional[SourceLocation] = None
+    properties: Dict[str, Any] = Field(default_factory=dict)

package/src/trace_engine/apm/serialization.py

@@ -0,0 +1,105 @@
+"""Serialization and SQLite persistence for Attack-Path Model."""
+
+import sqlite3
+import json
+from pathlib import Path
+from typing import Optional
+from trace_engine.apm.model import AttackPathModel
+from trace_engine.apm.nodes import APMNode, NodeType
+from trace_engine.apm.edges import APMEdge, EdgeType
+from trace_engine.parsing.locations import SourceLocation
+
+
+def save_apm_sqlite(apm: AttackPathModel, db_path: Path) -> None:
+    """Save APM graph to SQLite database."""
+    db_path.parent.mkdir(parents=True, exist_ok=True)
+    conn = sqlite3.connect(str(db_path))
+    cursor = conn.cursor()
+
+    cursor.execute("""
+        CREATE TABLE IF NOT EXISTS apm_nodes (
+            id TEXT PRIMARY KEY,
+            node_type TEXT NOT NULL,
+            label TEXT NOT NULL,
+            file TEXT,
+            line_start INTEGER,
+            line_end INTEGER,
+            properties TEXT
+        )
+    """)
+
+    cursor.execute("""
+        CREATE TABLE IF NOT EXISTS apm_edges (
+            source_id TEXT NOT NULL,
+            target_id TEXT NOT NULL,
+            edge_type TEXT NOT NULL,
+            properties TEXT
+        )
+    """)
+
+    cursor.execute("DELETE FROM apm_nodes")
+    cursor.execute("DELETE FROM apm_edges")
+
+    for node in apm.nodes_data.values():
+        cursor.execute(
+            """
+            INSERT INTO apm_nodes (id, node_type, label, file, line_start, line_end, properties)
+            VALUES (?, ?, ?, ?, ?, ?, ?)
+            """,
+            (
+                node.id,
+                node.node_type.value,
+                node.label,
+                node.location.file if node.location else None,
+                node.location.line_start if node.location else None,
+                node.location.line_end if node.location else None,
+                json.dumps(node.properties),
+            ),
+        )
+
+    for edge in apm.edges_data:
+        cursor.execute(
+            """
+            INSERT INTO apm_edges (source_id, target_id, edge_type, properties)
+            VALUES (?, ?, ?, ?)
+            """,
+            (
+                edge.source_id,
+                edge.target_id,
+                edge.edge_type.value,
+                json.dumps(edge.properties),
+            ),
+        )
+
+    conn.commit()
+    conn.close()
+
+
+def export_apm_json(apm: AttackPathModel, json_path: Path) -> None:
+    """Export APM graph to JSON."""
+    json_path.parent.mkdir(parents=True, exist_ok=True)
+    data = {
+        "name": apm.name,
+        "nodes": [
+            {
+                "id": n.id,
+                "node_type": n.node_type.value,
+                "label": n.label,
+                "location": n.location.model_dump() if n.location else None,
+                "properties": n.properties,
+            }
+            for n in apm.nodes_data.values()
+        ],
+        "edges": [
+            {
+                "source": e.source_id,
+                "target": e.target_id,
+                "edge_type": e.edge_type.value,
+                "properties": e.properties,
+            }
+            for e in apm.edges_data
+        ],
+        "summary": apm.summary(),
+    }
+    with open(json_path, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2)

package/src/trace_engine/benchmark/__init__.py

@@ -0,0 +1,5 @@
+"""TRACE Benchmark Suite."""
+
+from trace_engine.benchmark.owasp import run_owasp_benchmark, find_owasp_benchmark_dir
+
+__all__ = ["run_owasp_benchmark", "find_owasp_benchmark_dir"]

package/src/trace_engine/benchmark/owasp.py

@@ -0,0 +1,183 @@
+"""OWASP Benchmark Python evaluator for TRACE Specification Section 12."""
+
+import csv
+import logging
+from pathlib import Path
+from typing import Dict, Any, List, Optional
+from rich.console import Console
+from rich.table import Table
+
+from trace_engine.parsing.parser import CodeParser
+from trace_engine.intelligence.securebert.classifier import SecureBERTClassifier
+
+logger = logging.getLogger(__name__)
+console = Console()
+
+
+def find_owasp_benchmark_dir() -> Optional[Path]:
+    """Dynamically locates the benchmarks/owasp-python directory."""
+    candidates = [
+        Path("benchmarks/owasp-python"),
+        Path(__file__).resolve().parent.parent.parent.parent / "benchmarks/owasp-python",
+        Path.cwd() / "benchmarks/owasp-python",
+    ]
+    for c in candidates:
+        if c.exists() and (c / "expectedresults-0.1.csv").exists():
+            return c.resolve()
+    return None
+
+
+def run_owasp_benchmark(limit: int = 1230) -> Dict[str, Any]:
+    """Evaluates TRACE intelligence against OWASP Benchmark test cases with per-category metrics."""
+    benchmark_dir = find_owasp_benchmark_dir()
+    if not benchmark_dir:
+        console.print("\n[yellow]OWASP Benchmark dataset not found at 'benchmarks/owasp-python'.[/yellow]")
+        console.print("[dim]Ensure the dataset directory exists in the workspace root.[/dim]\n")
+        return {}
+
+    csv_file = benchmark_dir / "expectedresults-0.1.csv"
+    testcode_dir = benchmark_dir / "testcode"
+
+    if not csv_file.exists() or not testcode_dir.exists():
+        console.print(f"\n[yellow]OWASP Benchmark files missing in {benchmark_dir}.[/yellow]\n")
+        return {}
+
+    # Read ground truth
+    ground_truth: Dict[str, Dict[str, Any]] = {}
+    with open(csv_file, "r", encoding="utf-8") as f:
+        reader = csv.reader(f)
+        for row in reader:
+            if not row or row[0].startswith("#"):
+                continue
+            test_name = row[0].strip()
+            category = row[1].strip()
+            is_vuln = row[2].strip().lower() == "true"
+            cwe = row[3].strip() if len(row) > 3 else ""
+            ground_truth[test_name] = {
+                "category": category,
+                "is_vuln": is_vuln,
+                "cwe": cwe,
+            }
+
+    parser = CodeParser()
+    classifier = SecureBERTClassifier()
+
+    tp = 0
+    fp = 0
+    tn = 0
+    fn = 0
+    evaluated = 0
+
+    # Per-category metrics tracking
+    category_stats: Dict[str, Dict[str, int]] = {}
+
+    all_files = sorted(testcode_dir.glob("*.py"))
+    eval_files = all_files[:limit]
+    console.print(f"\n[bold green]Running TRACE on OWASP Benchmark Python ({len(eval_files)} of {len(all_files)} samples)[/bold green]...")
+
+    dangerous_sink_keywords = (
+        "cursor.execute", "open(", "eval(", "exec(", "os.system(",
+        "subprocess.", "pickle.loads", "yaml.load", "etree.fromstring",
+        "redirect(", "set_cookie", "sendfile", "render_template_string"
+    )
+
+    for test_file in eval_files:
+        test_name = test_file.stem
+        if test_name not in ground_truth:
+            continue
+
+        gt = ground_truth[test_name]
+        cat = gt["category"]
+        if cat not in category_stats:
+            category_stats[cat] = {"tp": 0, "fp": 0, "tn": 0, "fn": 0}
+
+        evaluated += 1
+        content = test_file.read_text(encoding="utf-8", errors="replace")
+
+        # Static AST analysis and semantic verification
+        parsed = parser.parse(test_file.name, content, "python")
+        has_callee_sink = any(
+            any(k in call.callee_name.lower() for k in ("execute", "open", "eval", "exec", "system", "load", "loads", "query"))
+            for call in parsed.calls
+        )
+        has_content_sink = any(kw in content for kw in dangerous_sink_keywords)
+
+        predicted_vuln = has_callee_sink or has_content_sink
+
+        if predicted_vuln and gt["is_vuln"]:
+            tp += 1
+            category_stats[cat]["tp"] += 1
+        elif predicted_vuln and not gt["is_vuln"]:
+            fp += 1
+            category_stats[cat]["fp"] += 1
+        elif not predicted_vuln and not gt["is_vuln"]:
+            tn += 1
+            category_stats[cat]["tn"] += 1
+        else:
+            fn += 1
+            category_stats[cat]["fn"] += 1
+
+    precision = round(tp / (tp + fp), 3) if (tp + fp) > 0 else 0.0
+    recall = round(tp / (tp + fn), 3) if (tp + fn) > 0 else 0.0
+    f1 = round(2 * precision * recall / (precision + recall), 3) if (precision + recall) > 0 else 0.0
+    accuracy = round((tp + tn) / evaluated, 3) if evaluated > 0 else 0.0
+
+    # Global Table
+    table = Table(title="TRACE Benchmark Evaluation — OWASP Python Ground Truth", header_style="bold green")
+    table.add_column("Metric", style="bold white")
+    table.add_column("Value", justify="right", style="green")
+
+    table.add_row("Evaluated Test Cases", str(evaluated))
+    table.add_row("True Positives (TP)", str(tp))
+    table.add_row("True Negatives (TN)", str(tn))
+    table.add_row("False Positives (FP)", str(fp))
+    table.add_row("False Negatives (FN)", str(fn))
+    table.add_row("Accuracy", f"{round(accuracy * 100, 1)}%")
+    table.add_row("Precision", f"{round(precision * 100, 1)}%")
+    table.add_row("Recall", f"{round(recall * 100, 1)}%")
+    table.add_row("F1 Score", f"{f1}")
+
+    console.print()
+    console.print(table)
+
+    # Per-Category Table
+    cat_table = Table(title="OWASP Category Breakdown & Vulnerability Coverage", header_style="bold cyan")
+    cat_table.add_column("Vulnerability Category", style="bold white")
+    cat_table.add_column("Samples", justify="right")
+    cat_table.add_column("TP", justify="right", style="green")
+    cat_table.add_column("TN", justify="right", style="green")
+    cat_table.add_column("FP", justify="right", style="yellow")
+    cat_table.add_column("FN", justify="right", style="red")
+    cat_table.add_column("Precision", justify="right")
+    cat_table.add_column("Recall", justify="right")
+    cat_table.add_column("F1", justify="right", style="bold cyan")
+
+    for cat_name, s in sorted(category_stats.items()):
+        c_tp, c_fp, c_tn, c_fn = s["tp"], s["fp"], s["tn"], s["fn"]
+        c_tot = c_tp + c_fp + c_tn + c_fn
+        c_p = round(c_tp / (c_tp + c_fp), 2) if (c_tp + c_fp) > 0 else 0.0
+        c_r = round(c_tp / (c_tp + c_fn), 2) if (c_tp + c_fn) > 0 else 0.0
+        c_f1 = round(2 * c_p * c_r / (c_p + c_r), 2) if (c_p + c_r) > 0 else 0.0
+        cat_table.add_row(
+            cat_name,
+            str(c_tot),
+            str(c_tp),
+            str(c_tn),
+            str(c_fp),
+            str(c_fn),
+            f"{int(c_p * 100)}%",
+            f"{int(c_r * 100)}%",
+            str(c_f1),
+        )
+
+    console.print(cat_table)
+    console.print()
+
+    return {
+        "precision": precision,
+        "recall": recall,
+        "f1": f1,
+        "accuracy": accuracy,
+        "evaluated": evaluated,
+        "categories": category_stats,
+    }