243 lines
9.7 KiB
Python
243 lines
9.7 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import re
|
|
from typing import Iterable, Protocol
|
|
|
|
from .models import Artifact, Finding, Severity
|
|
|
|
|
|
class Rule(Protocol):
|
|
rule_id: str
|
|
title: str
|
|
description: str
|
|
default_severity: Severity
|
|
|
|
def evaluate(self, artifact: Artifact) -> Iterable[Finding]: ...
|
|
|
|
|
|
def _location(text: str, offset: int) -> tuple[int, int]:
|
|
line = text.count("\n", 0, offset) + 1
|
|
last_break = text.rfind("\n", 0, offset)
|
|
return line, offset - last_break
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class RegexSignal:
|
|
name: str
|
|
pattern: re.Pattern[str]
|
|
severity: Severity
|
|
|
|
|
|
class SecretRule:
|
|
rule_id = "AAISA-SECRET-001"
|
|
title = "Potential embedded secret"
|
|
description = "Finds credential-shaped values and reports only a redacted signal."
|
|
default_severity = Severity.HIGH
|
|
artifact_kinds = ("text", "configuration", "source-code")
|
|
standards = ("CWE-798", "NIST AI 600-1")
|
|
references = (
|
|
"https://cwe.mitre.org/data/definitions/798.html",
|
|
"https://doi.org/10.6028/NIST.AI.600-1",
|
|
)
|
|
signals = (
|
|
RegexSignal("private key", re.compile(r"-----BEGIN (?:RSA |EC |OPENSSH )?PRIVATE KEY-----"), Severity.CRITICAL),
|
|
RegexSignal("AWS access key", re.compile(r"\bAKIA[0-9A-Z]{16}\b"), Severity.HIGH),
|
|
RegexSignal("GitHub token", re.compile(r"\bgh[pousr]_[A-Za-z0-9]{20,255}\b"), Severity.HIGH),
|
|
RegexSignal(
|
|
"credential assignment",
|
|
re.compile(r"(?im)^\s*(?:api[_-]?key|secret|token|password)\s*[:=]\s*['\"]?([^\s'\"#]{12,})"),
|
|
Severity.HIGH,
|
|
),
|
|
)
|
|
|
|
def evaluate(self, artifact: Artifact) -> Iterable[Finding]:
|
|
if artifact.text is None:
|
|
return ()
|
|
findings: list[Finding] = []
|
|
for signal in self.signals:
|
|
for match in signal.pattern.finditer(artifact.text):
|
|
line, column = _location(artifact.text, match.start())
|
|
findings.append(Finding(
|
|
self.rule_id,
|
|
self.title,
|
|
signal.severity,
|
|
f"Detected a {signal.name} pattern. The matched value is redacted.",
|
|
artifact.relative_path,
|
|
line,
|
|
column,
|
|
f"{signal.name}: [REDACTED]",
|
|
"Remove the value, rotate it if it may be live, and use an approved secret store.",
|
|
("credentials", "redaction"),
|
|
))
|
|
return findings
|
|
|
|
|
|
class PromptInjectionRule:
|
|
rule_id = "AAISA-PROMPT-001"
|
|
title = "Prompt injection or tool-exfiltration instruction"
|
|
description = "Finds instructions commonly used to override controls or disclose protected context."
|
|
default_severity = Severity.MEDIUM
|
|
artifact_kinds = ("prompt", "configuration", "document", "dataset-metadata")
|
|
standards = ("MITRE ATLAS AML.T0051", "OWASP Agentic AI Threats and Mitigations")
|
|
references = (
|
|
"https://atlas.mitre.org/techniques/AML.T0051",
|
|
"https://genai.owasp.org/resource/agentic-ai-threats-and-mitigations/",
|
|
)
|
|
pattern = re.compile(
|
|
r"(?i)\b(ignore|disregard|override)\b.{0,50}\b(previous|prior|system|developer)\b.{0,30}\b(instruction|prompt|message)s?\b"
|
|
r"|\b(reveal|print|exfiltrate|send)\b.{0,50}\b(system prompt|secret|credential|token)s?\b",
|
|
re.DOTALL,
|
|
)
|
|
|
|
def evaluate(self, artifact: Artifact) -> Iterable[Finding]:
|
|
if artifact.text is None or artifact.kind not in {"prompt", "configuration", "document", "dataset-metadata"}:
|
|
return ()
|
|
results: list[Finding] = []
|
|
for match in self.pattern.finditer(artifact.text):
|
|
line, column = _location(artifact.text, match.start())
|
|
results.append(Finding(
|
|
self.rule_id,
|
|
self.title,
|
|
Severity.MEDIUM,
|
|
"Artifact contains language associated with instruction override or protected-context disclosure.",
|
|
artifact.relative_path,
|
|
line,
|
|
column,
|
|
"Matched instruction category; surrounding content intentionally omitted.",
|
|
"Review the trust boundary, quote or neutralize untrusted instructions, and enforce tool authorization outside the prompt.",
|
|
("prompt-security", "human-review"),
|
|
))
|
|
return results
|
|
|
|
|
|
class DependencyPinRule:
|
|
rule_id = "AAISA-DEPENDENCY-001"
|
|
title = "Unpinned Python dependency"
|
|
description = "Finds requirements entries without an exact version or immutable reference."
|
|
default_severity = Severity.MEDIUM
|
|
artifact_kinds = ("dependency-manifest",)
|
|
standards = ("SLSA",)
|
|
references = ("https://slsa.dev/spec/v1.2/",)
|
|
|
|
def evaluate(self, artifact: Artifact) -> Iterable[Finding]:
|
|
if artifact.text is None or artifact.path.name.lower() != "requirements.txt":
|
|
return ()
|
|
results: list[Finding] = []
|
|
for number, raw in enumerate(artifact.text.splitlines(), 1):
|
|
line = raw.strip()
|
|
if not line or line.startswith(("#", "-r", "--", "git+")):
|
|
continue
|
|
if "==" not in line or any(token in line for token in (">=", "<=", "~=", "!=", "*")):
|
|
package = re.split(r"[<>=!~\s\[]", line, maxsplit=1)[0]
|
|
results.append(Finding(
|
|
self.rule_id,
|
|
self.title,
|
|
Severity.MEDIUM,
|
|
f"Dependency {package!r} is not pinned to one exact version.",
|
|
artifact.relative_path,
|
|
number,
|
|
1,
|
|
f"requirement: {package or '[unparsed]'}",
|
|
"Pin the dependency and preserve a reviewed lock file or hash set.",
|
|
("supply-chain", "dependencies"),
|
|
))
|
|
return results
|
|
|
|
|
|
class ContainerTagRule:
|
|
rule_id = "AAISA-CONTAINER-001"
|
|
title = "Mutable container image reference"
|
|
description = "Finds latest or untagged images in Dockerfiles and Compose files."
|
|
default_severity = Severity.MEDIUM
|
|
artifact_kinds = ("deployment",)
|
|
standards = ("SLSA",)
|
|
references = ("https://slsa.dev/spec/v1.2/",)
|
|
from_pattern = re.compile(r"(?im)^\s*FROM\s+([^\s]+)")
|
|
compose_pattern = re.compile(r"(?im)^\s*image:\s*([^\s#]+)")
|
|
|
|
def evaluate(self, artifact: Artifact) -> Iterable[Finding]:
|
|
if artifact.text is None or artifact.kind != "deployment":
|
|
return ()
|
|
findings: list[Finding] = []
|
|
for pattern in (self.from_pattern, self.compose_pattern):
|
|
for match in pattern.finditer(artifact.text):
|
|
image = match.group(1)
|
|
base = image.split("@", 1)[0]
|
|
last_segment = base.rsplit("/", 1)[-1]
|
|
mutable = "@sha256:" not in image and (":" not in last_segment or last_segment.endswith(":latest"))
|
|
if mutable:
|
|
line, column = _location(artifact.text, match.start(1))
|
|
findings.append(Finding(
|
|
self.rule_id,
|
|
self.title,
|
|
Severity.MEDIUM,
|
|
f"Container image {image!r} is untagged or uses the mutable latest tag.",
|
|
artifact.relative_path,
|
|
line,
|
|
column,
|
|
f"image: {image}",
|
|
"Use a reviewed immutable digest and document the update process.",
|
|
("supply-chain", "containers"),
|
|
))
|
|
return findings
|
|
|
|
|
|
class ProvenanceRule:
|
|
rule_id = "AAISA-PROVENANCE-001"
|
|
title = "Incomplete AI artifact provenance"
|
|
description = "Checks model and dataset metadata for source, version, license, and SHA-256 evidence."
|
|
default_severity = Severity.MEDIUM
|
|
artifact_kinds = ("model-metadata", "dataset-metadata")
|
|
standards = ("NIST AI 600-1", "CycloneDX ML-BOM")
|
|
references = (
|
|
"https://doi.org/10.6028/NIST.AI.600-1",
|
|
"https://cyclonedx.org/guides/OWASP_CycloneDX-Authoritative-Guide-to-AI-ML-BOM-en.pdf",
|
|
)
|
|
required = ("source", "version", "license", "sha256")
|
|
|
|
def evaluate(self, artifact: Artifact) -> Iterable[Finding]:
|
|
if artifact.text is None or artifact.kind not in {"model-metadata", "dataset-metadata"}:
|
|
return ()
|
|
lowered = artifact.text.lower()
|
|
missing = [field for field in self.required if not re.search(rf"(?m)^\s*[\"']?{field}[\"']?\s*[:=]", lowered)]
|
|
if not missing:
|
|
return ()
|
|
return (Finding(
|
|
self.rule_id,
|
|
self.title,
|
|
Severity.MEDIUM,
|
|
"AI artifact metadata is missing provenance fields: " + ", ".join(missing) + ".",
|
|
artifact.relative_path,
|
|
1,
|
|
1,
|
|
"missing fields: " + ", ".join(missing),
|
|
"Record the authorized source, immutable version, license, and verified SHA-256 digest.",
|
|
("provenance", artifact.kind),
|
|
),)
|
|
|
|
|
|
from .structured_rules import AgentToolPermissionRule, CycloneDxMLBomRule, SlsaProvenanceRule
|
|
|
|
|
|
BUILTIN_RULES: tuple[Rule, ...] = (
|
|
SecretRule(),
|
|
PromptInjectionRule(),
|
|
DependencyPinRule(),
|
|
ContainerTagRule(),
|
|
ProvenanceRule(),
|
|
AgentToolPermissionRule(),
|
|
SlsaProvenanceRule(),
|
|
CycloneDxMLBomRule(),
|
|
)
|
|
|
|
|
|
def select_rules(enabled: tuple[str, ...] | None) -> tuple[Rule, ...]:
|
|
if enabled is None:
|
|
return BUILTIN_RULES
|
|
known = {rule.rule_id: rule for rule in BUILTIN_RULES}
|
|
unknown = sorted(set(enabled) - known.keys())
|
|
if unknown:
|
|
raise ValueError("unknown enabled rule(s): " + ", ".join(unknown))
|
|
return tuple(known[rule_id] for rule_id in enabled)
|