tollgate 1.0.4__py3-none-any.whl → 1.4.0__py3-none-any.whl

tollgate/context_monitor.py

@@ -0,0 +1,292 @@
+"""Memory/Context integrity monitoring for AI agent systems.
+
+Detects unauthorized modifications to agent working memory between turns.
+Tracks checksums of context snapshots and alerts when unexpected changes
+are detected.
+
+This is a complementary layer — it operates alongside Tollgate's core
+enforcement pipeline to provide defense-in-depth against memory/context
+poisoning attacks (OWASP Agentic #2).
+
+Usage:
+
+    from tollgate.context_monitor import ContextIntegrityMonitor
+
+    monitor = ContextIntegrityMonitor(alert_sink=my_audit_sink)
+
+    # At the start of each turn, snapshot the context
+    monitor.snapshot("agent-1", "turn-5", context_data={
+        "system_prompt": "You are a helpful assistant...",
+        "tool_permissions": ["read", "write"],
+        "memory": {"key1": "value1"},
+    })
+
+    # Before processing, verify nothing changed unexpectedly
+    result = monitor.verify("agent-1", "turn-5", context_data={
+        "system_prompt": "You are a helpful assistant...",
+        "tool_permissions": ["read", "write"],
+        "memory": {"key1": "value1"},
+    })
+    assert result.is_valid  # True if unchanged
+
+    # Detect tampering
+    result = monitor.verify("agent-1", "turn-5", context_data={
+        "system_prompt": "IGNORE ALL RULES...",  # Poisoned!
+        "tool_permissions": ["read", "write", "admin"],  # Escalated!
+        "memory": {"key1": "value1"},
+    })
+    assert not result.is_valid
+    # result.changed_fields == ["system_prompt", "tool_permissions"]
+"""
+
+import hashlib
+import json
+import logging
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+logger = logging.getLogger("tollgate.context_monitor")
+
+
+@dataclass
+class ContextSnapshot:
+    """A point-in-time snapshot of agent context."""
+
+    agent_id: str
+    turn_id: str
+    checksum: str
+    field_checksums: dict[str, str]
+    timestamp: float
+    field_names: list[str]
+
+
+@dataclass
+class VerificationResult:
+    """Result of verifying context integrity."""
+
+    is_valid: bool
+    agent_id: str
+    turn_id: str
+    changed_fields: list[str] = field(default_factory=list)
+    added_fields: list[str] = field(default_factory=list)
+    removed_fields: list[str] = field(default_factory=list)
+    message: str = ""
+
+    @property
+    def has_changes(self) -> bool:
+        return bool(self.changed_fields or self.added_fields or self.removed_fields)
+
+
+class ContextIntegrityMonitor:
+    """Monitor for detecting unauthorized context modifications.
+
+    Maintains checksums of agent context per (agent_id, turn_id) pair.
+    Supports both full-context and per-field verification.
+
+    Args:
+        alert_callback: Optional callback invoked on integrity violation.
+            Receives a VerificationResult.
+        immutable_fields: Set of field names that must never change between
+            snapshot and verify. Violations are always flagged.
+        max_snapshots: Maximum number of snapshots to retain (per agent).
+            Older snapshots are evicted when the limit is reached.
+    """
+
+    def __init__(
+        self,
+        *,
+        alert_callback: Any | None = None,
+        immutable_fields: set[str] | None = None,
+        max_snapshots: int = 1000,
+    ):
+        self._alert_callback = alert_callback
+        self._immutable_fields = immutable_fields or {
+            "system_prompt",
+            "tool_permissions",
+            "security_level",
+            "role",
+        }
+        self._max_snapshots = max_snapshots
+        self._snapshots: dict[str, ContextSnapshot] = {}
+        self._agent_snapshot_keys: dict[str, list[str]] = {}
+
+    @staticmethod
+    def _compute_checksum(data: Any) -> str:
+        """Compute a deterministic SHA-256 checksum of arbitrary data."""
+        serialized = json.dumps(data, sort_keys=True, default=str)
+        return hashlib.sha256(serialized.encode()).hexdigest()
+
+    @staticmethod
+    def _compute_field_checksums(context_data: dict[str, Any]) -> dict[str, str]:
+        """Compute per-field checksums."""
+        checksums = {}
+        for key, value in context_data.items():
+            serialized = json.dumps(value, sort_keys=True, default=str)
+            checksums[key] = hashlib.sha256(serialized.encode()).hexdigest()
+        return checksums
+
+    def _snapshot_key(self, agent_id: str, turn_id: str) -> str:
+        return f"{agent_id}:{turn_id}"
+
+    def _evict_old_snapshots(self, agent_id: str):
+        """Evict oldest snapshots for an agent if over the limit."""
+        keys = self._agent_snapshot_keys.get(agent_id, [])
+        while len(keys) > self._max_snapshots:
+            old_key = keys.pop(0)
+            self._snapshots.pop(old_key, None)
+
+    def snapshot(
+        self,
+        agent_id: str,
+        turn_id: str,
+        context_data: dict[str, Any],
+    ) -> ContextSnapshot:
+        """Take a snapshot of the current context.
+
+        Args:
+            agent_id: The agent whose context is being snapshotted.
+            turn_id: A unique identifier for the current turn/step.
+            context_data: The context data to snapshot (dict of named fields).
+
+        Returns:
+            The created ContextSnapshot.
+        """
+        checksum = self._compute_checksum(context_data)
+        field_checksums = self._compute_field_checksums(context_data)
+
+        snap = ContextSnapshot(
+            agent_id=agent_id,
+            turn_id=turn_id,
+            checksum=checksum,
+            field_checksums=field_checksums,
+            timestamp=time.time(),
+            field_names=list(context_data.keys()),
+        )
+
+        key = self._snapshot_key(agent_id, turn_id)
+        self._snapshots[key] = snap
+
+        if agent_id not in self._agent_snapshot_keys:
+            self._agent_snapshot_keys[agent_id] = []
+        self._agent_snapshot_keys[agent_id].append(key)
+        self._evict_old_snapshots(agent_id)
+
+        return snap
+
+    def verify(
+        self,
+        agent_id: str,
+        turn_id: str,
+        context_data: dict[str, Any],
+    ) -> VerificationResult:
+        """Verify context integrity against a previous snapshot.
+
+        Args:
+            agent_id: The agent whose context is being verified.
+            turn_id: The turn_id used when the snapshot was taken.
+            context_data: The current context data to verify.
+
+        Returns:
+            VerificationResult with details of any changes detected.
+        """
+        key = self._snapshot_key(agent_id, turn_id)
+        snap = self._snapshots.get(key)
+
+        if snap is None:
+            return VerificationResult(
+                is_valid=True,
+                agent_id=agent_id,
+                turn_id=turn_id,
+                message="No snapshot found — nothing to verify against.",
+            )
+
+        # Quick check: full checksum
+        current_checksum = self._compute_checksum(context_data)
+        if current_checksum == snap.checksum:
+            return VerificationResult(
+                is_valid=True,
+                agent_id=agent_id,
+                turn_id=turn_id,
+                message="Context integrity verified.",
+            )
+
+        # Detailed check: per-field
+        current_field_checksums = self._compute_field_checksums(context_data)
+
+        changed_fields: list[str] = []
+        added_fields: list[str] = []
+        removed_fields: list[str] = []
+
+        # Check changed fields
+        for field_name in snap.field_checksums:
+            if field_name not in current_field_checksums:
+                removed_fields.append(field_name)
+            elif current_field_checksums[field_name] != snap.field_checksums[field_name]:
+                changed_fields.append(field_name)
+
+        # Check added fields
+        for field_name in current_field_checksums:
+            if field_name not in snap.field_checksums:
+                added_fields.append(field_name)
+
+        # Determine if immutable fields were violated
+        immutable_violations = [
+            f for f in changed_fields if f in self._immutable_fields
+        ] + [
+            f for f in removed_fields if f in self._immutable_fields
+        ]
+
+        is_valid = len(immutable_violations) == 0
+
+        message_parts = []
+        if changed_fields:
+            message_parts.append(f"Changed: {changed_fields}")
+        if added_fields:
+            message_parts.append(f"Added: {added_fields}")
+        if removed_fields:
+            message_parts.append(f"Removed: {removed_fields}")
+        if immutable_violations:
+            message_parts.append(f"IMMUTABLE VIOLATIONS: {immutable_violations}")
+
+        message = "; ".join(message_parts) if message_parts else "No changes."
+
+        result = VerificationResult(
+            is_valid=is_valid,
+            agent_id=agent_id,
+            turn_id=turn_id,
+            changed_fields=changed_fields,
+            added_fields=added_fields,
+            removed_fields=removed_fields,
+            message=message,
+        )
+
+        # Fire alert callback for violations
+        if not is_valid and self._alert_callback is not None:
+            try:
+                self._alert_callback(result)
+            except Exception:
+                logger.exception("Alert callback failed for context violation")
+
+        if not is_valid:
+            logger.warning(
+                "Context integrity violation for agent=%s turn=%s: %s",
+                agent_id, turn_id, message,
+            )
+
+        return result
+
+    def get_snapshot(self, agent_id: str, turn_id: str) -> ContextSnapshot | None:
+        """Retrieve a stored snapshot."""
+        key = self._snapshot_key(agent_id, turn_id)
+        return self._snapshots.get(key)
+
+    def clear(self, agent_id: str | None = None):
+        """Clear snapshots. If agent_id given, clear only that agent's snapshots."""
+        if agent_id is not None:
+            keys = self._agent_snapshot_keys.pop(agent_id, [])
+            for key in keys:
+                self._snapshots.pop(key, None)
+        else:
+            self._snapshots.clear()
+            self._agent_snapshot_keys.clear()

tollgate/exceptions.py

@@ -26,3 +26,23 @@ class TollgateDeferred(TollgateError):  # noqa: N818
     def __init__(self, approval_id: str):
         self.approval_id = approval_id
         super().__init__(f"Tool call deferred. Approval ID: {approval_id}")
+
+
+class TollgateRateLimited(TollgateError):  # noqa: N818
+    """Raised when a tool call is rejected due to rate limiting."""
+
+    def __init__(self, reason: str, retry_after: float | None = None):
+        self.reason = reason
+        self.retry_after = retry_after
+        msg = f"Rate limited: {reason}"
+        if retry_after is not None:
+            msg += f" (retry after {retry_after:.1f}s)"
+        super().__init__(msg)
+
+
+class TollgateConstraintViolation(TollgateError):  # noqa: N818
+    """Raised when tool parameters violate manifest constraints."""
+
+    def __init__(self, reason: str):
+        self.reason = reason
+        super().__init__(f"Constraint violation: {reason}")

tollgate/grants.py

@@ -1,9 +1,55 @@
 import asyncio
 import time
+from typing import Protocol, runtime_checkable
 
 from .types import AgentContext, Grant, ToolRequest
 
 
+@runtime_checkable
+class GrantStore(Protocol):
+    """Protocol for grant storage backends.
+
+    Implement this protocol to use a custom storage backend (Redis, SQLite, etc.).
+
+    Example Redis implementation:
+
+        class RedisGrantStore:
+            def __init__(self, redis_client):
+                self.redis = redis_client
+
+            async def create_grant(self, grant: Grant) -> str:
+                await self.redis.hset(f"grant:{grant.id}", mapping=grant.to_dict())
+                await self.redis.expireat(f"grant:{grant.id}", int(grant.expires_at))
+                return grant.id
+
+            async def find_matching_grant(self, agent_ctx, tool_request) -> Grant | None:
+                # Implement matching logic with Redis SCAN or secondary indexes
+                ...
+
+    All methods must be async. The InMemoryGrantStore serves as the reference implementation.
+    """
+
+    async def create_grant(self, grant: Grant) -> str:
+        ...
+
+    async def find_matching_grant(
+        self, agent_ctx: AgentContext, tool_request: ToolRequest
+    ) -> Grant | None:
+        ...
+
+    async def revoke_grant(self, grant_id: str) -> bool:
+        ...
+
+    async def list_active_grants(self, agent_id: str | None = None) -> list[Grant]:
+        ...
+
+    async def cleanup_expired(self) -> int:
+        ...
+
+    async def get_usage_count(self, grant_id: str) -> int:
+        ...
+
+
 class InMemoryGrantStore:
     """In-memory store for action grants with thread-safe matching logic."""
 

tollgate/manifest_signing.py

@@ -0,0 +1,90 @@
+"""Cryptographic manifest signing and verification.
+
+Uses HMAC-SHA256 for manifest integrity verification. This ensures that
+the manifest file has not been tampered with since it was signed by a
+trusted party.
+
+For production use with asymmetric keys (Ed25519), install the optional
+``cryptography`` dependency and use the Ed25519 variants.
+
+Usage (HMAC — zero dependencies):
+
+    # Sign a manifest (CI/build step):
+    from tollgate.manifest_signing import sign_manifest, verify_manifest
+
+    sign_manifest("manifest.yaml", secret_key=b"build-secret")
+    # Creates manifest.yaml.sig alongside the manifest
+
+    # Verify at load time:
+    valid = verify_manifest("manifest.yaml", secret_key=b"build-secret")
+    # Returns True if signature matches, False otherwise
+
+    # Use with ToolRegistry:
+    registry = ToolRegistry("manifest.yaml", signing_key=b"build-secret")
+    # Raises ValueError if signature is missing or invalid
+"""
+
+import hashlib
+import hmac
+from pathlib import Path
+
+
+def _compute_hmac(content: bytes, secret_key: bytes) -> str:
+    """Compute HMAC-SHA256 hex digest of content."""
+    return hmac.new(secret_key, content, hashlib.sha256).hexdigest()
+
+
+def sign_manifest(
+    manifest_path: str | Path, *, secret_key: bytes
+) -> Path:
+    """Sign a manifest file, writing the signature to ``<path>.sig``.
+
+    Args:
+        manifest_path: Path to the manifest YAML file.
+        secret_key: Shared secret key for HMAC-SHA256.
+
+    Returns:
+        Path to the created signature file.
+    """
+    manifest_path = Path(manifest_path)
+    if not manifest_path.exists():
+        raise FileNotFoundError(f"Manifest not found: {manifest_path}")
+
+    content = manifest_path.read_bytes()
+    signature = _compute_hmac(content, secret_key)
+
+    sig_path = manifest_path.with_suffix(manifest_path.suffix + ".sig")
+    sig_path.write_text(signature, encoding="utf-8")
+    return sig_path
+
+
+def verify_manifest(
+    manifest_path: str | Path, *, secret_key: bytes
+) -> bool:
+    """Verify a manifest file against its ``.sig`` signature file.
+
+    Args:
+        manifest_path: Path to the manifest YAML file.
+        secret_key: Shared secret key for HMAC-SHA256.
+
+    Returns:
+        True if the signature is valid, False otherwise.
+        Returns False if the signature file doesn't exist.
+    """
+    manifest_path = Path(manifest_path)
+    sig_path = manifest_path.with_suffix(manifest_path.suffix + ".sig")
+
+    if not manifest_path.exists() or not sig_path.exists():
+        return False
+
+    content = manifest_path.read_bytes()
+    expected = _compute_hmac(content, secret_key)
+
+    stored = sig_path.read_text(encoding="utf-8").strip()
+    return hmac.compare_digest(expected, stored)
+
+
+def get_manifest_hash(manifest_path: str | Path) -> str:
+    """Compute a SHA-256 content hash of the manifest (for audit trails)."""
+    content = Path(manifest_path).read_bytes()
+    return hashlib.sha256(content).hexdigest()

tollgate/network_guard.py

@@ -0,0 +1,114 @@
+"""Global network policy enforcement for AI agent tool calls.
+
+Provides a NetworkGuard that validates any URL-like parameters against
+a global allow/blocklist, independent of per-tool constraints in the
+manifest. This is the systematic solution for network-level security.
+
+Configuration is typically loaded from ``policy.yaml``:
+
+    network_policy:
+      default: deny       # "deny" or "allow"
+      allowlist:
+        - pattern: "https://api.github.com/*"
+        - pattern: "https://arxiv.org/*"
+      blocklist:
+        - pattern: "http://*"          # No plaintext HTTP
+        - pattern: "*.internal.*"      # No internal hosts
+      param_fields_to_check:           # Which param keys to inspect
+        - url
+        - endpoint
+        - target
+        - href
+        - uri
+"""
+
+import fnmatch
+from typing import Any
+
+
+class NetworkGuard:
+    """Global URL policy enforcement.
+
+    Inspects tool parameters for URL values and validates them against
+    allow/blocklists. Works alongside per-tool constraints in the manifest
+    (roadmap 1.4) to provide defense-in-depth.
+
+    Args:
+        default: "deny" (block unlisted URLs) or "allow" (permit unless blocked).
+        allowlist: List of dicts with ``pattern`` key (glob patterns).
+        blocklist: List of dicts with ``pattern`` key (glob patterns).
+        param_fields_to_check: List of parameter names to inspect for URLs.
+            If None, all string params starting with http(s):// are checked.
+    """
+
+    def __init__(
+        self,
+        *,
+        default: str = "deny",
+        allowlist: list[dict[str, str]] | None = None,
+        blocklist: list[dict[str, str]] | None = None,
+        param_fields_to_check: list[str] | None = None,
+    ):
+        if default not in ("deny", "allow"):
+            raise ValueError(f"default must be 'deny' or 'allow', got '{default}'")
+
+        self.default = default
+        self._allow_patterns = [e["pattern"] for e in (allowlist or []) if "pattern" in e]
+        self._block_patterns = [e["pattern"] for e in (blocklist or []) if "pattern" in e]
+        self._param_fields = set(param_fields_to_check) if param_fields_to_check else None
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> "NetworkGuard":
+        """Create a NetworkGuard from a policy.yaml ``network_policy`` dict."""
+        return cls(
+            default=config.get("default", "deny"),
+            allowlist=config.get("allowlist"),
+            blocklist=config.get("blocklist"),
+            param_fields_to_check=config.get("param_fields_to_check"),
+        )
+
+    def check(self, params: dict[str, Any]) -> list[str]:
+        """Check tool parameters against the network policy.
+
+        Returns a list of violation strings (empty = OK).
+        """
+        violations: list[str] = []
+
+        for key, value in params.items():
+            if not isinstance(value, str):
+                continue
+            if not (value.startswith("http://") or value.startswith("https://")):
+                continue
+
+            # If specific fields are configured, only check those
+            if self._param_fields is not None and key not in self._param_fields:
+                continue
+
+            # Check blocklist first (always wins)
+            for pattern in self._block_patterns:
+                if fnmatch.fnmatch(value, pattern):
+                    violations.append(
+                        f"Parameter '{key}': URL '{value}' blocked by "
+                        f"network policy (matches '{pattern}')"
+                    )
+                    break  # One block match is enough
+
+            # Check allowlist
+            if self._allow_patterns:
+                if any(fnmatch.fnmatch(value, p) for p in self._allow_patterns):
+                    continue  # Explicitly allowed
+
+                # Not in allowlist
+                if self.default == "deny":
+                    violations.append(
+                        f"Parameter '{key}': URL '{value}' not in "
+                        f"network policy allowlist"
+                    )
+            elif self.default == "deny":
+                # No allowlist defined + default deny = block all URLs
+                violations.append(
+                    f"Parameter '{key}': URL '{value}' blocked by "
+                    f"network policy (default: deny, no allowlist defined)"
+                )
+
+        return violations

tollgate/policy.py

@@ -22,6 +22,13 @@ class YamlPolicyEvaluator:
 
     # Security: Whitelist of allowed attributes for agent_ctx and intent matching
     ALLOWED_AGENT_ATTRS = frozenset({"agent_id", "version", "environment", "role"})
+    # Delegation-aware matching keys (checked separately from ALLOWED_AGENT_ATTRS)
+    DELEGATION_KEYS = frozenset({
+        "max_delegation_depth",
+        "deny_delegated",
+        "allowed_delegators",
+        "blocked_delegators",
+    })
     ALLOWED_INTENT_ATTRS = frozenset({"action", "reason", "session_id"})
 
     def __init__(
@@ -116,11 +123,41 @@ class YamlPolicyEvaluator:
         if "agent" in rule:
             for key, expected_val in rule["agent"].items():
                 # Security: Only allow whitelisted attributes
+                if key in self.DELEGATION_KEYS:
+                    continue  # Handled separately below
                 if key not in self.ALLOWED_AGENT_ATTRS:
                     continue
                 if getattr(agent_ctx, key, None) != expected_val:
                     return False
 
+        # Match delegation constraints (3.4)
+        if "agent" in rule:
+            agent_rule = rule["agent"]
+
+            # deny_delegated: true → block all delegated calls
+            if agent_rule.get("deny_delegated") and agent_ctx.is_delegated:
+                return False
+
+            # max_delegation_depth: N → block if chain is too deep
+            max_depth = agent_rule.get("max_delegation_depth")
+            if max_depth is not None and agent_ctx.delegation_depth > max_depth:
+                return False
+
+            # allowed_delegators: [...] → only delegated agents from these
+            # sources can match this rule. Non-delegated agents are excluded.
+            allowed = agent_rule.get("allowed_delegators")
+            if allowed is not None:
+                if not agent_ctx.is_delegated:
+                    return False  # Non-delegated agents don't match
+                if not any(d in allowed for d in agent_ctx.delegated_by):
+                    return False
+
+            # blocked_delegators: [...] → these agents cannot delegate
+            blocked = agent_rule.get("blocked_delegators")
+            if blocked is not None and agent_ctx.is_delegated:
+                if any(d in blocked for d in agent_ctx.delegated_by):
+                    return False
+
         # Match Intent (with attribute whitelist)
         if "intent" in rule:
             for key, expected_val in rule["intent"].items():