diff --git a/src/python/mitmproxy_addon.py b/src/python/mitmproxy_addon.py index 416e64a..e37d4ba 100644 --- a/src/python/mitmproxy_addon.py +++ b/src/python/mitmproxy_addon.py @@ -4,14 +4,16 @@ import websockets import json + async def send_ready_notification() -> None: uri = "ws://localhost:8001" - logging.info(f'sending mitmproxyReady JSON-RPC notification to {uri}') + logging.info(f"sending mitmproxyReady JSON-RPC notification to {uri}") async with websockets.connect(uri) as websocket: - notification_dto = { "jsonrpc": "2.0", "method": "mitmproxyReady" } + notification_dto = {"jsonrpc": "2.0", "method": "mitmproxyReady"} data = json.dumps(notification_dto) await websocket.send(data) + class MitmproxyAddon: def running(self) -> None: # tell the control API that we’re ready to receive traffic @@ -26,7 +28,9 @@ def request(self, flow: http.HTTPFlow) -> None: # (b'Connection', b'Upgrade'), (b'Upgrade', b'websocket') intercept = MitmproxyAddon.is_websocket_upgrade_request(flow.request) - logging.info(f'MitmproxyAddon {"intercepting" if intercept else "not intercepting"} `request` {flow.request.url}, headers {flow.request.headers}') + logging.info( + f'MitmproxyAddon {"intercepting" if intercept else "not intercepting"} `request` {flow.request.url}, headers {flow.request.headers}' + ) # pretty_host takes the "Host" header of the request into account, # which is useful in transparent mode where we usually only have the IP # otherwise. @@ -38,18 +42,26 @@ def request(self, flow: http.HTTPFlow) -> None: flow.request.host = "localhost" flow.request.port = 8002 - flow.request.scheme = 'http' + flow.request.scheme = "http" # TODO understand how port fits into this - flow.request.headers['Ably-Test-Host'] = original_host + flow.request.headers["Ably-Test-Host"] = original_host match original_scheme: - case 'http': - flow.request.headers['Ably-Test-Proto'] = 'ws' - case 'https': - flow.request.headers['Ably-Test-Proto'] = 'wss' + case "http": + flow.request.headers["Ably-Test-Proto"] = "ws" + case "https": + flow.request.headers["Ably-Test-Proto"] = "wss" @staticmethod def is_websocket_upgrade_request(request: http.Request) -> bool: # TODO this request handling is a bit fragile, the special case for `split` is just to handle the fact that Firefox sends 'Connection: keep-alive, Upgrade' - return True if 'Connection' in request.headers and ('Upgrade' in request.headers['Connection'].split(", ")) and 'Upgrade' in request.headers and request.headers['Upgrade'] == 'websocket' else False + return ( + True + if "Connection" in request.headers + and ("Upgrade" in request.headers["Connection"].split(", ")) + and "Upgrade" in request.headers + and request.headers["Upgrade"] == "websocket" + else False + ) + addons = [MitmproxyAddon()]