feat(dlp): websocket scanning, response headers, extended encoding variants, sk-proj pattern (PRD 0053)

bot_bottle/egress_addon.py

@@ -16,6 +16,7 @@ from mitmproxy import http  # type: ignore[import-not-found]
 
 from egress_addon_core import (  # type: ignore[import-not-found]
     Route,
+    build_inbound_scan_text,
     build_outbound_scan_text,
     decide,
     is_git_push_request,
@@ -148,16 +149,18 @@ class EgressAddon:
             flow.request.headers["authorization"] = decision.inject_authorization
 
     def response(self, flow: http.HTTPFlow) -> None:
-        """DLP inbound scan on response bodies (PRD 0053)."""
+        """DLP inbound scan on response headers and body."""
         route = match_route(self.routes, flow.request.pretty_host)
         if route is None:
             return
         if flow.response is None:
             return
+        resp_headers = {k.lower(): v for k, v in flow.response.headers.items()}
         body = flow.response.get_text(strict=False) or ""
-        if not body:
+        scan_text = build_inbound_scan_text(resp_headers, body)
+        if not scan_text:
             return
-        result = scan_inbound(route, body)
+        result = scan_inbound(route, scan_text)
         if result is None:
             return
         if result.severity == "block":
@@ -165,5 +168,34 @@ class EgressAddon:
         elif result.severity == "warn":
             sys.stderr.write(f"egress DLP warn: {result.reason}\n")
 
+    def websocket_message(self, flow: http.HTTPFlow) -> None:
+        """DLP scan on WebSocket frames.
+
+        Outbound frames (from_client) are scanned for credential leakage;
+        inbound frames are scanned for prompt injection.  On a block the
+        entire connection is killed — there is no HTTP response surface to
+        write to after the upgrade.
+        """
+        if flow.websocket is None:  # type: ignore[union-attr]
+            return
+        route = match_route(self.routes, flow.request.pretty_host)
+        if route is None:
+            return
+        message = flow.websocket.messages[-1]  # type: ignore[union-attr]
+        content = message.content.decode("utf-8", errors="replace")
+        if message.from_client:
+            result = scan_outbound(route, content, os.environ)
+            if result is not None and result.severity == "block":
+                sys.stderr.write(f"egress DLP: {result.reason}\n")
+                flow.kill()  # type: ignore[union-attr]
+        else:
+            result = scan_inbound(route, content)
+            if result is not None:
+                if result.severity == "block":
+                    sys.stderr.write(f"egress DLP: {result.reason}\n")
+                    flow.kill()  # type: ignore[union-attr]
+                elif result.severity == "warn":
+                    sys.stderr.write(f"egress DLP warn: {result.reason}\n")
+
 
 addons = [EgressAddon()]