diff --git a/README.md b/README.md index 0a6c8e42..9477e40e 100644 --- a/README.md +++ b/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). diff --git a/packages/darnit-baseline/src/darnit_baseline/threat_model/discovery_models.py b/packages/darnit-baseline/src/darnit_baseline/threat_model/discovery_models.py index 1304ed28..fcb40675 100644 --- a/packages/darnit-baseline/src/darnit_baseline/threat_model/discovery_models.py +++ b/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): diff --git a/packages/darnit-baseline/src/darnit_baseline/threat_model/queries/python.py b/packages/darnit-baseline/src/darnit_baseline/threat_model/queries/python.py index 902c8fce..8d08d00b 100644 --- a/packages/darnit-baseline/src/darnit_baseline/threat_model/queries/python.py +++ b/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", diff --git a/packages/darnit-baseline/src/darnit_baseline/threat_model/ts_discovery.py b/packages/darnit-baseline/src/darnit_baseline/threat_model/ts_discovery.py index fea9324b..b9a145f2 100644 --- a/packages/darnit-baseline/src/darnit_baseline/threat_model/ts_discovery.py +++ b/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, diff --git a/packages/darnit-hello/ml_test.py b/packages/darnit-hello/ml_test.py new file mode 100644 index 00000000..ef9e326d --- /dev/null +++ b/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") diff --git a/test_ml.py b/test_ml.py new file mode 100644 index 00000000..de127206 --- /dev/null +++ b/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() diff --git a/tests/darnit_baseline/test_handler_dispatch_integration.py b/tests/darnit_baseline/test_handler_dispatch_integration.py index b0c7c2bb..af87dccb 100644 --- a/tests/darnit_baseline/test_handler_dispatch_integration.py +++ b/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): diff --git a/tests/darnit_baseline/threat_model/test_file_discovery.py b/tests/darnit_baseline/threat_model/test_file_discovery.py index e9e2e9cd..6ed5d557 100644 --- a/tests/darnit_baseline/threat_model/test_file_discovery.py +++ b/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) diff --git a/tests/darnit_baseline/threat_model/test_sidecar.py b/tests/darnit_baseline/threat_model/test_sidecar.py index 91956682..f9d63123 100644 --- a/tests/darnit_baseline/threat_model/test_sidecar.py +++ b/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) diff --git a/tests/darnit_baseline/threat_model/test_ts_discovery.py b/tests/darnit_baseline/threat_model/test_ts_discovery.py index 3c398921..12cacbc4 100644 --- a/tests/darnit_baseline/threat_model/test_ts_discovery.py +++ b/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.""" diff --git a/uv.lock b/uv.lock index e1c47284..ed39613a 100644 --- a/uv.lock +++ b/uv.lock @@ -440,7 +440,7 @@ dependencies = [ requires-dist = [ { name = "darnit-core", editable = "packages/darnit" }, { name = "tree-sitter", specifier = ">=0.25" }, - { name = "tree-sitter-language-pack", specifier = ">=1.5" }, + { name = "tree-sitter-language-pack", specifier = ">=1.5,<1.8" }, ] [[package]]