Skip to content

Commit

Permalink
ruff format
Browse files Browse the repository at this point in the history
  • Loading branch information
lawrence-forooghian committed Jul 25, 2024
1 parent 482399f commit 777fc9a
Showing 1 changed file with 22 additions and 10 deletions.
32 changes: 22 additions & 10 deletions src/python/mitmproxy_addon.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
Expand All @@ -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()]

0 comments on commit 777fc9a

Please sign in to comment.