From 0740c36655b8dff151cfa980cf239c35efd335d0 Mon Sep 17 00:00:00 2001 From: Lawrence Forooghian Date: Thu, 18 Jul 2024 17:20:34 -0300 Subject: [PATCH] fixing some mypy stuff --- mitmproxy_addon_2.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/mitmproxy_addon_2.py b/mitmproxy_addon_2.py index 4ae36a9..77fbc9e 100644 --- a/mitmproxy_addon_2.py +++ b/mitmproxy_addon_2.py @@ -1,10 +1,10 @@ -import mitmproxy +from mitmproxy import http import logging import asyncio import websockets import json -async def send_ready_notification(): +async def send_ready_notification() -> None: uri = "ws://localhost:8001" logging.info(f'sending mitmproxyReady JSON-RPC notification to {uri}') async with websockets.connect(uri) as websocket: @@ -13,12 +13,12 @@ async def send_ready_notification(): await websocket.send(data) class MitmproxyAddon2: - def running(self): + def running(self) -> None: # tell the control API that we’re ready to receive traffic asyncio.get_running_loop().create_task(send_ready_notification()) # Copied from https://docs.mitmproxy.org/stable/addons-examples/#http-redirect-requests - def request(self, flow: mitmproxy.http.HTTPFlow) -> None: + def request(self, flow: http.HTTPFlow) -> None: # To make sure that when running in local redirect mode (and hence intercepting all traffic from the test process) we don’t mess with traffic from the test process to the control API # TODO see extended comments re this in test-node.yml and why it hasn’t yet been an issue in practice on macOS if not flow.request.port in [80, 443]: @@ -47,7 +47,8 @@ def request(self, flow: mitmproxy.http.HTTPFlow) -> None: case 'https': flow.request.headers['Ably-Test-Proto'] = 'wss' - def is_websocket_upgrade_request(request: mitmproxy.http.Request): + @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