Merge 02f08956a6 into 703eb4fa68

litellm/proxy/custom_hooks/x_trust_hook.py

@@ -0,0 +1,72 @@
+"""
+X-Trust Hook for LiteLLM Proxy
+Annotates requests with human presence score via X-Trust header.
+Zero KYC, zero PII. github.com/htl-syterme/htl-core
+"""
+import hmac
+import hashlib
+import base64
+import json
+import time
+from typing import Optional
+from litellm.integrations.custom_logger import CustomLogger
+
+
+class XTrustHook(CustomLogger):
+    """
+    LiteLLM hook that validates X-Trust headers.
+    Annotates requests with human presence score.
+    Doctrine AIR: annotate, never block.
+    """
+
+    def __init__(self, secret: str, min_score: float = 0.0):
+        self.secret = secret
+        self.min_score = min_score
+
+    def _verify(self, token: str) -> Optional[dict]:
+        try:
+            parts = token.split(".")
+            if len(parts) != 3 or parts[0] != "v1":
+                return None
+            _, payload_b64, sig_b64 = parts
+            padding = 4 - len(payload_b64) % 4
+            payload_bytes = base64.b64decode(
+                payload_b64.replace("-", "+").replace("_", "/") + "=" * padding
+            )
+            sig_bytes = base64.b64decode(
+                sig_b64.replace("-", "+").replace("_", "/") + "=" * (4 - len(sig_b64) % 4)
+            )
+            expected = hmac.new(
+                self.secret.encode(), payload_b64.encode(), hashlib.sha256
+            ).digest()
+            if not hmac.compare_digest(expected, sig_bytes):
+                return None
+            payload = json.loads(payload_bytes)
+            now = int(time.time())
+            if now > payload.get("exp", 0):
+                return None
+            if now - payload.get("iat", 0) > 120:
+                return None
+            score = payload.get("score", 0)
+            if not (0 <= score <= 1):
+                return None
+            return payload
+        except Exception:
+            return None
+
+    async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
+        token = data.get("metadata", {}).get("headers", {}).get("x-trust", "")
+        payload = self._verify(token) if token else None
+        score = payload["score"] if payload else 0.0
+        trusted = payload is not None and score >= self.min_score
+        # Redact token after verification to prevent replay via logging
+        if "headers" in data.get("metadata", {}):
+            data["metadata"]["headers"].pop("x-trust", None)
+        if "metadata" not in data:
+            data["metadata"] = {}
+        data["metadata"]["x_trust"] = {
+            "trusted": trusted,
+            "score": score,
+            "annotated": True,
+        }
+        return data

tests/test_litellm/test_x_trust_hook.py

@@ -0,0 +1,56 @@
+import pytest
+import time
+import hmac
+import hashlib
+import base64
+import json
+
+
+def make_token(secret: str, score: float = 0.88, expired: bool = False) -> str:
+    now = int(time.time())
+    payload = {"sub": "test", "score": score, "iat": now, "exp": now + (- 1 if expired else 120)}
+    payload_b64 = base64.urlsafe_b64encode(json.dumps(payload).encode()).rstrip(b"=").decode()
+    sig = hmac.new(secret.encode(), payload_b64.encode(), hashlib.sha256).digest()
+    sig_b64 = base64.urlsafe_b64encode(sig).rstrip(b"=").decode()
+    return f"v1.{payload_b64}.{sig_b64}"
+
+
+SECRET = "test_secret_htl"
+
+
+@pytest.fixture
+def hook():
+    from litellm.proxy.custom_hooks.x_trust_hook import XTrustHook
+    return XTrustHook(secret=SECRET, min_score=0.0)
+
+
+@pytest.mark.asyncio
+async def test_valid_token_annotates_trusted(hook):
+    token = make_token(SECRET, score=0.88)
+    data = {"metadata": {"headers": {"x-trust": token}}}
+    result = await hook.async_pre_call_hook(None, None, data, None)
+    assert result["metadata"]["x_trust"]["trusted"] is True
+    assert result["metadata"]["x_trust"]["score"] == 0.88
+
+
+@pytest.mark.asyncio
+async def test_wrong_secret_annotates_untrusted(hook):
+    token = make_token("wrong_secret", score=0.88)
+    data = {"metadata": {"headers": {"x-trust": token}}}
+    result = await hook.async_pre_call_hook(None, None, data, None)
+    assert result["metadata"]["x_trust"]["trusted"] is False
+
+
+@pytest.mark.asyncio
+async def test_no_token_annotates_untrusted(hook):
+    data = {"metadata": {"headers": {}}}
+    result = await hook.async_pre_call_hook(None, None, data, None)
+    assert result["metadata"]["x_trust"]["trusted"] is False
+
+
+@pytest.mark.asyncio
+async def test_expired_token_annotates_untrusted(hook):
+    token = make_token(SECRET, score=0.88, expired=True)
+    data = {"metadata": {"headers": {"x-trust": token}}}
+    result = await hook.async_pre_call_hook(None, None, data, None)
+    assert result["metadata"]["x_trust"]["trusted"] is False