cloakwall 0.1.0__py3-none-any.whl

cloakwall/audit.py

@@ -0,0 +1,231 @@
+"""
+cloakwall.audit — tamper-evident local audit trail, plus SIEM export.
+
+What a compliance officer asks for after an incident is not "were you
+redacting" but "prove what happened on 14 March and prove the record has not
+been edited since". A plain log file cannot answer the second half.
+
+Every entry commits to the SHA-256 of the entry before it, so rewriting any
+line breaks the chain from that point on and `verify()` names the line.
+
+A bare hash chain does not survive truncation, though: a prefix of a valid
+chain is itself a valid chain, so dropping the last N entries verifies
+clean. Two things close that gap.
+
+  1. Each entry carries a monotonic sequence number, so gaps in the middle
+     are caught even when the hashes are recomputed.
+
+  2. The head (sequence number plus hash) is mirrored to a separate anchor
+     file after every append, and optionally pushed to your SIEM. Verifying
+     against the anchor detects truncation, because the anchor still knows
+     how many entries there should be.
+
+The anchor is only as good as its separation from the log. On the same
+filesystem it stops accidental truncation and a careless attacker; pushed to
+a SIEM under different access control, it stops a determined one. That trust
+boundary is the operator's to choose, so both are supported and neither is
+claimed to be the other.
+
+The log never stores the values that were redacted. It stores what kind of
+entity was found and how many -- enough to prove the control was working,
+without becoming a second copy of the data you were protecting.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import socket
+import threading
+import time
+import urllib.error
+import urllib.request
+from datetime import datetime, timezone
+
+GENESIS = "0" * 64
+
+
+class AuditLog:
+    def __init__(self, path: str = "./cloakwall-audit.log"):
+        self.path = path
+        self._lock = threading.Lock()   # hooks run concurrently under uvicorn
+        os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True)
+
+    @property
+    def anchor_path(self) -> str:
+        return self.path + ".head"
+
+    def _last(self) -> tuple[int, str]:
+        """Sequence number and hash of the final entry."""
+        if not os.path.exists(self.path):
+            return 0, GENESIS
+        # Read the tail rather than the whole file: this log grows fast and
+        # loading it on every request would be O(n) per request.
+        with open(self.path, "rb") as fh:
+            try:
+                fh.seek(-4096, os.SEEK_END)
+            except OSError:
+                fh.seek(0)
+            tail = fh.read().decode("utf-8", "replace").splitlines()
+        for line in reversed(tail):
+            if line.strip():
+                try:
+                    e = json.loads(line)
+                    return int(e.get("seq", 0)), e["hash"]
+                except Exception:
+                    continue
+        return 0, GENESIS
+
+    def _write_anchor(self, seq: int, head: str) -> None:
+        """Written atomically: a torn anchor would be worse than none."""
+        tmp = self.anchor_path + ".tmp"
+        with open(tmp, "w") as fh:
+            json.dump({"seq": seq, "hash": head,
+                       "ts": datetime.now(timezone.utc).isoformat()}, fh)
+            fh.flush()
+            os.fsync(fh.fileno())
+        os.replace(tmp, self.anchor_path)
+
+    def read_anchor(self) -> dict | None:
+        try:
+            with open(self.anchor_path) as fh:
+                return json.load(fh)
+        except (OSError, ValueError):
+            return None
+
+    def append(self, event: str, **fields) -> dict:
+        with self._lock:
+            prev_seq, prev = self._last()
+            entry = {
+                "ts": datetime.now(timezone.utc).isoformat(),
+                "seq": prev_seq + 1,
+                "event": event,
+                "host": socket.gethostname(),
+                "prev": prev,
+                **fields,
+            }
+            body = json.dumps(entry, sort_keys=True)
+            entry["hash"] = hashlib.sha256((prev + body).encode()).hexdigest()
+            with open(self.path, "a") as fh:
+                fh.write(json.dumps(entry, sort_keys=True) + "\n")
+                fh.flush()
+                os.fsync(fh.fileno())   # survive a pod kill mid-write
+            self._write_anchor(entry["seq"], entry["hash"])
+            return entry
+
+    def verify(self, anchor: dict | None = None) -> tuple[bool, str]:
+        """Verify the chain, and if an anchor is available, verify the log has
+        not been truncated. Pass an anchor retrieved from your SIEM to check
+        against a copy the log's own filesystem cannot reach."""
+        if not os.path.exists(self.path):
+            return True, "no audit log yet"
+
+        prev = GENESIS
+        expected_seq = 0
+        last_hash = GENESIS
+        with open(self.path) as fh:
+            for n, line in enumerate(fh, 1):
+                if not line.strip():
+                    continue
+                e = json.loads(line)
+                claimed = e.pop("hash", None)
+                if e.get("prev") != prev:
+                    return False, f"line {n}: chain break"
+                expected_seq += 1
+                if int(e.get("seq", -1)) != expected_seq:
+                    return False, (f"line {n}: sequence gap "
+                                   f"(expected {expected_seq}, found {e.get('seq')})")
+                recomputed = hashlib.sha256(
+                    (prev + json.dumps(e, sort_keys=True)).encode()).hexdigest()
+                if recomputed != claimed:
+                    return False, f"line {n}: entry altered after writing"
+                prev = last_hash = claimed
+
+        a = anchor if anchor is not None else self.read_anchor()
+        if a:
+            if expected_seq < int(a.get("seq", 0)):
+                return False, (f"truncated: anchor records {a['seq']} entries, "
+                               f"log holds {expected_seq}")
+            if expected_seq == int(a.get("seq", 0)) and last_hash != a.get("hash"):
+                return False, "head does not match anchor"
+            return True, f"chain intact, {expected_seq} entries, anchor matches"
+        return True, f"chain intact, {expected_seq} entries (no anchor to check truncation)"
+
+
+# ------------------------------------------------------------------ #
+# SIEM export                                                         #
+#                                                                     #
+# Fire-and-forget on a worker thread. An audit sink that is slow or   #
+# down must never add latency to, or fail, an LLM request -- the log  #
+# on disk is the system of record and the SIEM is a copy.             #
+# ------------------------------------------------------------------ #
+
+class SIEMExporter:
+    def __init__(self, kind: str | None = None, url: str = "", token: str = "",
+                 timeout: float = 3.0, verify_tls: bool = True):
+        self.kind = (kind or "").lower()
+        self.url = url
+        self.token = token
+        self.timeout = timeout
+        self.verify_tls = verify_tls
+        self.failures = 0
+
+    @classmethod
+    def from_env(cls) -> "SIEMExporter":
+        return cls(
+            kind=os.environ.get("CLOAKWALL_SIEM", ""),
+            url=os.environ.get("CLOAKWALL_SIEM_URL", ""),
+            token=os.environ.get("CLOAKWALL_SIEM_TOKEN", ""),
+        )
+
+    @property
+    def enabled(self) -> bool:
+        return bool(self.kind and self.url and self.token)
+
+    def _payload(self, entry: dict) -> tuple[bytes, dict]:
+        if self.kind == "splunk":
+            body = {
+                "time": time.time(),
+                "host": entry.get("host"),
+                "source": "cloakwall",
+                "sourcetype": "_json",
+                "event": entry,
+            }
+            return json.dumps(body).encode(), {
+                "Authorization": f"Splunk {self.token}",
+                "Content-Type": "application/json",
+            }
+        if self.kind == "datadog":
+            body = [{
+                "ddsource": "cloakwall",
+                "service": "cloakwall",
+                "hostname": entry.get("host"),
+                "message": json.dumps(entry),
+            }]
+            return json.dumps(body).encode(), {
+                "DD-API-KEY": self.token,
+                "Content-Type": "application/json",
+            }
+        # elasticsearch and anything else that accepts a bare JSON document
+        return json.dumps(entry).encode(), {
+            "Authorization": f"Bearer {self.token}",
+            "Content-Type": "application/json",
+        }
+
+    def _post(self, entry: dict) -> None:
+        try:
+            data, headers = self._payload(entry)
+            req = urllib.request.Request(self.url, data=data, headers=headers)
+            ctx = None
+            if not self.verify_tls:
+                import ssl
+                ctx = ssl._create_unverified_context()   # self-signed on-prem SIEM
+            urllib.request.urlopen(req, timeout=self.timeout, context=ctx).read()
+        except (urllib.error.URLError, OSError, ValueError):
+            self.failures += 1     # counted and surfaced in /healthz, never raised
+
+    def send(self, entry: dict) -> None:
+        if not self.enabled:
+            return
+        threading.Thread(target=self._post, args=(entry,), daemon=True).start()

cloakwall/guardrail.py

@@ -0,0 +1,392 @@
+"""
+cloakwall.guardrail — the LiteLLM adapter.
+
+LiteLLM's OSS guardrail framework lets you register a CustomGuardrail without
+an enterprise licence. This class implements the hooks that matter:
+
+    async_pre_call_hook          runs before the request leaves. This is where
+                                 redaction has to happen -- LiteLLM documents
+                                 that pre_call is the mode to use when you
+                                 modify content, because during_call runs in
+                                 parallel with the LLM call and your edits may
+                                 not land before the request goes out.
+                                 Covers messages, the top-level system field,
+                                 prompt (completions, image generation) and
+                                 input (embeddings, Responses API).
+
+    async_post_call_success_hook runs on a non-streamed response. Always
+                                 audited; redacted too when redact_output is
+                                 on, so PII the model generates is masked
+                                 before the caller sees it.
+
+    async_post_call_streaming_iterator_hook
+                                 runs on a streamed response, chunk by chunk.
+                                 With redact_output on, text and tool-call
+                                 arguments go through a StreamRedactor so an
+                                 entity split across chunks is still caught.
+                                 (LiteLLM runs the success hook on streams
+                                 only after delivery, too late to redact.)
+
+Config (litellm config.yaml):
+
+    guardrails:
+      - guardrail_name: cloakwall
+        litellm_params:
+          guardrail: cloakwall.guardrail.Cloakwall
+          mode: [pre_call, post_call]   # post_call only needed for redact_output
+          redaction_mode: mask          # mask | hash | partial
+          redact_output: true           # also mask PII in model responses
+          audit_path: /var/log/cloakwall/audit.log
+          fail_closed: true
+
+Everything runs in-process. No sidecar, no external endpoint, no model
+download -- which is the point: it works unchanged in an air-gapped cluster.
+"""
+
+from __future__ import annotations
+
+import copy
+import os
+from typing import Any, AsyncGenerator, Optional, Union
+
+from .audit import AuditLog, SIEMExporter
+from .redact import Redaction, Redactor, StreamRedactor
+
+try:
+    from litellm.integrations.custom_guardrail import CustomGuardrail
+except ImportError:      # allows the core to be imported and tested standalone
+    class CustomGuardrail:  # type: ignore
+        def __init__(self, **kwargs): pass
+
+
+# Request fields that carry text, across every endpoint the proxy forwards.
+# Anything else in the request body (model, tools, temperature...) is config.
+REQUEST_FIELDS = ("messages", "system", "prompt", "input")
+
+
+class Cloakwall(CustomGuardrail):
+    def __init__(self,
+                 redaction_mode: str = "mask",
+                 entities: Optional[list] = None,
+                 audit_path: str = "./cloakwall-audit.log",
+                 fail_closed: bool = True,
+                 anchor_every: int = 100,
+                 block_on: Optional[list] = None,
+                 redact_output: bool = False,
+                 stream_holdback: int = 64,
+                 **kwargs):
+        self.redactor = Redactor(
+            mode=redaction_mode,
+            entities=set(entities) if entities else None,
+        )
+        self.audit = AuditLog(audit_path)
+        self.siem = SIEMExporter.from_env()
+        # fail_closed decides what happens if redaction itself throws. In a
+        # regulated deployment the safe answer is to reject the request rather
+        # than forward text that may not have been scrubbed.
+        self.fail_closed = fail_closed
+        self.anchor_every = max(1, int(anchor_every))
+        # Entities that must never reach the model at all. Redaction is the
+        # default because it keeps the request useful, but some policies want
+        # the call rejected outright -- a card number in a tool-call argument
+        # is a bug upstream, and silently masking it hides the bug.
+        self.block_on = set(block_on or [])
+        # Off by default: masking output changes what callers receive, and an
+        # app that asks the model to draft an email to a known address would
+        # break. Turn it on where the model may echo or invent identifiers.
+        self.redact_output = bool(redact_output)
+        self.stream_holdback = int(stream_holdback)
+        self.blocked = 0
+        self.requests = 0
+        self.redacted = 0
+        self.responses_redacted = 0
+        super().__init__(**kwargs)
+
+    def _record(self, event: str, **fields) -> None:
+        entry = self.audit.append(event, **fields)
+        self.siem.send(entry)
+        # Mirror the chain head to the SIEM on a cadence. The SIEM sits under
+        # different access control from the log's filesystem, so an anchor
+        # held there is what makes end-truncation detectable by someone who
+        # can write to the pod. Every N entries rather than every entry, to
+        # keep the export volume sane.
+        if self.siem.enabled and entry["seq"] % self.anchor_every == 0:
+            self.siem.send({
+                "event": "cloakwall.anchor",
+                "seq": entry["seq"],
+                "hash": entry["hash"],
+                "ts": entry["ts"],
+                "host": entry["host"],
+            })
+
+    async def async_pre_call_hook(self, user_api_key_dict: Any, cache: Any,
+                                  data: dict, call_type: Any
+                                  ) -> Optional[Union[Exception, str, dict]]:
+        self.requests += 1
+        present = [f for f in REQUEST_FIELDS if data.get(f) is not None]
+        if not present:
+            return data
+
+        try:
+            cleaned, redactions = {}, []
+            for f in present:
+                if f == "messages" and isinstance(data[f], list):
+                    cleaned[f], r = self.redactor.redact_messages(data[f])
+                else:
+                    cleaned[f], r = self.redactor.redact_value(data[f])
+                redactions.extend(r)
+        except Exception as exc:
+            self._record("cloakwall.error", stage="pre_call", error=str(exc))
+            if self.fail_closed:
+                raise
+            return data
+
+        hits = {r.entity for r in redactions}
+        blocked = hits & self.block_on
+        if blocked:
+            self.blocked += 1
+            self._record("cloakwall.blocked",
+                         model=data.get("model"),
+                         key_alias=getattr(user_api_key_dict, "key_alias", None),
+                         entities=_counts(redactions), blocked_on=sorted(blocked))
+            raise ValueError(
+                "Cloakwall blocked this request: "
+                f"{', '.join(sorted(blocked))} present in prompt or tool-call arguments")
+
+        if redactions:
+            self.redacted += 1
+            data.update(cleaned)
+            # Note what was found, never what it was.
+            self._record(
+                "cloakwall.redacted",
+                model=data.get("model"),
+                key_alias=getattr(user_api_key_dict, "key_alias", None),
+                team_id=getattr(user_api_key_dict, "team_id", None),
+                entities=_counts(redactions),
+                total=len(redactions),
+                mode=self.redactor.mode,
+            )
+        return data
+
+    async def async_post_call_success_hook(self, data: dict,
+                                           user_api_key_dict: Any,
+                                           response: Any) -> Any:
+        # A streamed response was already redacted chunk by chunk; LiteLLM
+        # hands this hook the assembled copy after delivery.
+        if self.redact_output and not data.get("stream"):
+            try:
+                redactions = self._redact_response(response)
+            except Exception as exc:
+                self._record("cloakwall.error", stage="post_call", error=str(exc))
+                if self.fail_closed:
+                    raise
+                redactions = []
+            self._record_output(data, user_api_key_dict, redactions, stream=False)
+        self._record(
+            "cloakwall.response",
+            model=data.get("model"),
+            key_alias=getattr(user_api_key_dict, "key_alias", None),
+            finish_reason=_finish_reason(response),
+        )
+        return response
+
+    async def async_post_call_streaming_iterator_hook(
+            self, user_api_key_dict: Any, response: Any, request_data: dict
+    ) -> AsyncGenerator[Any, None]:
+        if not self.redact_output:
+            async for chunk in response:
+                yield chunk
+            return
+
+        # One StreamRedactor per text stream: each choice's content, and each
+        # tool call's arguments, arrive as separate interleaved fragments.
+        streams: dict[tuple, StreamRedactor] = {}
+        templates: dict[tuple, Any] = {}
+        redactions: list[Redaction] = []
+
+        def feed(key, chunk, text):
+            sr = streams.get(key)
+            if sr is None:
+                sr = streams[key] = StreamRedactor(self.redactor, self.stream_holdback)
+            templates[key] = chunk
+            r = sr.feed(text)
+            redactions.extend(r.redactions)
+            return r.text
+
+        def flush(choice_index=None):
+            # Held-back text goes out in a chunk of its own, cloned from the
+            # last chunk that carried that stream, before the finish chunk.
+            out = []
+            for key in list(streams):
+                if choice_index is not None and key[1] != choice_index:
+                    continue
+                sr = streams.pop(key)
+                template = templates.pop(key)
+                if not sr.pending:
+                    continue
+                r = sr.flush()
+                redactions.extend(r.redactions)
+                out.append(_flush_chunk(template, key, r.text))
+            return out
+
+        try:
+            async for chunk in response:
+                finishing = []
+                for ch in _get(chunk, "choices") or []:
+                    idx = _get(ch, "index") or 0
+                    delta = _get(ch, "delta")
+                    if delta is not None:
+                        text = _get(delta, "content")
+                        if isinstance(text, str) and text:
+                            _set(delta, "content", feed(("content", idx), chunk, text))
+                        for tc in _get(delta, "tool_calls") or []:
+                            fn = _get(tc, "function")
+                            args = _get(fn, "arguments") if fn is not None else None
+                            if isinstance(args, str) and args:
+                                key = ("tool", idx, _get(tc, "index") or 0)
+                                _set(fn, "arguments", feed(key, chunk, args))
+                    if _get(ch, "finish_reason"):
+                        finishing.append(idx)
+                for idx in finishing:
+                    for extra in flush(idx):
+                        yield extra
+                yield chunk
+            for extra in flush():
+                yield extra
+        except Exception as exc:
+            self._record("cloakwall.error", stage="stream", error=str(exc))
+            raise
+        finally:
+            self._record_output(request_data, user_api_key_dict, redactions, stream=True)
+
+    def _redact_response(self, response: Any) -> list[Redaction]:
+        """Mask PII in a completed response, in place. Handles LiteLLM's
+        ModelResponse objects and plain dicts: chat choices (content and
+        tool-call arguments), text-completion choices, and Anthropic-style
+        content blocks."""
+        reds: list[Redaction] = []
+        for ch in _get(response, "choices") or []:
+            msg = _get(ch, "message")
+            if msg is not None:
+                reds += self._redact_field(msg, "content")
+                for tc in _get(msg, "tool_calls") or []:
+                    fn = _get(tc, "function")
+                    if fn is not None:
+                        reds += self._redact_field(fn, "arguments")
+            reds += self._redact_field(ch, "text")
+        content = _get(response, "content")
+        if isinstance(content, list):
+            for block in content:
+                kind = _get(block, "type")
+                if kind == "text":
+                    reds += self._redact_field(block, "text")
+                elif kind == "tool_use":
+                    new, r = self.redactor.redact_value(
+                        {"type": "tool_use", "input": _get(block, "input")})
+                    if r:
+                        _set(block, "input", new["input"])
+                        reds += r
+        return reds
+
+    def _redact_field(self, obj: Any, key: str) -> list[Redaction]:
+        value = _get(obj, key)
+        if value is None:
+            return []
+        new, r = self.redactor.redact_value(value)
+        if r:
+            _set(obj, key, new)
+        return r
+
+    def _record_output(self, data: dict, user_api_key_dict: Any,
+                       redactions: list[Redaction], stream: bool) -> None:
+        if not redactions:
+            return
+        self.responses_redacted += 1
+        self._record(
+            "cloakwall.redacted_output",
+            model=(data or {}).get("model"),
+            key_alias=getattr(user_api_key_dict, "key_alias", None),
+            team_id=getattr(user_api_key_dict, "team_id", None),
+            entities=_counts(redactions),
+            total=len(redactions),
+            mode=self.redactor.mode,
+            stream=stream,
+        )
+
+    def stats(self) -> dict:
+        ok, msg = self.audit.verify()
+        anchor = self.audit.read_anchor()
+        return {
+            "requests": self.requests,
+            "requests_with_redactions": self.redacted,
+            "requests_blocked": self.blocked,
+            "responses_with_redactions": self.responses_redacted,
+            "redact_output": self.redact_output,
+            "block_on": sorted(self.block_on),
+            "redaction_mode": self.redactor.mode,
+            "entities_enabled": sorted(self.redactor.entities),
+            "audit_chain": msg,
+            "audit_intact": ok,
+            "audit_entries": (anchor or {}).get("seq", 0),
+            "anchor_every": self.anchor_every,
+            "siem": self.siem.kind if self.siem.enabled else "disabled",
+            "siem_failures": self.siem.failures,
+        }
+
+
+def _counts(redactions: list[Redaction]) -> dict[str, int]:
+    out: dict[str, int] = {}
+    for r in redactions:
+        out[r.entity] = out.get(r.entity, 0) + 1
+    return out
+
+
+def _get(obj: Any, key: str) -> Any:
+    """Field access that works on LiteLLM's response objects and on dicts."""
+    if isinstance(obj, dict):
+        return obj.get(key)
+    return getattr(obj, key, None)
+
+
+def _set(obj: Any, key: str, value: Any) -> None:
+    if isinstance(obj, dict):
+        obj[key] = value
+    else:
+        setattr(obj, key, value)
+
+
+def _flush_chunk(template: Any, key: tuple, text: str) -> Any:
+    """A copy of `template` carrying only the flushed text for one stream:
+    a single choice, no finish_reason, and for tool calls only the arguments
+    (id and name were already sent; repeating them makes clients that
+    concatenate deltas double the function name)."""
+    chunk = copy.deepcopy(template)
+    choice_index = key[1]
+    choices = [c for c in (_get(chunk, "choices") or [])
+               if (_get(c, "index") or 0) == choice_index]
+    _set(chunk, "choices", choices[:1])
+    for ch in choices[:1]:
+        _set(ch, "finish_reason", None)
+        delta = _get(ch, "delta")
+        if key[0] == "content":
+            _set(delta, "content", text)
+            _set(delta, "tool_calls", None)
+        else:
+            tool_index = key[2]
+            calls = [tc for tc in (_get(delta, "tool_calls") or [])
+                     if (_get(tc, "index") or 0) == tool_index]
+            for tc in calls[:1]:
+                _set(tc, "id", None)
+                fn = _get(tc, "function")
+                _set(fn, "name", None)
+                _set(fn, "arguments", text)
+            _set(delta, "tool_calls", calls[:1])
+            _set(delta, "content", None)
+    return chunk
+
+
+def _finish_reason(response: Any) -> Optional[str]:
+    try:
+        return response.choices[0].finish_reason
+    except Exception:
+        return None