feat: Add ML pipeline STRIDE coverage

README.md

@@ -103,7 +103,8 @@ The tree-sitter discovery pipeline is tuned for **web-service shapes**. What wor
 | YAML / GitHub Actions workflows | Some — overly-broad permissions and similar config issues |
 | CLI tools (cobra, click, argparse, urfave/cli) | Not modeled — produces structurally complete but empty output |
 | Crypto/signing **client** libraries (sigstore-python, in-toto) | Out of scope — these call out rather than receive; entry-point queries don't fire |
-| Systems software, daemons, libraries, ML pipelines | Not modeled |
+| ML pipelines (PyTorch, Transformers, ONNX, Safetensors) | Modeled — mapped to TAMPERING and ELEVATION_OF_PRIVILEGE |
+| Systems software, daemons, libraries | Not modeled |
 
 If your project doesn't match a supported shape, the generator will still write a report — but it will likely show "Total findings: 0" because no entry points were discovered. That's a coverage gap on our side, not a clean bill of health. Expanding the query set is [tracked in our issue tracker](https://github.com/kusari-oss/darnit/issues?q=is%3Aissue+threat-model+coverage).
 

packages/darnit-baseline/src/darnit_baseline/threat_model/discovery_models.py

@@ -75,6 +75,7 @@ class EntryPointKind(str, Enum):
     MCP_TOOL = "mcp_tool"
     CLI_COMMAND = "cli_command"
     MESSAGE_HANDLER = "message_handler"
+    ML_PIPELINE = "ml_pipeline"
 
 
 class DataStoreKind(str, Enum):

packages/darnit-baseline/src/darnit_baseline/threat_model/queries/python.py

@@ -82,6 +82,35 @@
 """,
 )
 
+# ---------------------------------------------------------------------------
+# ML Pipeline queries
+# ---------------------------------------------------------------------------
+
+#: ML module imports
+ML_MODULE_IMPORT = make_query(
+    "python",
+    """
+(import_statement
+  name: (dotted_name) @name
+  (#match? @name "^(torch|transformers|onnx|safetensors)$"))
+(import_from_statement
+  module_name: (dotted_name) @module
+  (#match? @module "^(torch|transformers|onnx|safetensors)$"))
+""",
+)
+
+#: ML model loading execution blocks
+ML_MODEL_LOADER = make_query(
+    "python",
+    """
+(call
+  function: (attribute
+    object: (identifier) @obj
+    attribute: (identifier) @method)
+  (#match? @method "^(load|load_state_dict|from_pretrained|safe_open|load_file|load_weights|load_model)$")) @call
+""",
+)
+
 # ---------------------------------------------------------------------------
 # Subprocess / dangerous-call queries
 # ---------------------------------------------------------------------------
@@ -324,6 +353,16 @@ class PythonQuery:
         query=HTTP_ROUTE_IMPERATIVE,
         intent="constructor_call",
     ),
+    "python.entry.ml_import": PythonQuery(
+        id="python.entry.ml_import",
+        query=ML_MODULE_IMPORT,
+        intent="import",
+    ),
+    "python.entry.ml_loader": PythonQuery(
+        id="python.entry.ml_loader",
+        query=ML_MODEL_LOADER,
+        intent="bare_call",
+    ),
     "python.sink.dangerous_attr": PythonQuery(
         id="python.sink.dangerous_attr",
         query=DANGEROUS_ATTRIBUTE_CALL,
@@ -399,6 +438,8 @@ class PythonQuery:
     "MCP_TOOL_DECORATOR",
     "MCP_TOOL_IMPERATIVE",
     "HTTP_ROUTE_IMPERATIVE",
+    "ML_MODULE_IMPORT",
+    "ML_MODEL_LOADER",
     "DANGEROUS_ATTRIBUTE_CALL",
     "DANGEROUS_BARE_CALL",
     "DATASTORE_ATTRIBUTE_CALL",

packages/darnit-baseline/src/darnit_baseline/threat_model/ts_discovery.py

@@ -548,6 +548,54 @@ def _extract_python_entry_points(
             )
         )
 
+    # ML Pipelines - Imports
+    query = py_queries.QUERY_REGISTRY["python.entry.ml_import"].query
+    for caps in run_query(query, tree.root_node):
+        module_name = ""
+        whole = None
+        if "name" in caps:
+            name_node = caps["name"][0]
+            module_name = _text(name_node, source).split(".", 1)[0]
+            whole = name_node.parent
+        elif "module" in caps:
+            module_node = caps["module"][0]
+            module_name = _text(module_node, source).split(".", 1)[0]
+            whole = module_node.parent
+        if module_name and whole:
+            entries.append(
+                DiscoveredEntryPoint(
+                    kind=EntryPointKind.ML_PIPELINE,
+                    name=f"ML Import ({module_name})",
+                    location=_build_location(whole, file.relpath),
+                    language="python",
+                    framework=module_name,
+                    route_path=None,
+                    http_method=None,
+                    has_auth_decorator=False,
+                    source_query="python.entry.ml_import",
+                )
+            )
+
+    # ML Pipelines - Loaders
+    query = py_queries.QUERY_REGISTRY["python.entry.ml_loader"].query
+    for caps in run_query(query, tree.root_node):
+        call_node = caps["call"][0]
+        obj = _text(caps["obj"][0], source)
+        method = _text(caps["method"][0], source)
+        entries.append(
+            DiscoveredEntryPoint(
+                kind=EntryPointKind.ML_PIPELINE,
+                name=f"ML Loader ({obj}.{method})",
+                location=_build_location(call_node, file.relpath),
+                language="python",
+                framework=obj,
+                route_path=None,
+                http_method=None,
+                has_auth_decorator=False,
+                source_query="python.entry.ml_loader",
+            )
+        )
+
     return entries
 
 
@@ -1726,6 +1774,8 @@ def _spoofing_findings_from_entry_points(
     findings: list[CandidateFinding] = []
 
     for ep in entry_points:
+        if ep.kind == EntryPointKind.ML_PIPELINE:
+            continue
         if ep.has_auth_decorator:
             continue  # decorated with auth — no spoofing concern
         source = _get_source_for(ep.location.file, scanned_files, file_content_cache)
@@ -1808,6 +1858,72 @@ def _info_disclosure_findings_from_data_stores(
     return findings
 
 
+def _tampering_findings_from_ml_pipelines(
+    entry_points: list[DiscoveredEntryPoint],
+    scanned_files: list[ScannedFile],
+) -> list[CandidateFinding]:
+    """Emit Tampering and Elevation of Privilege findings for ML pipelines.
+
+    ML models and data pipelines load untrusted formats (like pickles) which
+    are prone to arbitrary code execution, and unverified data loads allow
+    model poisoning.
+    """
+    file_content_cache: dict[str, bytes] = {}
+    findings: list[CandidateFinding] = []
+
+    for ep in entry_points:
+        if ep.kind != EntryPointKind.ML_PIPELINE:
+            continue
+
+        source = _get_source_for(ep.location.file, scanned_files, file_content_cache)
+        if source is None:
+            continue
+
+        severity_tamp = severity_for(StrideCategory.TAMPERING, has_taint_trace=False)
+        confidence = confidence_for(FindingSource.TREE_SITTER_STRUCTURAL, query_intent="bare_call")
+        snippet = _build_snippet(source, ep.location.line)
+
+        findings.append(
+            CandidateFinding(
+                category=StrideCategory.TAMPERING,
+                title=f"Potential ML Model Poisoning via {ep.name}",
+                source=FindingSource.TREE_SITTER_STRUCTURAL,
+                primary_location=ep.location,
+                related_assets=(ep.id,),
+                code_snippet=snippet,
+                severity=severity_tamp,
+                confidence=confidence,
+                rationale=(
+                    f"A Machine Learning pipeline using {ep.framework} was detected. "
+                    "If the model or data is loaded from an untrusted source, attackers "
+                    "can manipulate the model outputs (Model Poisoning / Adversarial Tampering)."
+                ),
+                query_id=ep.source_query,
+            )
+        )
+
+        severity_eop = severity_for(StrideCategory.ELEVATION_OF_PRIVILEGE, has_taint_trace=False)
+        findings.append(
+            CandidateFinding(
+                category=StrideCategory.ELEVATION_OF_PRIVILEGE,
+                title=f"Potential Arbitrary Code Execution in ML loader {ep.name}",
+                source=FindingSource.TREE_SITTER_STRUCTURAL,
+                primary_location=ep.location,
+                related_assets=(ep.id,),
+                code_snippet=snippet,
+                severity=severity_eop,
+                confidence=confidence,
+                rationale=(
+                    f"Loading models via {ep.framework} (e.g. {ep.name}) often relies on "
+                    "insecure deserialization (like Python's pickle module). Loading untrusted "
+                    "files can lead to arbitrary code execution."
+                ),
+                query_id=ep.source_query,
+            )
+        )
+    return findings
+
+
 def _get_source_for(
     relpath: str,
     scanned_files: list[ScannedFile],
@@ -1991,6 +2107,7 @@ def discover_all(
     # would otherwise be empty when only subprocess/eval findings exist.
     findings.extend(_spoofing_findings_from_entry_points(entry_points, scanned_files))
     findings.extend(_info_disclosure_findings_from_data_stores(data_stores, scanned_files))
+    findings.extend(_tampering_findings_from_ml_pipelines(entry_points, scanned_files))
 
     # Opengrep enrichment: run bundled rules for taint analysis + structural
     # pattern matching when the binary is available. On absence or failure,

packages/darnit-hello/ml_test.py

@@ -0,0 +1,14 @@
+import safetensors.torch
+import torch
+from transformers import AutoModelForSequenceClassification
+
+
+def load_models():
+    # Load PyTorch model (EOP/Tampering)
+    torch.load("model.pt")
+
+    # Load via safetensors
+    safetensors.torch.load_file("model.safetensors")
+
+    # Load via Transformers
+    AutoModelForSequenceClassification.from_pretrained("bert-base-uncased")

test_ml.py

@@ -0,0 +1,19 @@
+from pathlib import Path
+
+from darnit_baseline.threat_model.ts_discovery import discover_all
+
+
+def main():
+    repo_root = Path("packages/darnit-hello")
+    res = discover_all(repo_root)
+    print(f"Total findings: {len(res.findings)}")
+    print("Findings:")
+    for f in res.findings:
+        print(f"[{f.category.value}] {f.title} (confidence: {f.confidence.value})")
+
+    print("\nEntry points:")
+    for e in res.entry_points:
+        print(f"[{e.kind.value}] {e.name}")
+
+if __name__ == "__main__":
+    main()

tests/darnit_baseline/test_handler_dispatch_integration.py

@@ -546,7 +546,10 @@ def test_osps_do_02_01_passes_for_bug_template(self, tmp_path, template_filename
         result = orchestrator.verify(control, context)
 
         assert result.status == "PASS"
-        assert result.evidence.get("relative_path") == f".github/ISSUE_TEMPLATE/{template_filename}"
+        evidence_path = result.evidence.get("relative_path")
+        if evidence_path:
+            evidence_path = evidence_path.replace("\\", "/")
+        assert evidence_path == f".github/ISSUE_TEMPLATE/{template_filename}"
 
     @pytest.mark.unit
     def test_osps_do_02_01_fails_for_feature_only_template(self, tmp_path):

tests/darnit_baseline/threat_model/test_file_discovery.py

@@ -46,7 +46,7 @@ def test_baseline_exclusions_are_honored(tmp_path: Path) -> None:
     _write(tmp_path / "__pycache__/cached.pyc", "")
 
     files, stats = walk_repo(tmp_path)
-    relpaths = {f.relpath for f in files}
+    relpaths = {f.relpath.replace("\\", "/") for f in files}
 
     assert "src/app.py" in relpaths
     assert not any("node_modules" in p for p in relpaths)
@@ -65,7 +65,7 @@ def test_gitignore_dirs_excluded(tmp_path: Path) -> None:
     _write(tmp_path / "custom_cache/blob.py", "x = 3\n")
 
     files, stats = walk_repo(tmp_path)
-    relpaths = {f.relpath for f in files}
+    relpaths = {f.relpath.replace("\\", "/") for f in files}
 
     assert "src/a.py" in relpaths
     assert not any("my_build" in p for p in relpaths)
@@ -93,7 +93,7 @@ def test_extra_excludes_are_additive(tmp_path: Path) -> None:
     _write(tmp_path / "custom_vendor/x.py", "pass\n")  # user exclusion
 
     files, _ = walk_repo(tmp_path, extra_excludes=["custom_vendor"])
-    relpaths = {f.relpath for f in files}
+    relpaths = {f.relpath.replace("\\", "/") for f in files}
 
     assert "src/app.py" in relpaths
     assert not any("node_modules" in p for p in relpaths)
@@ -167,7 +167,7 @@ def test_relpath_is_repo_relative(tmp_path: Path) -> None:
     _write(tmp_path / "deep/nested/file.py", "pass\n")
     files, _ = walk_repo(tmp_path)
     assert len(files) == 1
-    assert files[0].relpath == "deep/nested/file.py"
+    assert files[0].relpath.replace("\\", "/") == "deep/nested/file.py"
     assert files[0].path.is_absolute()
     assert files[0].language == "python"
     assert isinstance(files[0], ScannedFile)

tests/darnit_baseline/threat_model/test_sidecar.py

@@ -47,34 +47,34 @@ def _make_finding(
 
 class TestComputeFingerprint:
     def test_deterministic(self, tmp_path: Path) -> None:
-        f = _make_finding()
+        f = _make_finding(file=str(tmp_path / "src/app.py"))
         fp1 = compute_fingerprint(f, tmp_path)
         fp2 = compute_fingerprint(f, tmp_path)
         assert fp1 == fp2
 
     def test_starts_with_sha256_prefix(self, tmp_path: Path) -> None:
-        fp = compute_fingerprint(_make_finding(), tmp_path)
+        fp = compute_fingerprint(_make_finding(file=str(tmp_path / "src/app.py")), tmp_path)
         assert fp.startswith("sha256:")
         assert len(fp) == len("sha256:") + 16
 
     def test_changes_with_file_path(self, tmp_path: Path) -> None:
-        f1 = _make_finding(file="src/a.py")
-        f2 = _make_finding(file="src/b.py")
+        f1 = _make_finding(file=str(tmp_path / "src/a.py"))
+        f2 = _make_finding(file=str(tmp_path / "src/b.py"))
         assert compute_fingerprint(f1, tmp_path) != compute_fingerprint(f2, tmp_path)
 
     def test_changes_with_snippet(self, tmp_path: Path) -> None:
-        f1 = _make_finding(snippet_lines=("x = foo()",))
-        f2 = _make_finding(snippet_lines=("x = bar()",))
+        f1 = _make_finding(file=str(tmp_path / "src/app.py"), snippet_lines=("x = foo()",))
+        f2 = _make_finding(file=str(tmp_path / "src/app.py"), snippet_lines=("x = bar()",))
         assert compute_fingerprint(f1, tmp_path) != compute_fingerprint(f2, tmp_path)
 
     def test_stable_across_whitespace_reformatting(self, tmp_path: Path) -> None:
-        f1 = _make_finding(snippet_lines=("  x = foo()",))
-        f2 = _make_finding(snippet_lines=("x = foo()",))
+        f1 = _make_finding(file=str(tmp_path / "src/app.py"), snippet_lines=("  x = foo()",))
+        f2 = _make_finding(file=str(tmp_path / "src/app.py"), snippet_lines=("x = foo()",))
         assert compute_fingerprint(f1, tmp_path) == compute_fingerprint(f2, tmp_path)
 
     def test_changes_with_query_id(self, tmp_path: Path) -> None:
-        f1 = _make_finding(query_id="python.sink.a")
-        f2 = _make_finding(query_id="python.sink.b")
+        f1 = _make_finding(file=str(tmp_path / "src/app.py"), query_id="python.sink.a")
+        f2 = _make_finding(file=str(tmp_path / "src/app.py"), query_id="python.sink.b")
         assert compute_fingerprint(f1, tmp_path) != compute_fingerprint(f2, tmp_path)
 
 

tests/darnit_baseline/threat_model/test_ts_discovery.py

@@ -774,6 +774,9 @@ def test_finds_entry_points_sc001a(self, result) -> None:
         mcp_tools = [ep for ep in result.entry_points if ep.kind == EntryPointKind.MCP_TOOL]
         assert len(mcp_tools) >= 2, f"Expected at least 2 MCP_TOOL entry points; got {len(mcp_tools)}"
 
+        ml_pipelines = [ep for ep in result.entry_points if ep.kind == EntryPointKind.ML_PIPELINE]
+        assert len(ml_pipelines) >= 1, f"Expected at least 1 ML_PIPELINE entry point; got {len(ml_pipelines)}"
+
     def test_subprocess_scores_differentiated_sc001b(self, result) -> None:
         """SC-001b: subprocess findings must NOT all have identical scores.
 
@@ -808,6 +811,9 @@ def test_stride_coverage_sc009(self, result) -> None:
             f"got {len(categories_with_findings)}: "
             f"{[c.value for c in categories_with_findings]}"
         )
+        assert StrideCategory.ELEVATION_OF_PRIVILEGE in categories_with_findings, (
+            "Expected ELEVATION_OF_PRIVILEGE findings from ML pipeline deserialization"
+        )
 
     def test_imperative_fixture_in_curated_suite_sc002a(self) -> None:
         """SC-002a: imperative registration fixture must produce entry points."""