mirror of
https://github.com/vee1e/mitmproxy.git
synced 2026-09-01 10:18:26 +00:00
chore: add a sample flow generator for mitmweb (#8342)
test/helper_tools/gen_sample_flows.py writes a flow file with one flow per type and state the mitmweb flow table can display, including the TCP, UDP, DNS and QUIC flows a regular HTTP proxy never sees. This makes it possible to eyeball the whole table without having to capture the traffic first. Passing --upload pushes the flows into a running mitmweb instead of writing a file.
This commit is contained in:
Sleeyax
GitHub
parent
2209baf32d
commit
43ee31cbc5
1 changed files with 284 additions and 0 deletions
284
test/helper_tools/gen_sample_flows.py
Executable file
284
test/helper_tools/gen_sample_flows.py
Executable file
|
|
@ -0,0 +1,284 @@
|
|||
#!/usr/bin/env python3
|
||||
"""Generate a flow file with one sample flow per type and state mitmweb can display.
|
||||
|
||||
python test/helper_tools/gen_sample_flows.py # writes sample-flows.mitm
|
||||
mitmweb -r sample-flows.mitm # ...then open it
|
||||
|
||||
Or push the flows straight into a running mitmweb, replacing the ones it has:
|
||||
|
||||
python test/helper_tools/gen_sample_flows.py --upload --token TOKEN
|
||||
|
||||
The token is the `?token=...` in the URL mitmweb prints on startup, or your `web_password`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import http.cookiejar
|
||||
import sys
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
from mitmproxy import io
|
||||
from mitmproxy.flow import Flow
|
||||
from mitmproxy.http import Headers
|
||||
from mitmproxy.test.tflow import tclient_conn
|
||||
from mitmproxy.test.tflow import tdnsflow
|
||||
from mitmproxy.test.tflow import terr
|
||||
from mitmproxy.test.tflow import tflow
|
||||
from mitmproxy.test.tflow import ttcpflow
|
||||
from mitmproxy.test.tflow import tudpflow
|
||||
from mitmproxy.test.tflow import twebsocket
|
||||
from mitmproxy.test.tutils import treq
|
||||
from mitmproxy.test.tutils import tresp
|
||||
|
||||
XSRF_COOKIE = "_mitmproxy_xsrf"
|
||||
|
||||
|
||||
def _req(path: str, host: str = "example.com", tls: bool = True, **kwargs):
|
||||
return treq(
|
||||
host=host,
|
||||
port=443 if tls else 80,
|
||||
scheme=b"https" if tls else b"http",
|
||||
path=path.encode(),
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _resp(status_code: int = 200, content_type: str | None = None, **headers):
|
||||
if content_type:
|
||||
headers["content_type"] = content_type
|
||||
headers.setdefault("content_length", "7")
|
||||
return tresp(status_code=status_code, headers=Headers(**headers))
|
||||
|
||||
|
||||
def _plaintext_client():
|
||||
"""A client connection that never completed a TLS handshake, so the TLS column stays grey."""
|
||||
conn = tclient_conn()
|
||||
conn.timestamp_tls_setup = None
|
||||
conn.tls = False
|
||||
return conn
|
||||
|
||||
|
||||
def _quic_client():
|
||||
conn = tclient_conn()
|
||||
conn.tls_version = "QUICv1"
|
||||
conn.tls = True
|
||||
return conn
|
||||
|
||||
|
||||
def resource_type_flows() -> list[tuple[str, Flow]]:
|
||||
"""One flow per `ResourceType`, in the order the icons are worth comparing."""
|
||||
return [
|
||||
(
|
||||
"html",
|
||||
tflow(
|
||||
req=_req("/"),
|
||||
resp=_resp(content_type="text/html; charset=utf-8"),
|
||||
),
|
||||
),
|
||||
(
|
||||
"js",
|
||||
tflow(
|
||||
req=_req("/static/app.bundle.js"),
|
||||
resp=_resp(content_type="application/javascript"),
|
||||
),
|
||||
),
|
||||
(
|
||||
"css",
|
||||
tflow(
|
||||
req=_req("/static/main.css"),
|
||||
resp=_resp(content_type="text/css"),
|
||||
),
|
||||
),
|
||||
(
|
||||
"image",
|
||||
tflow(
|
||||
req=_req("/img/logo.png"),
|
||||
resp=_resp(content_type="image/png"),
|
||||
),
|
||||
),
|
||||
(
|
||||
"plain (unknown content type)",
|
||||
tflow(
|
||||
req=_req("/api/telemetry.bin"),
|
||||
resp=_resp(content_type="application/octet-stream"),
|
||||
),
|
||||
),
|
||||
(
|
||||
"plain (still in flight)",
|
||||
tflow(req=_req("/api/slow-endpoint"), live=True),
|
||||
),
|
||||
(
|
||||
"not-modified",
|
||||
tflow(
|
||||
req=_req("/static/app.bundle.js"),
|
||||
resp=_resp(304, content_length="0"),
|
||||
),
|
||||
),
|
||||
(
|
||||
"redirect",
|
||||
tflow(
|
||||
client_conn=_plaintext_client(),
|
||||
req=_req("/old-path", tls=False),
|
||||
resp=_resp(301, location="https://example.com/new-path"),
|
||||
),
|
||||
),
|
||||
(
|
||||
"websocket",
|
||||
tflow(
|
||||
req=_req("/socket", method=b"GET"),
|
||||
resp=tresp(status_code=101, reason=b"Switching Protocols"),
|
||||
ws=twebsocket(),
|
||||
),
|
||||
),
|
||||
("tcp", ttcpflow()),
|
||||
("udp", tudpflow()),
|
||||
("dns", tdnsflow(resp=True)),
|
||||
("quic", tudpflow(client_conn=_quic_client())),
|
||||
]
|
||||
|
||||
|
||||
def flow_state_flows() -> list[tuple[str, Flow]]:
|
||||
"""Row states that draw their own icons in the path and quickactions columns."""
|
||||
intercepted = tflow(req=_req("/checkout/submit", method=b"POST"))
|
||||
intercepted.intercepted = True
|
||||
|
||||
killed = tflow(req=_req("/tracking/beacon"), err=terr("Connection killed."))
|
||||
|
||||
errored = tflow(req=_req("/unreachable"), err=terr("Server unreachable"))
|
||||
|
||||
replayed = tflow(req=_req("/api/session"), resp=_resp(content_type="text/html"))
|
||||
replayed.is_replay = "request"
|
||||
|
||||
marked = tflow(req=_req("/interesting"), resp=_resp(content_type="text/html"))
|
||||
marked.marked = ":star:"
|
||||
marked.comment = "Worth a second look"
|
||||
|
||||
return [
|
||||
("intercepted", intercepted),
|
||||
("killed", killed),
|
||||
("error", errored),
|
||||
("replayed", replayed),
|
||||
("marked + comment", marked),
|
||||
]
|
||||
|
||||
|
||||
def stagger(flows: list[Flow], base: float) -> None:
|
||||
"""Spread flows over the last few minutes so the time columns are not all identical."""
|
||||
for i, f in enumerate(flows):
|
||||
start = base + i * 3
|
||||
end = start + 0.05 * (i + 1)
|
||||
f.timestamp_created = start
|
||||
f.client_conn.timestamp_start = start
|
||||
if f.server_conn:
|
||||
f.server_conn.timestamp_start = start
|
||||
f.server_conn.timestamp_end = end
|
||||
match f.type:
|
||||
case "http":
|
||||
f.request.timestamp_start = start
|
||||
f.request.timestamp_end = start + 0.01
|
||||
if f.response:
|
||||
f.response.timestamp_start = start + 0.02
|
||||
f.response.timestamp_end = end
|
||||
if f.websocket:
|
||||
f.websocket.timestamp_end = end
|
||||
case "tcp" | "udp":
|
||||
for j, msg in enumerate(f.messages):
|
||||
msg.timestamp = start + 0.01 * (j + 1)
|
||||
case "dns":
|
||||
f.request.timestamp = start
|
||||
if f.response:
|
||||
f.response.timestamp = end
|
||||
|
||||
|
||||
def build() -> list[tuple[str, Flow]]:
|
||||
labelled = resource_type_flows() + flow_state_flows()
|
||||
stagger([f for _, f in labelled], time.time() - 300)
|
||||
return labelled
|
||||
|
||||
|
||||
def upload(data: bytes, web_url: str, token: str) -> None:
|
||||
"""POST the flows to a running mitmweb, mimicking what the UI's File > Open does."""
|
||||
jar = http.cookiejar.CookieJar()
|
||||
opener = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(jar))
|
||||
|
||||
# Only the /updates handler mints an XSRF cookie (see IndexHandler's neighbours in tools/web/app.py),
|
||||
# and it does so in prepare() so the botched websocket handshake this provokes still leaves us with both the auth and XSRF cookies.
|
||||
try:
|
||||
opener.open(f"{web_url}/updates?token={token}").read()
|
||||
except urllib.error.HTTPError as e:
|
||||
if e.code != 400:
|
||||
raise
|
||||
try:
|
||||
xsrf = next(c.value for c in jar if c.name == XSRF_COOKIE)
|
||||
except StopIteration:
|
||||
sys.exit(
|
||||
f"mitmweb did not set a {XSRF_COOKIE} cookie; is {web_url} really mitmweb?"
|
||||
)
|
||||
|
||||
request = urllib.request.Request(
|
||||
f"{web_url}/flows/dump",
|
||||
data=data,
|
||||
method="POST",
|
||||
headers={"X-XSRFToken": xsrf, "Content-Type": "application/octet-stream"},
|
||||
)
|
||||
opener.open(request).read()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
|
||||
)
|
||||
parser.add_argument(
|
||||
"-o",
|
||||
"--out",
|
||||
type=Path,
|
||||
default=Path("sample-flows.mitm"),
|
||||
help="output flow file",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--upload",
|
||||
action="store_true",
|
||||
help="also load the flows into a running mitmweb",
|
||||
)
|
||||
parser.add_argument("--token", help="mitmweb auth token, required with --upload")
|
||||
parser.add_argument(
|
||||
"--web-url",
|
||||
default="http://127.0.0.1:8081",
|
||||
help="mitmweb address (not the proxy port)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
labelled = build()
|
||||
with args.out.open("wb") as fp:
|
||||
writer = io.FlowWriter(fp)
|
||||
for _, f in labelled:
|
||||
writer.add(f)
|
||||
|
||||
for label, _ in labelled:
|
||||
print(f" {label}")
|
||||
print(f"\n{len(labelled)} flows -> {args.out}")
|
||||
|
||||
if not args.upload:
|
||||
print(f"\nOpen with: mitmweb -r {args.out}")
|
||||
print("or in a running mitmweb via File > Open.")
|
||||
return
|
||||
|
||||
if not args.token:
|
||||
sys.exit("--upload needs --token (the ?token=... mitmweb prints at startup)")
|
||||
try:
|
||||
upload(args.out.read_bytes(), args.web_url.rstrip("/"), args.token)
|
||||
except urllib.error.HTTPError as e:
|
||||
sys.exit(
|
||||
f"upload failed: {e.code} {e.reason} (a 403 usually means a stale token)"
|
||||
)
|
||||
except urllib.error.URLError as e:
|
||||
sys.exit(f"could not reach mitmweb at {args.web_url}: {e.reason}")
|
||||
print(f"\nUploaded to {args.web_url}, replacing the flows it had.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Add table
Add a link
Reference in a new issue