From a24fe0264d55a0dd79d880ca3f0cf8a0f607ae33 Mon Sep 17 00:00:00 2001 From: codex Date: Mon, 27 Jul 2026 02:33:18 +0000 Subject: [PATCH] refactor(egress): extract outbound DLP request stage --- bot_bottle/gateway/egress/addon.py | 62 ++------------ .../gateway/egress/outbound_pipeline.py | 83 +++++++++++++++++++ tests/unit/test_egress_outbound_pipeline.py | 81 ++++++++++++++++++ 3 files changed, 173 insertions(+), 53 deletions(-) create mode 100644 bot_bottle/gateway/egress/outbound_pipeline.py create mode 100644 tests/unit/test_egress_outbound_pipeline.py diff --git a/bot_bottle/gateway/egress/addon.py b/bot_bottle/gateway/egress/addon.py index aa6f6271..1ede2cc7 100644 --- a/bot_bottle/gateway/egress/addon.py +++ b/bot_bottle/gateway/egress/addon.py @@ -16,7 +16,7 @@ import typing from mitmproxy import http # type: ignore[import-not-found] # pylint: disable=import-error from bot_bottle.constants import IDENTITY_HEADER -from bot_bottle.gateway.egress.dlp_detectors import redact_tokens, strip_crlf +from bot_bottle.gateway.egress.dlp_detectors import redact_tokens from bot_bottle.gateway.egress.dlp_config import ( DEFAULT_OUTBOUND_ON_MATCH, ON_MATCH_BLOCK, @@ -25,12 +25,11 @@ from bot_bottle.gateway.egress.dlp_config import ( from bot_bottle.gateway.egress.context import resolve_client_context from bot_bottle.gateway.egress.dlp import ( build_inbound_scan_text, - build_outbound_scan_text, build_token_allow_payload, - outbound_scan_headers, scan_inbound, scan_outbound, ) +from bot_bottle.gateway.egress.outbound_pipeline import redact_request, scan_request from bot_bottle.gateway.egress.matching import ( decide, match_route, @@ -496,20 +495,12 @@ class EgressAddon: Loops so the supervise policy can re-scan after each approval — a second, un-approved token in the same request is still caught.""" while True: - request_path, _, query = flow.request.path.partition("?") - body = flow.request.get_text(strict=False) or "" - headers = outbound_scan_headers(route, dict(flow.request.headers)) - scan_text = build_outbound_scan_text( - flow.request.pretty_host, request_path, query, headers, body, - ) - # CRLF is scanned only over the request line + headers, never the - # body (see scan_outbound) — a body is not an injection vector. - crlf_text = build_outbound_scan_text( - flow.request.pretty_host, request_path, query, headers, "", - ) - result = scan_outbound( - route, scan_text, env, - safe_tokens=self._safe_tokens_for(slug), crlf_text=crlf_text, + request_path, _, _ = flow.request.path.partition("?") + result = scan_request( + flow.request, + route, + env, + safe_tokens=self._safe_tokens_for(slug), ) if result is None or result.severity != "block": return True @@ -519,7 +510,7 @@ class EgressAddon: # redact scrubs every detection (tokens and structural CRLF) and # forwards; it fails closed only if a match survives the scrub. if policy == ON_MATCH_REDACT: - if self._redact_outbound(flow, route, env): + if redact_request(flow.request, route, env): if self._flow_log(flow) >= LOG_BLOCKS: sys.stderr.write(json.dumps({ "event": "egress_redacted", @@ -552,41 +543,6 @@ class EgressAddon: return False # _supervise_token_block wrote the 403 response # loop: the approved value is now in safe_tokens; re-scan. - def _redact_outbound( - self, flow: http.HTTPFlow, route: Route, env: "typing.Mapping[str, str]", - ) -> bool: - """Scrub detected tokens (and CRLF injection sequences) from the mutable - request surfaces (body, headers, path/query) and re-scan. `env` is the - per-bottle env overlay. Returns True if the request is now clean; False - if a block-severity match remains on a surface redaction cannot rewrite - (the hostname) so the caller fails closed.""" - body = flow.request.get_text(strict=False) - if body: - redacted_body = redact_tokens(body, env=env) - if redacted_body != body: - flow.request.text = redacted_body - for name, value in list(flow.request.headers.items()): - if name.lower() == "host": - continue # routing-critical; never a legitimate token - redacted = strip_crlf(redact_tokens(value, env=env)) - if redacted != value: - flow.request.headers[name] = redacted - redacted_path = strip_crlf(redact_tokens(flow.request.path, env=env)) - if redacted_path != flow.request.path: - flow.request.path = redacted_path - - request_path, _, query = flow.request.path.partition("?") - new_body = flow.request.get_text(strict=False) or "" - headers = outbound_scan_headers(route, dict(flow.request.headers)) - scan_text = build_outbound_scan_text( - flow.request.pretty_host, request_path, query, headers, new_body, - ) - crlf_text = build_outbound_scan_text( - flow.request.pretty_host, request_path, query, headers, "", - ) - result = scan_outbound(route, scan_text, env, crlf_text=crlf_text) - return result is None or result.severity != "block" - async def _supervise_token_block( self, flow: http.HTTPFlow, diff --git a/bot_bottle/gateway/egress/outbound_pipeline.py b/bot_bottle/gateway/egress/outbound_pipeline.py new file mode 100644 index 00000000..bf7cfc64 --- /dev/null +++ b/bot_bottle/gateway/egress/outbound_pipeline.py @@ -0,0 +1,83 @@ +"""Outbound DLP request scanning and redaction for the egress pipeline.""" + +from __future__ import annotations + +from typing import ItemsView, Mapping, Protocol + +from .dlp import ( + build_outbound_scan_text, + outbound_scan_headers, + scan_outbound, +) +from .dlp_detectors import redact_tokens, strip_crlf +from .types import Route, ScanResult + + +class MutableHeaders(Protocol): + def items(self) -> ItemsView[str, str]: ... + def __getitem__(self, name: str, /) -> str: ... + def __setitem__(self, name: str, value: str, /) -> None: ... + + +class MutableRequest(Protocol): + pretty_host: str + path: str + headers: MutableHeaders + text: str + + def get_text(self, strict: bool = False) -> str | None: ... + + +def scan_request( + request: MutableRequest, + route: Route, + env: Mapping[str, str], + *, + safe_tokens: set[str] | None = None, +) -> ScanResult | None: + """Scan all mutable outbound request surfaces in their canonical order.""" + request_path, _, query = request.path.partition("?") + headers = outbound_scan_headers(route, dict(request.headers.items())) + body = request.get_text(strict=False) or "" + scan_text = build_outbound_scan_text( + request.pretty_host, request_path, query, headers, body, + ) + # Bodies cannot alter HTTP framing, so CRLF detection is deliberately + # restricted to the request line and headers. + crlf_text = build_outbound_scan_text( + request.pretty_host, request_path, query, headers, "", + ) + return scan_outbound( + route, + scan_text, + env, + safe_tokens=safe_tokens, + crlf_text=crlf_text, + ) + + +def redact_request( + request: MutableRequest, + route: Route, + env: Mapping[str, str], +) -> bool: + """Redact mutable request surfaces and return whether the result is clean.""" + body = request.get_text(strict=False) + if body: + redacted_body = redact_tokens(body, env=env) + if redacted_body != body: + request.text = redacted_body + for name, value in list(request.headers.items()): + if name.lower() == "host": + continue + redacted = strip_crlf(redact_tokens(value, env=env)) + if redacted != value: + request.headers[name] = redacted + redacted_path = strip_crlf(redact_tokens(request.path, env=env)) + if redacted_path != request.path: + request.path = redacted_path + result = scan_request(request, route, env) + return result is None or result.severity != "block" + + +__all__ = ["MutableRequest", "redact_request", "scan_request"] diff --git a/tests/unit/test_egress_outbound_pipeline.py b/tests/unit/test_egress_outbound_pipeline.py new file mode 100644 index 00000000..38a71e76 --- /dev/null +++ b/tests/unit/test_egress_outbound_pipeline.py @@ -0,0 +1,81 @@ +"""Unit tests for framework-neutral outbound DLP request stages.""" + +from __future__ import annotations + +import unittest + +from bot_bottle.gateway.egress.outbound_pipeline import ( + MutableHeaders, + redact_request, + scan_request, +) +from bot_bottle.gateway.egress.types import Route + + +class _Headers(dict[str, str]): + pass + + +class _Request: + def __init__( + self, + *, + host: str = "api.example.com", + path: str = "/v1/messages", + headers: dict[str, str] | None = None, + body: str = "", + ) -> None: + self.pretty_host = host + self.path = path + self.headers: MutableHeaders = _Headers(headers or {}) + self.text = body + + def get_text(self, strict: bool = False) -> str | None: + del strict + return self.text + + +class TestOutboundScan(unittest.TestCase): + def test_detects_secret_in_body(self) -> None: + request = _Request(body="token=sk-" + "a" * 48) + + result = scan_request(request, Route(host="api.example.com"), {}) + + self.assertIsNotNone(result) + self.assertEqual("block", result.severity if result else None) + + def test_safe_token_is_ignored(self) -> None: + token = "sk-" + "a" * 48 + request = _Request(body=f"token={token}") + + result = scan_request( + request, + Route(host="api.example.com"), + {}, + safe_tokens={token}, + ) + + self.assertIsNone(result) + + +class TestOutboundRedaction(unittest.TestCase): + def test_redacts_body_header_and_path_but_preserves_host(self) -> None: + token = "sk-" + "a" * 48 + request = _Request( + path=f"/v1/messages?key={token}", + headers={"Host": "api.example.com", "X-Token": token + "\r\nInjected: yes"}, + body=f"token={token}", + ) + + clean = redact_request(request, Route(host="api.example.com"), {}) + + self.assertTrue(clean) + self.assertNotIn(token, request.path) + self.assertNotIn(token, request.headers["X-Token"]) + self.assertNotIn("\r", request.headers["X-Token"]) + self.assertNotIn(token, request.text) + self.assertEqual("api.example.com", request.headers["Host"]) + + +if __name__ == "__main__": + unittest.main()