openscript-sdk 0.1.0__py3-none-any.whl

openscript_sdk/policies/compliance.py

@@ -0,0 +1,261 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from enum import StrEnum
+from typing import TYPE_CHECKING, Any, cast
+
+import structlog
+
+from openscript_sdk.contracts.interceptor import BasePolicy
+from openscript_sdk.contracts.server_types import Event, EventType
+from openscript_sdk.contracts.types import ActionContext, FailureMode, InterceptorDecision
+from openscript_sdk.policies.secrets import find_secrets
+
+if TYPE_CHECKING:
+    from openscript_sdk.events.writer import EventWriter
+
+logger = structlog.get_logger(__name__)
+
+# ---------------------------------------------------------------------------
+# PHI pattern bank
+# Structural patterns for protected-health-information categories a regex can
+# honestly detect: coded identifiers and identifier-labelled values, not
+# free-text medical meaning.
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class _PHIPattern:
+    regex: re.Pattern[str]
+    label: str
+
+
+def _p(pattern: str, label: str) -> _PHIPattern:
+    return _PHIPattern(re.compile(pattern, re.IGNORECASE), label)
+
+
+_PHI_PATTERNS: list[_PHIPattern] = [
+    # ICD-10 diagnosis codes — dotted form only (E11.9), the undotted form
+    # collides with too many product/part numbers
+    _PHIPattern(re.compile(r"\b[A-TV-Z]\d{2}\.\d{1,4}\b"), "icd10_code"),
+    # Medical record numbers introduced by their label
+    _p(r"\b(?:mrn|medical\s+record\s+(?:number|no\.?))\s*[:#]?\s*\d{5,}\b", "mrn"),
+    # Insurance / member identifiers introduced by their label
+    _p(
+        r"\b(?:member|policy|subscriber|insurance)\s+(?:id|number|no\.?)\s*[:#]?\s*[A-Z0-9-]{5,}\b",
+        "insurance_id",
+    ),
+    # Medication name (common drug-class suffixes) followed by a dosage
+    _p(
+        r"\b[a-z]+(?:cillin|mycin|statin|prazole|sartan|olol|pril|formin|azepam|oxetine|codone)\b"
+        r"\s+\d+\s*(?:mg|mcg|ml)\b",
+        "medication_dosage",
+    ),
+    # Explicit diagnosis statements
+    _p(r"\b(?:diagnosed\s+with|diagnosis\s+of)\s+[a-z]", "diagnosis_mention"),
+]
+
+_CREDENTIAL_RISK = 0.9
+_PHI_RISK = 0.6
+
+# Preset name -> description; the public contract of CompliancePolicy(rules=...)
+KNOWN_RULES: dict[str, str] = {
+    "phi_detection": "detect PHI identifiers (ICD-10 codes, MRNs, insurance IDs, "
+    "medication dosages, diagnosis statements) in input and output",
+    "credential_output_guard": "deny any output containing raw credential material "
+    "(reuses the secrets pattern bank)",
+    "data_access_audit": "emit an audit event for every tool call",
+}
+
+
+class PHIMode(StrEnum):
+    ANNOTATE = "annotate"
+    DENY = "deny"
+
+
+def find_phi(text: str) -> list[tuple[str, str]]:
+    """Return list of (matched_text, label) pairs of PHI found in *text*.
+
+    Standalone helper — callable outside a policy pipeline.
+    """
+    found: list[tuple[str, str]] = []
+    for pat in _PHI_PATTERNS:
+        for m in pat.regex.finditer(text):
+            found.append((m.group(), pat.label))
+    return found
+
+
+def _flatten_text(obj: Any) -> str:
+    if isinstance(obj, str):
+        return obj
+    if isinstance(obj, dict):
+        return " ".join(_flatten_text(v) for v in obj.values())
+    if isinstance(obj, list):
+        return " ".join(_flatten_text(item) for item in obj)
+    return ""
+
+
+class CompliancePolicy(BasePolicy):
+    """Concrete, honestly-scoped checks that assist compliance programs.
+
+    This policy ASSISTS compliance work (its checks map onto common
+    GDPR/HIPAA/SOC2 controls); it does not confer or certify compliance
+    with any regulation. Presets are named after what they check:
+
+    - "phi_detection": scan input (before_action) and output (after_action)
+      for PHI identifiers — ICD-10 codes, MRNs, insurance IDs, medication
+      dosages, diagnosis statements. Follows *phi_mode*: ANNOTATE (default)
+      records findings only; DENY blocks. Pair with PIIPolicy/SecretsPolicy
+      for redaction.
+    - "credential_output_guard": deny any output containing raw credential
+      material (secrets pattern bank, internal URLs excluded). Always
+      blocks — it is a guard, not a detector.
+    - "data_access_audit": emit a POLICY_EVALUATED audit event for every
+      tool call (requires a writer; annotates metadata either way).
+
+    Results accumulate in context.metadata["compliance"] using the
+    standardized shape (risk keeps the max across phases; credentials in
+    output score 0.9, PHI findings 0.6):
+      {"risk": float, "category": "compliance", "rules": [...],
+       "phi": [labels...], "credentials": [labels...], "audited": bool}
+    """
+
+    failure_mode: FailureMode = FailureMode.FAIL_OPEN
+
+    def __init__(
+        self,
+        rules: list[str],
+        phi_mode: PHIMode | str = PHIMode.ANNOTATE,
+        writer: EventWriter | None = None,
+    ) -> None:
+        unknown = set(rules) - set(KNOWN_RULES)
+        if unknown:
+            raise ValueError(
+                f"unknown compliance rule(s): {sorted(unknown)}; known rules: {sorted(KNOWN_RULES)}"
+            )
+        if not rules:
+            raise ValueError(f"rules must name at least one preset: {sorted(KNOWN_RULES)}")
+        self._rules = list(dict.fromkeys(rules))
+        self._phi_mode = PHIMode(phi_mode)
+        self._writer = writer
+        self._sequence_counters: dict[str, int] = {}
+
+    async def before_action(self, context: ActionContext) -> ActionContext:
+        meta = self._meta(context)
+
+        if "data_access_audit" in self._rules and context.action == "tool_call":
+            meta["audited"] = True
+            if self._writer is not None:
+                await self._emit_event(
+                    context,
+                    policy="data_access_audit",
+                    payload={
+                        "tool": context.input_data.get("name", ""),
+                        "arg_keys": sorted(context.input_data.get("args", {})),
+                    },
+                )
+
+        if "phi_detection" in self._rules:
+            await self._check_phi(context, context.input_data, phase="input")
+
+        return context
+
+    async def after_action(self, context: ActionContext) -> ActionContext:
+        meta = self._meta(context)
+
+        if "phi_detection" in self._rules:
+            await self._check_phi(context, context.output_data, phase="output")
+
+        if "credential_output_guard" in self._rules:
+            text = _flatten_text(context.output_data)
+            credentials = sorted(
+                {label for _, label in find_secrets(text) if label != "internal_url"}
+            )
+            if credentials:
+                meta["credentials"] = list(dict.fromkeys([*meta["credentials"], *credentials]))
+                meta["risk"] = max(meta["risk"], _CREDENTIAL_RISK)
+                context.decision = InterceptorDecision.DENY
+                context.decision_reason = f"Credential material in output: {', '.join(credentials)}"
+                logger.warning(
+                    "credential_output_blocked",
+                    session_id=context.session_id,
+                    agent_id=context.agent_id,
+                    types=credentials,
+                )
+                if self._writer is not None:
+                    await self._emit_event(
+                        context,
+                        policy="credential_output_guard",
+                        payload={"credential_types": credentials},
+                    )
+
+        return context
+
+    async def _check_phi(self, context: ActionContext, data: dict[str, Any], phase: str) -> None:
+        meta = self._meta(context)
+        text = _flatten_text(data)
+        found = sorted({label for _, label in find_phi(text)})
+        if not found:
+            return
+
+        meta["phi"] = list(dict.fromkeys([*meta["phi"], *found]))
+        meta["risk"] = max(meta["risk"], _PHI_RISK)
+
+        logger.warning(
+            "phi_detected",
+            session_id=context.session_id,
+            agent_id=context.agent_id,
+            types=found,
+            phase=phase,
+            mode=self._phi_mode.value,
+        )
+
+        if self._phi_mode == PHIMode.DENY:
+            context.decision = InterceptorDecision.DENY
+            context.decision_reason = f"PHI detected in {phase}: {', '.join(found)}"
+
+        if self._writer is not None:
+            await self._emit_event(
+                context,
+                policy="phi_detection",
+                payload={"phi_types": found, "phase": phase, "mode": self._phi_mode.value},
+            )
+
+    def _meta(self, context: ActionContext) -> dict[str, Any]:
+        return cast(
+            dict[str, Any],
+            context.metadata.setdefault(
+                "compliance",
+                {
+                    "risk": 0.0,
+                    "category": "compliance",
+                    "rules": self._rules,
+                    "phi": [],
+                    "credentials": [],
+                    "audited": False,
+                },
+            ),
+        )
+
+    async def _emit_event(
+        self, context: ActionContext, policy: str, payload: dict[str, Any]
+    ) -> None:
+        seq = self._next_sequence(context.session_id)
+        event = Event(
+            session_id=context.session_id,
+            agent_id=context.agent_id,
+            event_type=EventType.POLICY_EVALUATED,
+            payload={
+                "policy": policy,
+                **payload,
+                "decision": context.decision.value,
+            },
+            sequence_num=seq,
+        )
+        await self._writer.write(event)  # type: ignore[union-attr]
+
+    def _next_sequence(self, session_id: str) -> int:
+        count = self._sequence_counters.get(session_id, 0) + 1
+        self._sequence_counters[session_id] = count
+        return count

openscript_sdk/policies/config.py

@@ -0,0 +1,244 @@
+from __future__ import annotations
+
+import importlib
+from collections.abc import Callable
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import yaml
+from pydantic import BaseModel, Field
+
+from openscript_sdk.contracts.interceptor import Policy
+from openscript_sdk.interceptors.base import NoopInterceptor
+from openscript_sdk.interceptors.event_writer import AuditPolicy
+from openscript_sdk.interceptors.pii import PIIMode, PIIPolicy
+from openscript_sdk.interceptors.threat import PromptInjectionPolicy
+from openscript_sdk.policies.compliance import CompliancePolicy, PHIMode
+from openscript_sdk.policies.harmful_request import HarmfulRequestPolicy
+from openscript_sdk.policies.output_schema import OutputSchemaPolicy
+from openscript_sdk.policies.secrets import InternalURLMode, SecretsPolicy
+from openscript_sdk.policies.session_taint import SessionTaintPolicy
+from openscript_sdk.policies.tool_firewall import ToolFirewallPolicy, ToolRules
+from openscript_sdk.policies.toxicity import ToxicityPolicy
+
+if TYPE_CHECKING:
+    from openscript_sdk.events.writer import EventWriter
+
+
+class PolicyConfig(BaseModel):
+    """Declarative policy pipeline: ordered mapping of policy name -> params.
+
+    Example YAML (see policies_example.yaml):
+
+        policies:
+          prompt_injection:
+            threshold: 0.6
+          pii:
+            mode: redact
+          secrets:
+            mode: deny
+          tool_firewall:
+            rules_path: tools.yaml
+
+    Order is significant — policies run in declaration order.
+    """
+
+    policies: dict[str, dict[str, Any] | None] = Field(default_factory=dict)
+
+
+# A factory receives the (possibly empty) params dict for its entry plus the
+# shared EventWriter (None when the caller doesn't persist events).
+PolicyFactory = Callable[[dict[str, Any], "EventWriter | None"], Policy]
+
+
+def _take(params: dict[str, Any], name: str, allowed: set[str]) -> None:
+    unknown = set(params) - allowed
+    if unknown:
+        raise ValueError(
+            f"unknown parameter(s) for policy '{name}': {sorted(unknown)}; "
+            f"allowed: {sorted(allowed)}"
+        )
+
+
+def _build_prompt_injection(params: dict[str, Any], writer: EventWriter | None) -> Policy:
+    _take(params, "prompt_injection", {"threshold", "tool_output"})
+    return PromptInjectionPolicy(
+        threshold=params.get("threshold", 0.5),
+        writer=writer,
+        tool_output=params.get("tool_output", "annotate"),
+    )
+
+
+def _build_pii(params: dict[str, Any], writer: EventWriter | None) -> Policy:
+    _take(params, "pii", {"mode"})
+    return PIIPolicy(mode=PIIMode(params.get("mode", "redact")), writer=writer)
+
+
+def _build_secrets(params: dict[str, Any], writer: EventWriter | None) -> Policy:
+    _take(params, "secrets", {"mode", "internal_url_mode", "internal_url_allowlist"})
+    return SecretsPolicy(
+        mode=PIIMode(params.get("mode", "redact")),
+        writer=writer,
+        internal_url_mode=InternalURLMode(params.get("internal_url_mode", "annotate")),
+        internal_url_allowlist=params.get("internal_url_allowlist", ()),
+    )
+
+
+def _build_tool_firewall(params: dict[str, Any], writer: EventWriter | None) -> Policy:
+    _take(params, "tool_firewall", {"rules", "rules_path"})
+    rules = params.get("rules")
+    return ToolFirewallPolicy(
+        rules=ToolRules.model_validate(rules) if rules is not None else None,
+        rules_path=params.get("rules_path"),
+    )
+
+
+def _build_session_taint(params: dict[str, Any], writer: EventWriter | None) -> Policy:
+    _take(params, "session_taint", {"sources", "sinks", "on_taint", "known", "updates"})
+    missing = sorted({"sources", "sinks"} - set(params))
+    if missing:
+        # No defaults: which tools read untrusted content and which change
+        # state is a fact about the application, not something to guess.
+        raise ValueError(f"policy 'session_taint' needs {' and '.join(missing)}")
+    return SessionTaintPolicy(
+        sources=params["sources"],
+        sinks=params["sinks"],
+        on_taint=params.get("on_taint", "require_approval"),
+        known=params.get("known"),
+        updates=params.get("updates"),
+    )
+
+
+def _resolve_model(path: str) -> type[BaseModel]:
+    """Import a pydantic model from a dotted path: "pkg.module:Class" or "pkg.module.Class"."""
+    module_name, _, attr = path.partition(":") if ":" in path else path.rpartition(".")
+    if not module_name or not attr:
+        raise ValueError(f"invalid model path {path!r}; expected 'package.module:ClassName'")
+    try:
+        model = getattr(importlib.import_module(module_name), attr)
+    except (ImportError, AttributeError) as exc:
+        raise ValueError(f"cannot import model {path!r}: {exc}") from exc
+    if not (isinstance(model, type) and issubclass(model, BaseModel)):
+        raise ValueError(f"{path!r} is not a pydantic BaseModel subclass")
+    return model
+
+
+def _build_output_schema(params: dict[str, Any], writer: EventWriter | None) -> Policy:
+    _take(
+        params,
+        "output_schema",
+        {
+            "model",
+            "on_invalid",
+            "grounding_source",
+            "hallucination_mode",
+            "hallucination_threshold",
+        },
+    )
+    model = params.get("model")
+    return OutputSchemaPolicy(
+        model=_resolve_model(model) if isinstance(model, str) else model,
+        on_invalid=params.get("on_invalid", "deny"),
+        grounding_source=params.get("grounding_source"),
+        hallucination_mode=params.get("hallucination_mode", "keyword"),
+        hallucination_threshold=params.get("hallucination_threshold", 0.5),
+        writer=writer,
+    )
+
+
+def _build_toxicity(params: dict[str, Any], writer: EventWriter | None) -> Policy:
+    _take(params, "toxicity", {"threshold"})
+    return ToxicityPolicy(threshold=params.get("threshold", 0.5), writer=writer)
+
+
+def _build_harmful_request(params: dict[str, Any], writer: EventWriter | None) -> Policy:
+    _take(params, "harmful_request", {"threshold"})
+    return HarmfulRequestPolicy(threshold=params.get("threshold", 0.5), writer=writer)
+
+
+def _build_semantic(params: dict[str, Any], writer: EventWriter | None) -> Policy:
+    _take(params, "semantic", {"threshold", "mode", "local_files_only"})
+    # Imported here, not at module scope: the policy is opt-in, and the module
+    # that defines it is the one place a model dependency is allowed to live.
+    from openscript_sdk.policies.semantic import SemanticPolicy
+
+    return SemanticPolicy(
+        threshold=params.get("threshold"),
+        mode=params.get("mode", "deny"),
+        local_files_only=params.get("local_files_only", False),
+        writer=writer,
+    )
+
+
+def _build_compliance(params: dict[str, Any], writer: EventWriter | None) -> Policy:
+    _take(params, "compliance", {"rules", "phi_mode"})
+    return CompliancePolicy(
+        rules=params.get("rules", []),
+        phi_mode=PHIMode(params.get("phi_mode", "annotate")),
+        writer=writer,
+    )
+
+
+def _build_audit(params: dict[str, Any], writer: EventWriter | None) -> Policy:
+    _take(params, "audit", set())
+    if writer is None:
+        raise ValueError("policy 'audit' requires load_policies(..., writer=<EventWriter>)")
+    return AuditPolicy(writer)
+
+
+def _build_noop(params: dict[str, Any], writer: EventWriter | None) -> Policy:
+    _take(params, "noop", set())
+    return NoopInterceptor()
+
+
+_REGISTRY: dict[str, PolicyFactory] = {
+    "prompt_injection": _build_prompt_injection,
+    "toxicity": _build_toxicity,
+    "harmful_request": _build_harmful_request,
+    "pii": _build_pii,
+    "secrets": _build_secrets,
+    "compliance": _build_compliance,
+    "semantic": _build_semantic,
+    "tool_firewall": _build_tool_firewall,
+    "session_taint": _build_session_taint,
+    "output_schema": _build_output_schema,
+    "audit": _build_audit,
+    "noop": _build_noop,
+}
+
+
+def register_policy(name: str, factory: PolicyFactory) -> None:
+    """Register a custom policy factory under *name* for use in load_policies configs."""
+    _REGISTRY[name] = factory
+
+
+def load_policies(
+    config: PolicyConfig | dict[str, Any] | str | Path,
+    writer: EventWriter | None = None,
+) -> list[Policy]:
+    """Assemble a ``policies=[...]`` list from a config.
+
+    *config* may be a PolicyConfig, a dict (with or without the top-level
+    "policies" key), or a path to a YAML file. *writer* is injected into
+    policies that emit events (prompt_injection, pii, secrets, audit).
+
+    Raises ValueError on unknown policy names or parameters.
+    """
+    if isinstance(config, (str, Path)):
+        with open(config) as fh:
+            raw = yaml.safe_load(fh) or {}
+        cfg = PolicyConfig.model_validate(raw)
+    elif isinstance(config, PolicyConfig):
+        cfg = config
+    else:
+        if "policies" not in config:
+            config = {"policies": config}
+        cfg = PolicyConfig.model_validate(config)
+
+    policies: list[Policy] = []
+    for name, params in cfg.policies.items():
+        factory = _REGISTRY.get(name)
+        if factory is None:
+            raise ValueError(f"unknown policy '{name}'; known policies: {sorted(_REGISTRY)}")
+        policies.append(factory(params or {}, writer))
+    return policies