agent-tool-firewall 0.2.0__py3-none-any.whl

agent_firewall/__init__.py

@@ -0,0 +1,47 @@
+"""Agent Firewall — human-in-the-loop tool call approval."""
+
+from agent_firewall.context import set_current_user
+from agent_firewall.decorator import guarded_tool
+from agent_firewall.policy_engine import PolicyEngine
+from agent_firewall.server.app import create_app
+
+# Guard modules — injection detection, PII masking, groundedness checking.
+from agent_firewall.guards.injection_detector import InjectionScore, score_chunk
+from agent_firewall.guards.pii_masker import (
+    PIIMasker,
+    mask_pii,
+    unmask_pii,
+    reset_pii_masker,
+    get_pii_mapping,
+)
+from agent_firewall.guards.groundedness_checker import (
+    ClaimResult,
+    GroundednessReport,
+    check_groundedness,
+    format_grounded_answer,
+)
+from agent_firewall.guards.pipeline import GuardPipeline, GuardResult
+
+__all__ = [
+    "guarded_tool",
+    "set_current_user",
+    "PolicyEngine",
+    "create_app",
+    # Guards
+    "InjectionScore",
+    "score_chunk",
+    "PIIMasker",
+    "mask_pii",
+    "unmask_pii",
+    "reset_pii_masker",
+    "get_pii_mapping",
+    "ClaimResult",
+    "GroundednessReport",
+    "check_groundedness",
+    "format_grounded_answer",
+    "GuardPipeline",
+    "GuardResult",
+]
+
+__version__ = "0.2.0"
+

agent_firewall/cli.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+import shutil
+
+import uvicorn
+
+from agent_firewall.config import get_settings
+
+
+def _example_policy_path() -> Path | None:
+    """Resolve firewall.example.yaml for wheel installs and local checkouts."""
+    candidates = [
+        Path(__file__).resolve().parent / "data" / "firewall.example.yaml",
+        Path.cwd() / "firewall.example.yaml",
+        Path(__file__).resolve().parents[1] / "firewall.example.yaml",
+    ]
+    return next((p for p in candidates if p.exists()), None)
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(prog="agent-firewall", description="Agent Firewall CLI")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    serve = sub.add_parser("serve", help="Start API + approval dashboard")
+    serve.add_argument("--host", default="127.0.0.1")
+    serve.add_argument("--port", type=int, default=8000)
+    serve.add_argument("--db", default=None, help="SQLite DB path")
+    serve.add_argument("--policy", default=None, help="firewall.yaml path")
+    serve.add_argument("--reload", action="store_true")
+
+    init = sub.add_parser("init-policy", help="Copy firewall.example.yaml to ./firewall.yaml")
+    init.add_argument("--force", action="store_true")
+
+    args = parser.parse_args(argv)
+
+    if args.command == "init-policy":
+        example = _example_policy_path()
+        if example is None:
+            raise SystemExit("firewall.example.yaml not found")
+        dest = Path.cwd() / "firewall.yaml"
+        if dest.exists() and not args.force:
+            raise SystemExit(f"{dest} already exists (use --force to overwrite)")
+        shutil.copy(example, dest)
+        print(f"Wrote {dest}")
+        return
+
+    if args.command == "serve":
+        settings = get_settings(db_path=args.db, policy_path=args.policy)
+        if args.db:
+            import os
+
+            os.environ["FIREWALL_DB_PATH"] = args.db
+        if args.policy:
+            import os
+
+            os.environ["FIREWALL_POLICY_PATH"] = args.policy
+
+        policy_path = Path(settings.policy_path)
+        if not policy_path.exists():
+            print(
+                f"Warning: policy file not found at {policy_path.resolve()}. "
+                f"Run `agent-firewall init-policy` or pass --policy."
+            )
+
+        uvicorn.run(
+            "agent_firewall.server.app:create_app",
+            factory=True,
+            host=args.host,
+            port=args.port,
+            reload=args.reload,
+        )
+
+
+if __name__ == "__main__":
+    main()

agent_firewall/client.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+from typing import Any
+
+import requests
+
+from agent_firewall.config import get_settings
+
+
+class FirewallClientError(RuntimeError):
+    pass
+
+
+class FirewallClient:
+    def __init__(self, base_url: str | None = None, timeout: float = 10.0):
+        settings = get_settings(firewall_url=base_url)
+        self.base_url = settings.firewall_url
+        self.timeout = timeout
+
+    def _url(self, path: str) -> str:
+        return f"{self.base_url}{path}"
+
+    def _request(self, method: str, path: str, **kwargs) -> Any:
+        try:
+            resp = requests.request(method, self._url(path), timeout=self.timeout, **kwargs)
+        except requests.RequestException as exc:
+            raise FirewallClientError(
+                f"Agent Firewall backend unreachable at {self.base_url}: {exc}"
+            ) from exc
+        if resp.status_code >= 400:
+            raise FirewallClientError(
+                f"{method} {path} failed ({resp.status_code}): {resp.text}"
+            )
+        if resp.content:
+            return resp.json()
+        return None
+
+    def create_pending(
+        self,
+        *,
+        tool_name: str,
+        payload: dict[str, Any],
+        risk_level: str,
+        matched_rule: str,
+        requested_by: str,
+        timeout_seconds: int,
+    ) -> int:
+        data = self._request(
+            "POST",
+            "/api/pending",
+            json={
+                "tool_name": tool_name,
+                "payload": payload,
+                "risk_level": risk_level,
+                "matched_rule": matched_rule,
+                "requested_by": requested_by,
+                "timeout_seconds": timeout_seconds,
+            },
+        )
+        return int(data["id"])
+
+    def get_pending(self, pending_id: int) -> dict[str, Any]:
+        return self._request("GET", f"/api/pending/{pending_id}")
+
+    def mark_executed(self, pending_id: int, error_message: str | None = None) -> dict[str, Any]:
+        return self._request(
+            "POST",
+            f"/api/pending/{pending_id}/executed",
+            json={"error_message": error_message},
+        )
+
+    def log_audit(
+        self,
+        *,
+        tool_name: str,
+        payload: dict[str, Any],
+        decision: str,
+        requested_by: str,
+        matched_rule: str,
+        pending_action_id: int | None = None,
+        decided_by: str | None = None,
+    ) -> None:
+        self._request(
+            "POST",
+            "/api/audit",
+            json={
+                "tool_name": tool_name,
+                "payload": payload,
+                "decision": decision,
+                "requested_by": requested_by,
+                "matched_rule": matched_rule,
+                "pending_action_id": pending_action_id,
+                "decided_by": decided_by,
+            },
+        )

agent_firewall/config.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+@dataclass(frozen=True)
+class Settings:
+    firewall_url: str = "http://127.0.0.1:8000"
+    db_path: str = "./firewall.db"
+    policy_path: str = "./firewall.yaml"
+    auth_mode: str = "local"  # local | passthrough
+    session_secret: str = "dev-secret-change-me"
+    passthrough_header: str = "X-Forwarded-User"
+    poll_interval_seconds: float = 1.5
+
+
+def get_settings(
+    db_path: str | None = None,
+    policy_path: str | None = None,
+    firewall_url: str | None = None,
+) -> Settings:
+    return Settings(
+        firewall_url=(firewall_url or os.getenv("FIREWALL_URL", "http://127.0.0.1:8000")).rstrip(
+            "/"
+        ),
+        db_path=db_path or os.getenv("FIREWALL_DB_PATH", "./firewall.db"),
+        policy_path=policy_path or os.getenv("FIREWALL_POLICY_PATH", "./firewall.yaml"),
+        auth_mode=os.getenv("FIREWALL_AUTH_MODE", "local").lower(),
+        session_secret=os.getenv("FIREWALL_SESSION_SECRET", "dev-secret-change-me"),
+        passthrough_header=os.getenv("FIREWALL_PASSTHROUGH_HEADER", "X-Forwarded-User"),
+        poll_interval_seconds=float(os.getenv("FIREWALL_POLL_INTERVAL", "1.5")),
+    )
+
+
+def resolve_path(path: str) -> Path:
+    return Path(path).expanduser().resolve()

agent_firewall/context.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from contextvars import ContextVar
+
+current_user: ContextVar[str] = ContextVar("current_user", default="unknown")
+
+
+def set_current_user(user_id: str) -> None:
+    """Set the requesting user identity for subsequent guarded tool calls."""
+    current_user.set(user_id)
+
+
+def get_current_user() -> str:
+    return current_user.get()

agent_firewall/data/firewall.example.yaml

@@ -0,0 +1,64 @@
+# Copy to firewall.yaml and adjust for your project.
+# Demo users (local auth). Change passwords before any real deployment.
+
+users:
+  - username: admin
+    password: admin123
+    display_name: Admin Approver
+  - username: alice
+    password: alice123
+    display_name: Alice
+
+rules:
+  - name: high-value-payments
+    match:
+      tool_name: send_payment
+      payload_contains:
+        field: amount
+        operator: ">"
+        value: 1000
+    action: require_approval
+    risk_level: high
+    timeout_seconds: 120
+
+  - name: low-value-payments
+    match:
+      tool_name: send_payment
+      payload_contains:
+        field: amount
+        operator: "<="
+        value: 1000
+    action: allow
+    risk_level: low
+
+  - name: any-deletion
+    match:
+      tool_name: "delete_*"
+    action: require_approval
+    risk_level: high
+    timeout_seconds: 90
+
+  - name: block-transfer-all
+    match:
+      tool_name: transfer_all_funds
+    action: block
+    risk_level: high
+
+# ---------------------------------------------------------------------------
+# Guard pipeline — injection detection and PII scanning.
+# These run BEFORE policy rules and can escalate the action (never downgrade).
+# Groundedness is a standalone post-answer API (check_groundedness /
+# format_grounded_answer), not part of this pre-tool pipeline.
+# ---------------------------------------------------------------------------
+guards:
+  injection_detection:
+    enabled: true
+    block_threshold: 0.8        # score >= 0.8 → auto-block
+    approval_threshold: 0.5     # score >= 0.5 → require approval
+    scan_fields: null           # null = scan all string fields
+
+  pii_policy:
+    enabled: true
+    action: require_approval    # allow | require_approval | block
+    scan_fields: null           # null = scan all string fields
+    types: [EMAIL, PHONE, CREDIT_CARD, SSN, IPV4]

agent_firewall/db/__init__.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from agent_firewall.db.models import make_session_factory
+
+__all__ = ["make_session_factory"]

agent_firewall/db/models.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+
+from sqlalchemy import DateTime, ForeignKey, Integer, String, Text, create_engine
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, sessionmaker
+
+
+def utcnow() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+class User(Base):
+    __tablename__ = "users"
+
+    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
+    username: Mapped[str] = mapped_column(String(128), unique=True, nullable=False)
+    password_hash: Mapped[str] = mapped_column(String(255), nullable=False)
+    display_name: Mapped[str] = mapped_column(String(255), nullable=False)
+
+
+class PendingAction(Base):
+    __tablename__ = "pending_actions"
+
+    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
+    tool_name: Mapped[str] = mapped_column(String(255), nullable=False)
+    payload_json: Mapped[str] = mapped_column(Text, nullable=False)
+    risk_level: Mapped[str] = mapped_column(String(32), nullable=False, default="medium")
+    matched_rule: Mapped[str] = mapped_column(String(255), nullable=False, default="")
+    status: Mapped[str] = mapped_column(String(32), nullable=False, default="pending")
+    requested_by: Mapped[str] = mapped_column(String(255), nullable=False, default="unknown")
+    requested_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+    decided_by: Mapped[str | None] = mapped_column(String(255), nullable=True)
+    decided_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
+    reason: Mapped[str | None] = mapped_column(Text, nullable=True)
+    timeout_seconds: Mapped[int] = mapped_column(Integer, nullable=False, default=120)
+    executed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
+    error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
+
+
+class AuditLog(Base):
+    __tablename__ = "audit_log"
+
+    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
+    pending_action_id: Mapped[int | None] = mapped_column(
+        Integer, ForeignKey("pending_actions.id"), nullable=True
+    )
+    tool_name: Mapped[str] = mapped_column(String(255), nullable=False)
+    payload_json: Mapped[str] = mapped_column(Text, nullable=False)
+    decision: Mapped[str] = mapped_column(String(32), nullable=False)
+    requested_by: Mapped[str] = mapped_column(String(255), nullable=False, default="unknown")
+    decided_by: Mapped[str | None] = mapped_column(String(255), nullable=True)
+    timestamp: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+    matched_rule: Mapped[str] = mapped_column(String(255), nullable=False, default="")
+    reason: Mapped[str | None] = mapped_column(Text, nullable=True)
+
+
+def make_engine(db_path: str):
+    # SQLite needs absolute-ish URL; keep simple file path.
+    url = f"sqlite:///{db_path}"
+    return create_engine(url, connect_args={"check_same_thread": False})
+
+
+def make_session_factory(db_path: str):
+    engine = make_engine(db_path)
+    return engine, sessionmaker(bind=engine, autoflush=False, autocommit=False)
+
+
+def _ensure_sqlite_columns(engine) -> None:
+    """Add columns introduced after first install (SQLite has no migrations)."""
+    if engine.dialect.name != "sqlite":
+        return
+    with engine.begin() as conn:
+        rows = conn.exec_driver_sql("PRAGMA table_info(audit_log)").fetchall()
+        columns = {row[1] for row in rows}
+        if rows and "reason" not in columns:
+            conn.exec_driver_sql("ALTER TABLE audit_log ADD COLUMN reason TEXT")
+
+
+def init_db(engine) -> None:
+    Base.metadata.create_all(bind=engine)
+    _ensure_sqlite_columns(engine)

agent_firewall/decorator.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+import json
+import logging
+from collections.abc import Callable
+from functools import wraps
+from typing import Any, TypeVar
+
+from agent_firewall.client import FirewallClient
+from agent_firewall.config import get_settings
+from agent_firewall.context import get_current_user
+from agent_firewall.models import PolicyAction
+from agent_firewall.policy_engine import PolicyEngine
+from agent_firewall.waiter import start_approval_waiter
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+logger = logging.getLogger(__name__)
+
+_policy_engine: PolicyEngine | None = None
+_guard_pipeline: Any = None  # lazy import to avoid circular imports
+
+
+def _get_policy_engine() -> PolicyEngine:
+    global _policy_engine
+    if _policy_engine is None:
+        settings = get_settings()
+        _policy_engine = PolicyEngine(settings.policy_path)
+    return _policy_engine
+
+
+def _get_guard_pipeline():
+    """Build or return the cached ``GuardPipeline`` from the policy YAML ``guards:`` section."""
+    global _guard_pipeline
+    if _guard_pipeline is None:
+        from agent_firewall.guards.pipeline import GuardPipeline
+
+        engine = _get_policy_engine()
+        guards_config = engine.guards_config
+        if guards_config:
+            _guard_pipeline = GuardPipeline.from_policy_config(guards_config)
+        else:
+            # No guards configured — return a no-op pipeline.
+            _guard_pipeline = GuardPipeline()
+    return _guard_pipeline
+
+
+def reload_policy() -> None:
+    global _guard_pipeline
+    _get_policy_engine().reload()
+    _guard_pipeline = None  # Force rebuild on next call.
+
+
+def guarded_tool(
+    require_approval: bool | None = None,
+    risk_level_hint: str | None = None,
+    scan_injection: bool | None = None,
+    scan_pii: bool | None = None,
+) -> Callable[[F], F]:
+    """
+    Intercept tool calls and apply YAML policy (+ optional decorator override).
+
+    Non-blocking on require_approval: returns a dashboard message immediately and
+    executes the real tool in a background waiter after approve.
+
+    Parameters
+    ----------
+    require_approval:
+        Force HITL regardless of YAML policy.
+    risk_level_hint:
+        Override the risk level for this tool.
+    scan_injection:
+        Override the YAML ``guards.injection_detection.enabled`` setting for
+        this specific tool.  ``True`` to force scanning, ``False`` to skip.
+    scan_pii:
+        Override the YAML ``guards.pii_policy.enabled`` setting for this
+        specific tool.  ``True`` to force scanning, ``False`` to skip.
+    """
+
+    def decorator(fn: F) -> F:
+        @wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            engine = _get_policy_engine()
+            # Prefer kwargs as payload (typical tool calling). Merge positional by name if needed.
+            payload = dict(kwargs)
+            decision = engine.evaluate(fn.__name__, payload)
+
+            if require_approval is True:
+                from agent_firewall.models import PolicyDecision
+
+                decision = PolicyDecision(
+                    action=PolicyAction.REQUIRE_APPROVAL,
+                    risk_level=risk_level_hint or decision.risk_level or "high",
+                    matched_rule=f"decorator-override:{decision.matched_rule}",
+                    timeout_seconds=decision.timeout_seconds or 120,
+                )
+
+            # --- Run guard pipeline (injection + PII) ---
+            guard_details: dict[str, Any] = {}
+            try:
+                pipeline = _get_guard_pipeline()
+
+                # Apply per-decorator overrides.
+                if scan_injection is not None:
+                    pipeline.injection_enabled = scan_injection
+                if scan_pii is not None:
+                    pipeline.pii_enabled = scan_pii
+
+                guard_result = pipeline.evaluate(fn.__name__, payload)
+                guard_details = guard_result.guard_details
+
+                # Guards can escalate the action but never downgrade it.
+                if guard_result.action == PolicyAction.BLOCK and decision.action != PolicyAction.BLOCK:
+                    from agent_firewall.models import PolicyDecision
+
+                    decision = PolicyDecision(
+                        action=PolicyAction.BLOCK,
+                        risk_level="critical",
+                        matched_rule=f"guard:{guard_result.reason}",
+                        timeout_seconds=decision.timeout_seconds,
+                    )
+                elif (
+                    guard_result.action == PolicyAction.REQUIRE_APPROVAL
+                    and decision.action == PolicyAction.ALLOW
+                ):
+                    from agent_firewall.models import PolicyDecision
+
+                    decision = PolicyDecision(
+                        action=PolicyAction.REQUIRE_APPROVAL,
+                        risk_level=risk_level_hint or "high",
+                        matched_rule=f"guard:{guard_result.reason}",
+                        timeout_seconds=decision.timeout_seconds or 120,
+                    )
+            except Exception:  # noqa: BLE001
+                logger.warning("Guard pipeline error for %s — proceeding with policy only", fn.__name__, exc_info=True)
+
+            client = FirewallClient()
+            requested_by = get_current_user()
+
+            # Merge guard scan details into the audit payload so the dashboard can show them.
+            audit_extra: dict[str, Any] = {}
+            if guard_details:
+                audit_extra["guard_scan"] = guard_details
+
+            if decision.action == PolicyAction.BLOCK:
+                client.log_audit(
+                    tool_name=fn.__name__,
+                    payload={**payload, **audit_extra},
+                    decision="auto_block",
+                    requested_by=requested_by,
+                    matched_rule=decision.matched_rule,
+                )
+                return (
+                    f"Blocked by Agent Firewall policy rule '{decision.matched_rule}'. "
+                    f"Tool '{fn.__name__}' was not executed."
+                )
+
+            if decision.action == PolicyAction.ALLOW:
+                client.log_audit(
+                    tool_name=fn.__name__,
+                    payload={**payload, **audit_extra},
+                    decision="auto_allow",
+                    requested_by=requested_by,
+                    matched_rule=decision.matched_rule,
+                )
+                return fn(*args, **kwargs)
+
+            # require_approval
+            pending_id = client.create_pending(
+                tool_name=fn.__name__,
+                payload={**payload, **audit_extra},
+                risk_level=risk_level_hint or decision.risk_level,
+                matched_rule=decision.matched_rule,
+                requested_by=requested_by,
+                timeout_seconds=decision.timeout_seconds or 120,
+            )
+            start_approval_waiter(
+                pending_id=pending_id,
+                timeout_seconds=decision.timeout_seconds or 120,
+                fn=fn,
+                args=args,
+                kwargs=kwargs,
+                client=client,
+            )
+            settings = get_settings()
+            return (
+                f"This tool requires human approval. Request #{pending_id} was added to the "
+                f"Agent Firewall dashboard ({settings.firewall_url}). "
+                f"Sign in and approve or reject it there. "
+                f"Payload: {json.dumps(payload, default=str)}"
+            )
+
+        return wrapper  # type: ignore[return-value]
+
+    return decorator

agent_firewall/guards/__init__.py

@@ -0,0 +1,33 @@
+"""Agent Firewall — pluggable guardrails (injection, PII, groundedness)."""
+
+from agent_firewall.guards.injection_detector import InjectionScore, score_chunk
+from agent_firewall.guards.pii_masker import (
+    PIIMasker,
+    get_pii_mapping,
+    mask_pii,
+    reset_pii_masker,
+    unmask_pii,
+)
+from agent_firewall.guards.groundedness_checker import (
+    ClaimResult,
+    GroundednessReport,
+    check_groundedness,
+    format_grounded_answer,
+)
+from agent_firewall.guards.pipeline import GuardPipeline, GuardResult
+
+__all__ = [
+    "InjectionScore",
+    "score_chunk",
+    "PIIMasker",
+    "mask_pii",
+    "unmask_pii",
+    "reset_pii_masker",
+    "get_pii_mapping",
+    "ClaimResult",
+    "GroundednessReport",
+    "check_groundedness",
+    "format_grounded_answer",
+    "GuardPipeline",
+    "GuardResult",
+]