agent-memory-guard 0.2.2__py3-none-any.whl

agent_memory_guard/__init__.py

@@ -0,0 +1,24 @@
+"""OWASP Agent Memory Guard — runtime defense against memory poisoning (ASI06)."""
+
+from agent_memory_guard.events import Action, SecurityEvent, Severity
+from agent_memory_guard.exceptions import (
+    IntegrityError,
+    MemoryGuardError,
+    PolicyViolation,
+)
+from agent_memory_guard.guard import MemoryGuard
+from agent_memory_guard.policies.policy import Policy
+
+__version__ = "0.2.2"
+
+__all__ = [
+    "MemoryGuard",
+    "Policy",
+    "SecurityEvent",
+    "Severity",
+    "Action",
+    "MemoryGuardError",
+    "PolicyViolation",
+    "IntegrityError",
+    "__version__",
+]

agent_memory_guard/detectors/__init__.py

@@ -0,0 +1,18 @@
+from agent_memory_guard.detectors.anomaly import (
+    RapidChangeDetector,
+    SizeAnomalyDetector,
+)
+from agent_memory_guard.detectors.base import DetectionResult, Detector
+from agent_memory_guard.detectors.injection import PromptInjectionDetector
+from agent_memory_guard.detectors.leakage import SensitiveDataDetector
+from agent_memory_guard.detectors.protected_keys import ProtectedKeyDetector
+
+__all__ = [
+    "Detector",
+    "DetectionResult",
+    "PromptInjectionDetector",
+    "SensitiveDataDetector",
+    "SizeAnomalyDetector",
+    "RapidChangeDetector",
+    "ProtectedKeyDetector",
+]

agent_memory_guard/detectors/anomaly.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import time
+from collections import deque
+from typing import Any
+
+from agent_memory_guard.detectors.base import DetectionResult
+from agent_memory_guard.detectors.injection import _stringify
+from agent_memory_guard.events import Severity
+
+
+class SizeAnomalyDetector:
+    """Flags memory writes that are unusually large or grow unusually fast."""
+
+    name = "size_anomaly"
+
+    def __init__(
+        self,
+        max_bytes: int = 64 * 1024,
+        growth_factor: float = 10.0,
+        severity: Severity = Severity.MEDIUM,
+    ) -> None:
+        self._max_bytes = max_bytes
+        self._growth_factor = growth_factor
+        self._last_size: dict[str, int] = {}
+        self._severity = severity
+
+    def inspect(self, key: str, value: Any, *, operation: str) -> DetectionResult:
+        size = len(_stringify(value).encode("utf-8"))
+        previous = self._last_size.get(key)
+        self._last_size[key] = size
+
+        if size > self._max_bytes:
+            return DetectionResult(
+                detector=self.name,
+                matched=True,
+                severity=self._severity,
+                message=f"Memory value for '{key}' exceeds size limit ({size} > {self._max_bytes} bytes)",
+                metadata={"size": size, "limit": self._max_bytes},
+            )
+
+        if previous and previous > 0 and size > previous * self._growth_factor:
+            return DetectionResult(
+                detector=self.name,
+                matched=True,
+                severity=self._severity,
+                message=f"Memory value for '{key}' grew {size / previous:.1f}x in one write",
+                metadata={"size": size, "previous": previous},
+            )
+
+        return DetectionResult(self.name, matched=False)
+
+
+class RapidChangeDetector:
+    """Flags suspiciously high write frequency on a single key (churn attack)."""
+
+    name = "rapid_change"
+
+    def __init__(
+        self,
+        window_seconds: float = 5.0,
+        max_writes: int = 20,
+        severity: Severity = Severity.MEDIUM,
+    ) -> None:
+        self._window = window_seconds
+        self._max = max_writes
+        self._writes: dict[str, deque[float]] = {}
+        self._severity = severity
+
+    def inspect(self, key: str, value: Any, *, operation: str) -> DetectionResult:
+        if operation != "write":
+            return DetectionResult(self.name, matched=False)
+
+        now = time.monotonic()
+        history = self._writes.setdefault(key, deque())
+        history.append(now)
+        cutoff = now - self._window
+        while history and history[0] < cutoff:
+            history.popleft()
+
+        if len(history) > self._max:
+            return DetectionResult(
+                detector=self.name,
+                matched=True,
+                severity=self._severity,
+                message=f"Rapid write churn on '{key}': {len(history)} writes in {self._window}s",
+                metadata={"writes": len(history), "window": self._window},
+            )
+        return DetectionResult(self.name, matched=False)

agent_memory_guard/detectors/base.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Protocol
+
+from agent_memory_guard.events import Severity
+
+
+@dataclass
+class DetectionResult:
+    """Verdict returned by a single detector."""
+
+    detector: str
+    matched: bool
+    severity: Severity = Severity.INFO
+    message: str = ""
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+class Detector(Protocol):
+    name: str
+
+    def inspect(self, key: str, value: Any, *, operation: str) -> DetectionResult: ...

agent_memory_guard/detectors/injection.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable
+from typing import Any
+
+from agent_memory_guard.detectors.base import DetectionResult
+from agent_memory_guard.events import Severity
+
+DEFAULT_INJECTION_PATTERNS: tuple[str, ...] = (
+    r"ignore (?:all |any |the )?(?:previous|prior|above) (?:instructions|messages|rules)",
+    r"disregard (?:all |any |the )?(?:previous|prior|above) (?:instructions|messages|rules)",
+    r"forget (?:all |any |the )?(?:previous|prior|above) (?:instructions|messages|rules)",
+    r"\byou are now\b.{0,40}(?:dan|jailbroken|admin|root|developer mode)",
+    r"\bsystem\s*[:\-]\s*you (?:are|must|will)",
+    r"</?\s*(?:system|assistant|tool)\s*>",
+    r"\bact as (?:an? )?(?:admin|root|system|developer|unrestricted)",
+    r"\b(?:reveal|print|leak|dump|exfiltrate)\s+(?:the\s+)?(?:system\s+)?(?:prompt|instructions|secrets|api[_\s-]?key)",
+    r"\bnew (?:instructions|directive|persona)\s*[:\-]",
+    r"\boverride\s+(?:safety|security|guardrails|policy)",
+)
+
+
+class PromptInjectionDetector:
+    """Regex-based screen for indirect prompt-injection markers in memory values."""
+
+    name = "prompt_injection"
+
+    def __init__(
+        self,
+        patterns: Iterable[str] = DEFAULT_INJECTION_PATTERNS,
+        severity: Severity = Severity.HIGH,
+    ) -> None:
+        self._patterns = [re.compile(p, re.IGNORECASE | re.DOTALL) for p in patterns]
+        self._severity = severity
+
+    def inspect(self, key: str, value: Any, *, operation: str) -> DetectionResult:
+        text = _stringify(value)
+        if not text:
+            return DetectionResult(self.name, matched=False)
+
+        hits: list[str] = []
+        for pattern in self._patterns:
+            match = pattern.search(text)
+            if match:
+                hits.append(match.group(0))
+
+        if not hits:
+            return DetectionResult(self.name, matched=False)
+
+        return DetectionResult(
+            detector=self.name,
+            matched=True,
+            severity=self._severity,
+            message=f"Possible prompt-injection markers in '{key}'",
+            metadata={"hits": hits[:5], "operation": operation},
+        )
+
+
+def _stringify(value: Any) -> str:
+    if value is None:
+        return ""
+    if isinstance(value, str):
+        return value
+    if isinstance(value, (list, tuple, set)):
+        return "\n".join(_stringify(v) for v in value)
+    if isinstance(value, dict):
+        return "\n".join(f"{k}: {_stringify(v)}" for k, v in value.items())
+    return str(value)

agent_memory_guard/detectors/leakage.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable
+from typing import Any
+
+from agent_memory_guard.detectors.base import DetectionResult
+from agent_memory_guard.detectors.injection import _stringify
+from agent_memory_guard.events import Severity
+
+DEFAULT_LEAKAGE_PATTERNS: dict[str, str] = {
+    "aws_access_key": r"\bAKIA[0-9A-Z]{16}\b",
+    "aws_secret_key": r"(?i)aws(.{0,20})?(secret|private)?[\s_-]?access[\s_-]?key[\s_-]?[:=][\s\"']*([A-Za-z0-9/+=]{40})",
+    "github_token": r"\bghp_[A-Za-z0-9]{36}\b",
+    "github_oauth": r"\bgho_[A-Za-z0-9]{36}\b",
+    "openai_key": r"\bsk-[A-Za-z0-9_-]{20,}\b",
+    "anthropic_key": r"\bsk-ant-[A-Za-z0-9_-]{20,}\b",
+    "google_api_key": r"\bAIza[0-9A-Za-z_-]{35}\b",
+    "slack_token": r"\bxox[abpr]-[0-9A-Za-z-]{10,}\b",
+    "private_key_pem": r"-----BEGIN (?:RSA |EC |OPENSSH |DSA )?PRIVATE KEY-----",
+    "jwt": r"\beyJ[A-Za-z0-9_-]{10,}\.eyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\b",
+    "credit_card": r"\b(?:\d[ -]*?){13,19}\b",
+    "ssn_us": r"\b\d{3}-\d{2}-\d{4}\b",
+    "email": r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b",
+}
+
+
+class SensitiveDataDetector:
+    """Flags secrets/PII present in memory values prior to write or after read."""
+
+    name = "sensitive_data"
+
+    def __init__(
+        self,
+        patterns: dict[str, str] | None = None,
+        ignore: Iterable[str] = ("email",),
+        severity: Severity = Severity.HIGH,
+    ) -> None:
+        merged = dict(DEFAULT_LEAKAGE_PATTERNS) if patterns is None else dict(patterns)
+        for name in ignore:
+            merged.pop(name, None)
+        self._patterns = {name: re.compile(p) for name, p in merged.items()}
+        self._severity = severity
+
+    def inspect(self, key: str, value: Any, *, operation: str) -> DetectionResult:
+        text = _stringify(value)
+        if not text:
+            return DetectionResult(self.name, matched=False)
+
+        findings: list[str] = []
+        for label, pattern in self._patterns.items():
+            if pattern.search(text):
+                findings.append(label)
+
+        if not findings:
+            return DetectionResult(self.name, matched=False)
+
+        return DetectionResult(
+            detector=self.name,
+            matched=True,
+            severity=self._severity,
+            message=f"Sensitive data ({', '.join(findings)}) detected in '{key}'",
+            metadata={"categories": findings, "operation": operation},
+        )
+
+    def redact(self, value: Any) -> Any:
+        text = _stringify(value)
+        for label, pattern in self._patterns.items():
+            text = pattern.sub(f"[REDACTED:{label}]", text)
+        return text

agent_memory_guard/detectors/protected_keys.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+import fnmatch
+from collections.abc import Iterable
+from typing import Any
+
+from agent_memory_guard.detectors.base import DetectionResult
+from agent_memory_guard.events import Severity
+
+
+class ProtectedKeyDetector:
+    """Flags writes targeting keys declared immutable by policy."""
+
+    name = "protected_key"
+
+    def __init__(
+        self,
+        protected: Iterable[str] = (),
+        severity: Severity = Severity.CRITICAL,
+    ) -> None:
+        self._patterns = list(protected)
+        self._severity = severity
+
+    def add(self, pattern: str) -> None:
+        self._patterns.append(pattern)
+
+    def matches(self, key: str) -> str | None:
+        for pattern in self._patterns:
+            if fnmatch.fnmatchcase(key, pattern):
+                return pattern
+        return None
+
+    def inspect(self, key: str, value: Any, *, operation: str) -> DetectionResult:
+        if operation != "write":
+            return DetectionResult(self.name, matched=False)
+        match = self.matches(key)
+        if not match:
+            return DetectionResult(self.name, matched=False)
+        return DetectionResult(
+            detector=self.name,
+            matched=True,
+            severity=self._severity,
+            message=f"Write to protected key '{key}' (matched pattern '{match}')",
+            metadata={"pattern": match},
+        )

agent_memory_guard/events.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+import time
+import uuid
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+
+class Severity(str, Enum):
+    INFO = "info"
+    LOW = "low"
+    MEDIUM = "medium"
+    HIGH = "high"
+    CRITICAL = "critical"
+
+
+class Action(str, Enum):
+    ALLOW = "allow"
+    REDACT = "redact"
+    BLOCK = "block"
+    QUARANTINE = "quarantine"
+
+
+@dataclass
+class SecurityEvent:
+    """Structured record of a guard decision, suitable for SIEM forwarding."""
+
+    detector: str
+    severity: Severity
+    action: Action
+    key: str
+    message: str
+    operation: str = "write"
+    metadata: dict[str, Any] = field(default_factory=dict)
+    timestamp: float = field(default_factory=time.time)
+    event_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "event_id": self.event_id,
+            "timestamp": self.timestamp,
+            "detector": self.detector,
+            "severity": self.severity.value,
+            "action": self.action.value,
+            "operation": self.operation,
+            "key": self.key,
+            "message": self.message,
+            "metadata": self.metadata,
+        }

agent_memory_guard/exceptions.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+
+class MemoryGuardError(Exception):
+    """Base exception for Agent Memory Guard."""
+
+
+class PolicyViolation(MemoryGuardError):
+    """Raised when a memory operation violates an enforcement policy."""
+
+    def __init__(self, message: str, rule: str | None = None, key: str | None = None):
+        super().__init__(message)
+        self.rule = rule
+        self.key = key
+
+
+class IntegrityError(MemoryGuardError):
+    """Raised when a memory entry fails its integrity baseline check."""
+
+    def __init__(self, message: str, key: str, expected: str, actual: str):
+        super().__init__(message)
+        self.key = key
+        self.expected = expected
+        self.actual = actual