plug-play-memory 0.2.0__py3-none-any.whl

plug_play_memory/__init__.py

@@ -0,0 +1,23 @@
+"""plug-play-memory: self-hosted, brain-inspired memory + permissions for AI agents."""
+
+__version__ = "0.2.0"
+
+from .storage import MemoryStore, Fact
+from .audit import AuditLog, AuditEntry
+from .policy import PolicyEngine, PolicyDecision
+from .events import AgentEvent
+from .judge import Judge, JudgeDecision, RulesJudge, HeuristicEnricher, SalienceScorer
+from .emitter import EventEmitter
+from .consolidation import Consolidator
+from .auth import AuthConfig, AuthContext, AuthError, ALL_SCOPES
+
+__all__ = [
+    "MemoryStore", "Fact",
+    "AuditLog", "AuditEntry",
+    "PolicyEngine", "PolicyDecision",
+    "AgentEvent",
+    "Judge", "JudgeDecision", "RulesJudge", "HeuristicEnricher", "SalienceScorer",
+    "EventEmitter",
+    "Consolidator",
+    "AuthConfig", "AuthContext", "AuthError", "ALL_SCOPES",
+]

plug_play_memory/audit.py

@@ -0,0 +1,209 @@
+"""
+Audit log for plug-play-memory — security-hardened version.
+
+Changes from v0.1:
+  - HMAC-signed entries (C4 fix) — appended fake entries fail signature check
+  - External chain tip verification method (C4 fix)
+  - Constructor accepts optional signing_key (passed from MemoryStore)
+  - Triggers still block UPDATE/DELETE on rows
+"""
+
+from __future__ import annotations
+
+import hashlib
+import hmac
+import json
+import os
+import sqlite3
+import time
+from dataclasses import dataclass
+from typing import Any, Callable
+
+
+@dataclass
+class AuditEntry:
+    id: int
+    agent_id: str
+    op_type: str
+    detail: dict[str, Any]
+    allowed: bool | None
+    timestamp: float
+    prev_hash: str | None
+    entry_hash: str
+    signature: str | None  # new in v2
+
+
+class AuditLog:
+    def __init__(
+        self,
+        get_conn: Callable[[], sqlite3.Connection],
+        owner: Any = None,
+        signing_key: str | None = None,
+    ):
+        self._get_conn = get_conn
+        # If no key provided, derive one from the OS (random per process).
+        # This means signatures are unverifiable across restarts unless user
+        # provides a stable key — which is exactly the security property we want.
+        if signing_key:
+            self._signing_key = signing_key.encode()
+        else:
+            self._signing_key = os.urandom(32)
+        self._init_schema()
+
+    def _init_schema(self) -> None:
+        conn = self._get_conn()
+        conn.executescript(
+            """
+            CREATE TABLE IF NOT EXISTS audit_log (
+                id         INTEGER PRIMARY KEY AUTOINCREMENT,
+                agent_id   TEXT NOT NULL,
+                op_type    TEXT NOT NULL,
+                detail     TEXT NOT NULL,
+                allowed    INTEGER,
+                timestamp  REAL NOT NULL,
+                prev_hash  TEXT,
+                entry_hash TEXT NOT NULL,
+                signature  TEXT
+            );
+            CREATE INDEX IF NOT EXISTS idx_audit_agent_time
+                ON audit_log(agent_id, timestamp);
+            CREATE INDEX IF NOT EXISTS idx_audit_op
+                ON audit_log(op_type, timestamp);
+
+            CREATE TRIGGER IF NOT EXISTS audit_no_update
+            BEFORE UPDATE ON audit_log
+            BEGIN
+                SELECT RAISE(ABORT, 'audit_log is append-only');
+            END;
+            CREATE TRIGGER IF NOT EXISTS audit_no_delete
+            BEFORE DELETE ON audit_log
+            BEGIN
+                SELECT RAISE(ABORT, 'audit_log is append-only');
+            END;
+            """
+        )
+        conn.commit()
+
+    def _sign(self, payload: str) -> str:
+        return hmac.new(self._signing_key, payload.encode(), hashlib.sha256).hexdigest()
+
+    def log(
+        self,
+        agent_id: str,
+        op_type: str,
+        detail: dict[str, Any],
+        allowed: bool | None = None,
+    ) -> None:
+        conn = self._get_conn()
+        now = time.time()
+        prev = conn.execute(
+            "SELECT entry_hash FROM audit_log ORDER BY id DESC LIMIT 1"
+        ).fetchone()
+        prev_hash = prev["entry_hash"] if prev else None
+        # Build the payload that determines entry_hash
+        hash_payload = (
+            f"{agent_id}|{op_type}|{json.dumps(detail, sort_keys=True)}|"
+            f"{allowed}|{now}|{prev_hash}"
+        )
+        entry_hash = hashlib.blake2b(hash_payload.encode(), digest_size=16).hexdigest()
+        # Sign the entry_hash with HMAC (C4 fix)
+        signature = self._sign(entry_hash)
+
+        conn.execute(
+            """INSERT INTO audit_log
+               (agent_id, op_type, detail, allowed, timestamp, prev_hash, entry_hash, signature)
+               VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
+            (
+                agent_id, op_type, json.dumps(detail, sort_keys=True),
+                1 if allowed else (0 if allowed is False else None),
+                now, prev_hash, entry_hash, signature,
+            ),
+        )
+        conn.commit()
+
+    def recent(self, agent_id: str, limit: int = 50) -> list[AuditEntry]:
+        limit = max(1, min(limit, 1000))  # H2 fix
+        rows = self._get_conn().execute(
+            "SELECT * FROM audit_log WHERE agent_id = ? ORDER BY id DESC LIMIT ?",
+            (agent_id, limit),
+        ).fetchall()
+        return [self._row_to_entry(r) for r in rows]
+
+    def recent_all(self, limit: int = 100) -> list[AuditEntry]:
+        limit = max(1, min(limit, 1000))  # H2 fix
+        rows = self._get_conn().execute(
+            "SELECT * FROM audit_log ORDER BY id DESC LIMIT ?", (limit,)
+        ).fetchall()
+        return [self._row_to_entry(r) for r in rows]
+
+    def coverage_manifest(self, agent_id: str) -> dict[str, Any]:
+        rows = self._get_conn().execute(
+            """SELECT detail, allowed FROM audit_log
+               WHERE agent_id = ? AND op_type = 'permission_check'""",
+            (agent_id,),
+        ).fetchall()
+
+        by_action: dict[str, dict[str, int]] = {}
+        for r in rows:
+            action = json.loads(r["detail"]).get("action", "unknown")
+            bucket = by_action.setdefault(action, {"allowed": 0, "denied": 0})
+            bucket["allowed" if r["allowed"] else "denied"] += 1
+
+        return {
+            "agent_id": agent_id,
+            "actions_checked": list(by_action.keys()),
+            "counts_by_action": by_action,
+            "total_permission_checks": len(rows),
+            "note": "Reflects only actions actually attempted — not a safety claim.",
+        }
+
+    def verify_chain(self) -> dict[str, Any]:
+        """
+        Verify the integrity of the audit chain (C4 fix).
+
+        Returns a report with:
+          - total_entries: count
+          - valid_signatures: entries whose signature matches our HMAC key
+          - invalid_signatures: entries with bad signatures (likely tampered)
+          - chain_broken: entries whose prev_hash doesn't match previous entry_hash
+        """
+        rows = self._get_conn().execute(
+            "SELECT id, entry_hash, prev_hash, signature FROM audit_log ORDER BY id"
+        ).fetchall()
+
+        valid = 0
+        invalid_sigs = []
+        broken_chain = []
+        prev_hash = None
+
+        for r in rows:
+            # Check signature
+            expected_sig = self._sign(r["entry_hash"])
+            if r["signature"] and hmac.compare_digest(expected_sig, r["signature"]):
+                valid += 1
+            else:
+                invalid_sigs.append(r["id"])
+            # Check chain linkage
+            if prev_hash is not None and r["prev_hash"] != prev_hash:
+                broken_chain.append(r["id"])
+            prev_hash = r["entry_hash"]
+
+        return {
+            "total_entries": len(rows),
+            "valid_signatures": valid,
+            "invalid_signatures": invalid_sigs,
+            "chain_broken_at": broken_chain,
+            "integrity_ok": (not invalid_sigs and not broken_chain),
+        }
+
+    @staticmethod
+    def _row_to_entry(r: sqlite3.Row) -> AuditEntry:
+        allowed_raw = r["allowed"]
+        return AuditEntry(
+            id=r["id"], agent_id=r["agent_id"], op_type=r["op_type"],
+            detail=json.loads(r["detail"]),
+            allowed=bool(allowed_raw) if allowed_raw is not None else None,
+            timestamp=r["timestamp"],
+            prev_hash=r["prev_hash"], entry_hash=r["entry_hash"],
+            signature=r["signature"],
+        )

plug_play_memory/auth.py

@@ -0,0 +1,155 @@
+"""
+Auth layer for plug-play-memory MCP server — security-hardened.
+
+Changes from v0.1:
+  - Tokens stored as SHA-256 hashes (H4 fix) — plaintext tokens never persisted
+  - Config schema validation with friendly errors (was KeyError before)
+  - Reserved agent_id validation (M4 fix)
+"""
+
+from __future__ import annotations
+
+import hashlib
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+@dataclass
+class AuthContext:
+    """The identity of the caller, established at connection time."""
+    agent_id: str
+    agent_type: str
+    scopes: list[str]
+    token_id: str
+
+
+ALL_SCOPES = {
+    "memory:read",
+    "memory:write",
+    "memory:delete",
+    "permission:check",
+    "admin:consolidate",
+    "admin:audit",
+}
+
+
+class AuthError(Exception):
+    pass
+
+
+def _hash_token(token: str) -> str:
+    """SHA-256 hash of a token. Stored at rest instead of plaintext (H4 fix)."""
+    return hashlib.sha256(token.encode()).hexdigest()
+
+
+class AuthConfig:
+    """
+    Loads auth config from YAML. Tokens can be specified in two ways:
+
+    1. Plaintext (convenient for first-time setup, but stored as hash internally):
+        tokens:
+          - id: token-coder-1
+            token: sk-coder-1-abc123...
+
+    2. Pre-hashed (recommended for production):
+        tokens:
+          - id: token-coder-1
+            token_hash: 5f2b09c...  # sha256(token)
+
+    Either form works. Plaintext is auto-hashed on load.
+    """
+
+    def __init__(self, config_path: str | Path | None = None):
+        self.config_path = Path(config_path) if config_path else None
+        self.enabled = True
+        self.default_agent_id = "local"
+        self.default_agent_type = "coding"
+        self.default_scopes = ["memory:read", "memory:write", "permission:check"]
+        self._tokens: dict[str, dict[str, Any]] = {}  # token_hash -> {id, agent_id, agent_type, scopes}
+        if self.config_path and self.config_path.exists():
+            self._load()
+
+    def _load(self) -> None:
+        try:
+            spec = yaml.safe_load(self.config_path.read_text()) or {}
+        except yaml.YAMLError as e:
+            raise AuthError(f"Invalid YAML in config: {e}")
+        auth = spec.get("auth", {})
+        self.enabled = auth.get("enabled", True)
+        self.default_agent_id = auth.get("default_agent_id", "local")
+        self.default_agent_type = auth.get("default_agent_type", "coding")
+        self.default_scopes = auth.get("default_scopes",
+                                        ["memory:read", "memory:write", "permission:check"])
+
+        # Validate default agent_id isn't reserved (M4 fix)
+        from .storage import validate_agent_id, RESERVED_AGENT_IDS
+        try:
+            validate_agent_id(self.default_agent_id)
+        except ValueError as e:
+            raise AuthError(
+                f"config.yaml: default_agent_id '{self.default_agent_id}' is reserved. "
+                f"Reserved: {sorted(RESERVED_AGENT_IDS)}"
+            )
+
+        for i, t in enumerate(auth.get("tokens", []) or []):
+            # Schema validation (was missing in v0.1)
+            if "id" not in t:
+                raise AuthError(f"config.yaml: tokens[{i}] missing 'id' field")
+            if "agent_id" not in t:
+                raise AuthError(f"config.yaml: tokens[{i}] missing 'agent_id' field")
+            try:
+                validate_agent_id(t["agent_id"])
+            except ValueError as e:
+                raise AuthError(f"config.yaml: tokens[{i}].agent_id: {e}")
+
+            # Accept either plaintext 'token' or pre-hashed 'token_hash' (H4 fix)
+            if "token_hash" in t:
+                token_hash = t["token_hash"]
+            elif "token" in t:
+                token_hash = _hash_token(t["token"])
+            else:
+                raise AuthError(
+                    f"config.yaml: tokens[{i}] must have 'token' or 'token_hash'"
+                )
+            self._tokens[token_hash] = {
+                "id": t["id"],
+                "agent_id": t["agent_id"],
+                "agent_type": t.get("agent_type", "default"),
+                "scopes": t.get("scopes", self.default_scopes),
+            }
+
+    def authenticate(self, token: str | None = None) -> AuthContext:
+        if not self.enabled:
+            return AuthContext(
+                agent_id=self.default_agent_id,
+                agent_type=self.default_agent_type,
+                scopes=list(ALL_SCOPES),
+                token_id="disabled",
+            )
+        if not token:
+            return AuthContext(
+                agent_id=os.environ.get("AGENTMEMORY_AGENT_ID", self.default_agent_id),
+                agent_type=os.environ.get("AGENTMEMORY_AGENT_TYPE",
+                                          self.default_agent_type),
+                scopes=self.default_scopes,
+                token_id="default",
+            )
+        # Hash the incoming token and compare (H4 fix — constant-time via dict lookup)
+        token_hash = _hash_token(token)
+        if token_hash not in self._tokens:
+            raise AuthError("unknown token")
+        t = self._tokens[token_hash]
+        return AuthContext(
+            agent_id=t["agent_id"],
+            agent_type=t["agent_type"],
+            scopes=t["scopes"],
+            token_id=t["id"],
+        )
+
+    @staticmethod
+    def has_scope(ctx: AuthContext, scope: str) -> bool:
+        return scope in ctx.scopes