fix(#1594): [REVIEW] model-supply-chain: add remote-code and final-artifact provenance gates

skills/ai-security/model-supply-chain/__init__.py

@@ -0,0 +1,3 @@
+from .model_supply_chain import ModelSupplyChain
+
+__all__ = ["ModelSupplyChain"]

skills/ai-security/model-supply-chain/model_supply_chain.py

@@ -0,0 +1,43 @@
+from skills import Skill
+from skills.helper import get_model_provenance
+
+class ModelSupplyChain(Skill):
+    def __init__(self):
+        super().__init__()
+        self.potential_provenance_gates = [
+            "trust_remote_code",
+            "final_artifact_provenance",
+            "signed_slsa_attestation",
+        ]
+
+    def audit(self, code):
+        # Existing code to check for model supply chain issues
+        # ...
+
+        # Check for trust_remote_code=True
+        if "from_pretrained" in code and "trust_remote_code=True" in code:
+            self.issues.append("Model loaded with trust_remote_code=True")
+
+        # Check for final artifact provenance
+        if "snapshot_download" in code:
+            model_dir = get_model_provenance(code)
+            if not model_dir:
+                self.issues.append("Model loaded without final artifact provenance")
+
+        # Check for signed SLSA attestation
+        if "signed_slsa_attestation" in code:
+            attestation = get_model_provenance(code)
+            if not attestation:
+                self.issues.append("Model loaded without signed SLSA attestation")
+
+        return self.issues
+
+    def get_provenance_gates(self):
+        return self.potential_provenance_gates
+
+def get_model_provenance(code):
+    # Implement logic to extract model provenance from code
+    # For example, parse the code to extract the model directory
+    # or the signed SLSA attestation
+    # ...
+    pass

skills/ai-security/model-supply-chain/tests/test_model_supply_chain.py

@@ -0,0 +1,23 @@
+import unittest
+from skills.ai_security.model_supply_chain.model_supply_chain import ModelSupplyChain
+
+class TestModelSupplyChain(unittest.TestCase):
+    def test_audit(self):
+        code = """
+        from transformers import AutoModelForCausalLM
+        model = AutoModelForCausalLM.from_pretrained(
+            "research-lab/custom-architecture-llm",
+            revision="main",
+            trust_remote_code=True,
+        )
+        """
+        skill = ModelSupplyChain()
+        issues = skill.audit(code)
+        self.assertIn("Model loaded with trust_remote_code=True", issues)
+
+    def test_get_provenance_gates(self):
+        skill = ModelSupplyChain()
+        gates = skill.get_provenance_gates()
+        self.assertIn("trust_remote_code", gates)
+        self.assertIn("final_artifact_provenance", gates)
+        self.assertIn("signed_slsa_attestation", gates)