Skip to content

Instantly share code, notes, and snippets.

@linw1995
Last active December 12, 2023 07:51
Show Gist options
  • Save linw1995/4630163c5f3fb2b575bb6fd50e89aa80 to your computer and use it in GitHub Desktop.
Save linw1995/4630163c5f3fb2b575bb6fd50e89aa80 to your computer and use it in GitHub Desktop.
A simple example shows how to hook response with mitmproxy in script.
# Standard Library
import asyncio
import contextlib
from pathlib import Path
from typing import Callable, Optional
# Third Party Library
import httpx
from mitmproxy.addons import core
from mitmproxy.http import HTTPFlow
from mitmproxy.master import Master
from mitmproxy.options import Options
from mitmproxy.proxy import ProxyConfig, ProxyServer
ResponseHook = Callable[[HTTPFlow], None]
CAFILE_PATH = Path("~/.mitmproxy/mitmproxy-ca.pem").expanduser().absolute()
class Interceptor:
def __init__(self):
self.response_hook: Optional[ResponseHook] = None
@contextlib.contextmanager
def hook_response(self, hook_function: ResponseHook):
self.response_hook = hook_function
try:
yield
finally:
self.response_hook = None
def response(self, flow: HTTPFlow):
if self.response_hook is not None:
self.response_hook(flow)
def control(interceptor: Interceptor, proxy_port: int):
def wait_resposne(flow: HTTPFlow):
# run in main thread
print("in sync control")
print("received response url:", flow.request.url)
with httpx.Client(
proxies=f"http://localhost:{proxy_port}", verify=CAFILE_PATH
) as client:
# run in another thread
with interceptor.hook_response(wait_resposne):
resp = client.get("https://httpbin.org/get")
resp.raise_for_status()
resp = client.get("https://httpbin.org/get")
resp.raise_for_status()
async def async_control(interceptor: Interceptor, proxy_port: int):
def wait_resposne(flow: HTTPFlow):
print("in async control")
print("received response url:", flow.request.url)
async with httpx.AsyncClient(
proxies=f"http://localhost:{proxy_port}",
verify=CAFILE_PATH,
) as client:
with interceptor.hook_response(wait_resposne):
resp = await client.get("https://httpbin.org/get")
resp.raise_for_status()
resp = await client.get("https://httpbin.org/get")
resp.raise_for_status()
@contextlib.asynccontextmanager
async def intercept(listen_port: int = 8080) -> None:
options = Options(listen_host="0.0.0.0", listen_port=listen_port)
master = Master(options)
master.server = server = ProxyServer(ProxyConfig(options))
# TODO: Keeps this line until the major version 7 is released
master.addons.add(core.Core())
master.start()
await master.running()
interceptor = Interceptor()
master.addons.add(interceptor)
try:
yield interceptor
finally:
# Don't use `Master.shutdown()`.
# It will make the process being unable to exit.
# And I don't know what is causing this.
master.should_exit.set()
server.shutdown()
async def main():
port = 8888
async with intercept(listen_port=port) as interceptor:
# run in other thread to avoid blocking main thread
await asyncio.to_thread(control, interceptor=interceptor, proxy_port=port)
# async function is invoked normally via await statement
await async_control(interceptor, port)
if __name__ == "__main__":
asyncio.run(main())
in sync control
received response url: http://httpbin.org/get
in async control
received response url: http://httpbin.org/get
[project]
name = ""
version = ""
description = ""
authors = [
{name = "林玮 (Jade Lin)", email = "linw1995@icloud.com"},
]
dependencies = [
"httpx~=0.17",
"mitmproxy~=6.0",
]
requires-python = ">=3.9"
dynamic = ["classifiers"]
license = {text = "MIT"}
[project.urls]
homepage = ""
[build-system]
requires = ["pdm-pep517"]
build-backend = "pdm.pep517.api"
[tool]
[tool.pdm]
@linw1995
Copy link
Author

@milahu This script has only been tested on the version "mitmproxy~=6.0."

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment