feat: add mitm_launcher.py and explicit scan_failed degradation in addon
- addon.py: split except block into specific httpx network errors (ConnectError, TimeoutException, HTTPStatusError) and generic Exception; both now log action="scan_failed" instead of "allow", and redact mode calls flow.kill() to block the upload when info-privacy is unreachable - mitm_launcher.py: programmatic DumpMaster entry point for PyInstaller, replaces `mitmdump -s addon.py` file-load so all source compiles to binary - tests/test_addon_degradation.py: 8 parametrized tests covering audit/redact degradation with ConnectError, TimeoutException, HTTPStatusError, and _audit=None edge cases Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -132,14 +132,33 @@ class PrivacyGatewayAddon:
|
||||
for f in files:
|
||||
try:
|
||||
result = await scan_and_redact(f.filename, f.data)
|
||||
except Exception as e:
|
||||
ctx.log.warn(f"[privacy-gw] scan failed {f.filename}: {e}")
|
||||
except (httpx.ConnectError, httpx.TimeoutException, httpx.HTTPStatusError) as e:
|
||||
logger.warning("[privacy-gw] info-privacy unreachable for %s: %s", f.filename, e)
|
||||
ctx.log.warn(f"[privacy-gw] scan_failed {f.filename}: {e}")
|
||||
if _audit is not None:
|
||||
_audit.log(
|
||||
host, {}, "allow", f.data, total_files,
|
||||
host, {}, "scan_failed", f.data, total_files,
|
||||
client_ip=client_ip, request_url=request_url,
|
||||
filename=f.filename, file_size=len(f.data),
|
||||
)
|
||||
if mode == "redact":
|
||||
# Block upload when info-privacy is unavailable in redact mode
|
||||
flow.kill()
|
||||
return
|
||||
# audit or off mode: allow through with scan_failed logged
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.warning("[privacy-gw] unexpected scan error for %s: %s", f.filename, e)
|
||||
ctx.log.warn(f"[privacy-gw] scan error {f.filename}: {e}")
|
||||
if _audit is not None:
|
||||
_audit.log(
|
||||
host, {}, "scan_failed", f.data, total_files,
|
||||
client_ip=client_ip, request_url=request_url,
|
||||
filename=f.filename, file_size=len(f.data),
|
||||
)
|
||||
if mode == "redact":
|
||||
flow.kill()
|
||||
return
|
||||
continue
|
||||
|
||||
if mode == "audit":
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
"""Programmatic mitmproxy entry point for PyInstaller compilation.
|
||||
|
||||
Replaces `mitmdump -s addon.py` file-load so all source is compiled into binary.
|
||||
Usage:
|
||||
python -m privacy_gateway.mitm_launcher
|
||||
# or after PyInstaller packaging:
|
||||
./kvm-mitm
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import signal
|
||||
|
||||
from mitmproxy import options
|
||||
from mitmproxy.tools.dump import DumpMaster
|
||||
|
||||
from privacy_gateway.addon import PrivacyGatewayAddon
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
opts = options.Options(
|
||||
listen_host="0.0.0.0",
|
||||
listen_port=8888,
|
||||
mode=["regular"],
|
||||
ssl_insecure=True,
|
||||
)
|
||||
master = DumpMaster(opts, with_termlog=False, with_dumper=False)
|
||||
master.addons.add(PrivacyGatewayAddon())
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
for sig in (signal.SIGINT, signal.SIGTERM):
|
||||
loop.add_signal_handler(sig, master.shutdown)
|
||||
|
||||
await master.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,128 @@
|
||||
"""Test that addon handles info-privacy being down gracefully (scan_failed degradation).
|
||||
|
||||
When scan_and_redact raises a network error:
|
||||
- audit mode : audit log called with action="scan_failed", request allowed through
|
||||
- redact mode: audit log called with action="scan_failed", flow killed/blocked
|
||||
"""
|
||||
import asyncio
|
||||
import pytest
|
||||
from unittest.mock import MagicMock, patch, AsyncMock
|
||||
|
||||
import httpx
|
||||
|
||||
import services.privacy_gateway.addon as _addon_mod
|
||||
|
||||
|
||||
def _make_mock_flow(ip="127.0.0.1"):
|
||||
flow = MagicMock()
|
||||
flow.client_conn.peername = (ip, 9999)
|
||||
flow.request.url = "https://api.openai.com/v1/files"
|
||||
flow.request.content = b"original"
|
||||
return flow
|
||||
|
||||
|
||||
def _make_mock_file(filename="secret.pdf", data=b"bytes"):
|
||||
f = MagicMock()
|
||||
f.filename = filename
|
||||
f.data = data
|
||||
f.field_name = "file"
|
||||
return f
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helper: build ScanError exceptions that map to httpx network failures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_NETWORK_ERRORS = [
|
||||
httpx.ConnectError("Connection refused"),
|
||||
httpx.TimeoutException("Timeout"),
|
||||
httpx.HTTPStatusError("503", request=MagicMock(), response=MagicMock()),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("exc", _NETWORK_ERRORS)
|
||||
async def test_audit_mode_scan_failed_logs_and_allows(exc):
|
||||
"""In audit mode, scan failure → action='scan_failed', request NOT blocked."""
|
||||
mock_audit = MagicMock()
|
||||
|
||||
with patch.object(_addon_mod, "_audit", mock_audit), \
|
||||
patch.object(_addon_mod, "scan_and_redact", AsyncMock(side_effect=exc)), \
|
||||
patch.object(_addon_mod, "ctx", MagicMock()):
|
||||
|
||||
addon = _addon_mod.PrivacyGatewayAddon()
|
||||
flow = _make_mock_flow()
|
||||
mock_file = _make_mock_file()
|
||||
|
||||
await addon._process_upload(flow, "api.openai.com", [mock_file], "audit")
|
||||
|
||||
# Must log with action="scan_failed" (not "allow")
|
||||
mock_audit.log.assert_called_once()
|
||||
call_args = mock_audit.log.call_args
|
||||
# positional: domain, pii_types, action, raw_bytes?, file_count?
|
||||
action = call_args[0][2] if len(call_args[0]) > 2 else call_args[1].get("action")
|
||||
assert action == "scan_failed", (
|
||||
f"Expected action='scan_failed' in audit mode, got '{action}'"
|
||||
)
|
||||
# Request must NOT be blocked/killed
|
||||
flow.kill.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("exc", _NETWORK_ERRORS)
|
||||
async def test_redact_mode_scan_failed_blocks_request(exc):
|
||||
"""In redact mode, scan failure → action='scan_failed', flow.kill() called."""
|
||||
mock_audit = MagicMock()
|
||||
|
||||
with patch.object(_addon_mod, "_audit", mock_audit), \
|
||||
patch.object(_addon_mod, "scan_and_redact", AsyncMock(side_effect=exc)), \
|
||||
patch.object(_addon_mod, "ctx", MagicMock()):
|
||||
|
||||
addon = _addon_mod.PrivacyGatewayAddon()
|
||||
flow = _make_mock_flow()
|
||||
mock_file = _make_mock_file()
|
||||
|
||||
await addon._process_upload(flow, "api.openai.com", [mock_file], "redact")
|
||||
|
||||
# Must log with action="scan_failed"
|
||||
mock_audit.log.assert_called_once()
|
||||
call_args = mock_audit.log.call_args
|
||||
action = call_args[0][2] if len(call_args[0]) > 2 else call_args[1].get("action")
|
||||
assert action == "scan_failed", (
|
||||
f"Expected action='scan_failed' in redact mode, got '{action}'"
|
||||
)
|
||||
# Request MUST be blocked
|
||||
flow.kill.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audit_mode_scan_failed_no_audit_object():
|
||||
"""If _audit is None, scan failure in audit mode should not raise."""
|
||||
with patch.object(_addon_mod, "_audit", None), \
|
||||
patch.object(_addon_mod, "scan_and_redact", AsyncMock(
|
||||
side_effect=httpx.ConnectError("down"))), \
|
||||
patch.object(_addon_mod, "ctx", MagicMock()):
|
||||
|
||||
addon = _addon_mod.PrivacyGatewayAddon()
|
||||
flow = _make_mock_flow()
|
||||
mock_file = _make_mock_file()
|
||||
|
||||
# Must not raise even if _audit is None
|
||||
await addon._process_upload(flow, "api.openai.com", [mock_file], "audit")
|
||||
flow.kill.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redact_mode_scan_failed_no_audit_object():
|
||||
"""If _audit is None, scan failure in redact mode should still kill the flow."""
|
||||
with patch.object(_addon_mod, "_audit", None), \
|
||||
patch.object(_addon_mod, "scan_and_redact", AsyncMock(
|
||||
side_effect=httpx.ConnectError("down"))), \
|
||||
patch.object(_addon_mod, "ctx", MagicMock()):
|
||||
|
||||
addon = _addon_mod.PrivacyGatewayAddon()
|
||||
flow = _make_mock_flow()
|
||||
mock_file = _make_mock_file()
|
||||
|
||||
await addon._process_upload(flow, "api.openai.com", [mock_file], "redact")
|
||||
flow.kill.assert_called_once()
|
||||
Reference in New Issue
Block a user