81 lines
3.8 KiB
Python
81 lines
3.8 KiB
Python
from pathlib import Path
|
|
import json
|
|
import unittest
|
|
|
|
from aaisa.models import Artifact
|
|
from aaisa.structured_rules import AgentToolPermissionRule, CycloneDxMLBomRule, SlsaProvenanceRule
|
|
|
|
|
|
def json_artifact(name: str, payload: object, kind: str) -> Artifact:
|
|
text = json.dumps(payload)
|
|
return Artifact(Path(name), name, kind, "application/json", len(text), "0" * 64, text)
|
|
|
|
|
|
class StructuredRuleTests(unittest.TestCase):
|
|
def test_unsafe_agent_tool_reports_independent_controls(self) -> None:
|
|
payload = {"tools": [{
|
|
"name": "shell-exec",
|
|
"permissions": ["*"],
|
|
"allow_network": True,
|
|
"allowed_hosts": [],
|
|
"requires_approval": False,
|
|
}]}
|
|
findings = list(AgentToolPermissionRule().evaluate(json_artifact("agent.json", payload, "agent-manifest")))
|
|
self.assertEqual(len(findings), 3)
|
|
self.assertEqual({finding.severity.label() for finding in findings}, {"critical", "high"})
|
|
|
|
def test_constrained_agent_tool_is_accepted(self) -> None:
|
|
payload = {"tools": [{
|
|
"name": "shell-exec",
|
|
"permissions": ["workspace:read"],
|
|
"allow_network": True,
|
|
"allowed_hosts": ["packages.example.invalid"],
|
|
"requires_approval": True,
|
|
}]}
|
|
self.assertFalse(list(AgentToolPermissionRule().evaluate(json_artifact("agent.json", payload, "agent-manifest"))))
|
|
|
|
def test_incomplete_slsa_provenance_is_reported(self) -> None:
|
|
payload = {"_type": "https://in-toto.io/Statement/v1", "predicateType": "https://slsa.dev/provenance/v1", "predicate": {}}
|
|
finding = list(SlsaProvenanceRule().evaluate(json_artifact("provenance.json", payload, "provenance")))[0]
|
|
self.assertIn("subject SHA-256 digest", finding.evidence or "")
|
|
self.assertIn("runDetails.builder.id", finding.evidence or "")
|
|
|
|
def test_complete_slsa_provenance_is_accepted(self) -> None:
|
|
payload = {
|
|
"_type": "https://in-toto.io/Statement/v1",
|
|
"subject": [{"name": "model.bin", "digest": {"sha256": "a" * 64}}],
|
|
"predicateType": "https://slsa.dev/provenance/v1",
|
|
"predicate": {
|
|
"buildDefinition": {
|
|
"buildType": "https://example.invalid/build/model/v1",
|
|
"externalParameters": {},
|
|
"resolvedDependencies": [{"uri": "git+https://example.invalid/repo@v1"}],
|
|
},
|
|
"runDetails": {"builder": {"id": "https://example.invalid/builders/research"}},
|
|
},
|
|
}
|
|
self.assertFalse(list(SlsaProvenanceRule().evaluate(json_artifact("provenance.json", payload, "provenance"))))
|
|
|
|
def test_non_hex_digest_is_not_accepted(self) -> None:
|
|
payload = {
|
|
"_type": "https://in-toto.io/Statement/v1",
|
|
"subject": [{"name": "model.bin", "digest": {"sha256": "z" * 64}}],
|
|
"predicateType": "https://slsa.dev/provenance/v1",
|
|
"predicate": {
|
|
"buildDefinition": {"buildType": "example", "externalParameters": {}, "resolvedDependencies": [{}]},
|
|
"runDetails": {"builder": {"id": "example"}},
|
|
},
|
|
}
|
|
finding = list(SlsaProvenanceRule().evaluate(json_artifact("provenance.json", payload, "provenance")))[0]
|
|
self.assertIn("subject SHA-256 digest", finding.evidence or "")
|
|
|
|
def test_ml_bom_requires_transparency_and_integrity(self) -> None:
|
|
payload = {"bomFormat": "CycloneDX", "specVersion": "1.7", "components": [{"type": "machine-learning-model", "name": "example"}]}
|
|
finding = list(CycloneDxMLBomRule().evaluate(json_artifact("ml-bom.json", payload, "ml-bom")))[0]
|
|
self.assertIn("modelCard", finding.evidence or "")
|
|
self.assertIn("SHA-256", finding.evidence or "")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|