mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-26 18:18:25 +00:00
Introduce cuts: a flow dissector
This PR introduces the cuts addon, a flow dissector that allows you to select and operate on specific components of flows. It also adds the first consumer for cuts - the cuts.save command. Save the content of the focus to /tmp/foo: cuts.save s.content|@focus /tmp/foo Save the URL and response content-type headers for all flows currently shown to file, comma-separated, one flow per line: cuts.save s.url,q.header[content-type]|@focus /tmp/foo We also use this to replace the body save shortcut in the console flowlist.
This commit is contained in:
parent
d439b34511
commit
4b568f99d6
@ -5,6 +5,7 @@ from mitmproxy.addons import check_ca
|
||||
from mitmproxy.addons import clientplayback
|
||||
from mitmproxy.addons import core_option_validation
|
||||
from mitmproxy.addons import core
|
||||
from mitmproxy.addons import cut
|
||||
from mitmproxy.addons import disable_h2c
|
||||
from mitmproxy.addons import onboarding
|
||||
from mitmproxy.addons import proxyauth
|
||||
@ -28,6 +29,7 @@ def default_addons():
|
||||
check_alpn.CheckALPN(),
|
||||
check_ca.CheckCA(),
|
||||
clientplayback.ClientPlayback(),
|
||||
cut.Cut(),
|
||||
disable_h2c.DisableH2C(),
|
||||
onboarding.Onboarding(),
|
||||
proxyauth.ProxyAuth(),
|
||||
|
112
mitmproxy/addons/cut.py
Normal file
112
mitmproxy/addons/cut.py
Normal file
@ -0,0 +1,112 @@
|
||||
import os
|
||||
import typing
|
||||
from mitmproxy import command
|
||||
from mitmproxy import exceptions
|
||||
from mitmproxy import flow
|
||||
from mitmproxy import ctx
|
||||
from mitmproxy.utils import strutils
|
||||
|
||||
|
||||
def headername(spec: str):
|
||||
if not (spec.startswith("header[") and spec.endswith("]")):
|
||||
raise exceptions.CommandError("Invalid header spec: %s" % spec)
|
||||
return spec[len("header["):-1].strip()
|
||||
|
||||
|
||||
def extract(cut: str, f: flow.Flow) -> typing.Union[str, bytes]:
|
||||
if cut.startswith("q."):
|
||||
req = getattr(f, "request", None)
|
||||
if not req:
|
||||
return ""
|
||||
rem = cut[len("q."):]
|
||||
if rem in ["method", "scheme", "host", "port", "path", "url"]:
|
||||
return str(getattr(req, rem))
|
||||
elif rem == "content":
|
||||
return req.content
|
||||
elif rem.startswith("header["):
|
||||
return req.headers.get(headername(rem), "")
|
||||
elif cut.startswith("s."):
|
||||
resp = getattr(f, "response", None)
|
||||
if not resp:
|
||||
return ""
|
||||
rem = cut[len("s."):]
|
||||
if rem in ["status_code", "reason"]:
|
||||
return str(getattr(resp, rem))
|
||||
elif rem == "content":
|
||||
return resp.content
|
||||
elif rem.startswith("header["):
|
||||
return resp.headers.get(headername(rem), "")
|
||||
raise exceptions.CommandError("Invalid cut specification: %s" % cut)
|
||||
|
||||
|
||||
def parse_cutspec(s: str) -> typing.Tuple[str, typing.Sequence[str]]:
|
||||
"""
|
||||
Returns (flowspec, [cuts]).
|
||||
|
||||
Raises exceptions.CommandError if input is invalid.
|
||||
"""
|
||||
parts = s.split("|", maxsplit=1)
|
||||
flowspec = "@all"
|
||||
if len(parts) == 2:
|
||||
flowspec = parts[1].strip()
|
||||
cuts = parts[0]
|
||||
cutparts = [i.strip() for i in cuts.split(",") if i.strip()]
|
||||
if len(cutparts) == 0:
|
||||
raise exceptions.CommandError("Invalid cut specification.")
|
||||
return flowspec, cutparts
|
||||
|
||||
|
||||
class Cut:
|
||||
@command.command("cut")
|
||||
def cut(self, cutspec: str) -> command.Cuts:
|
||||
"""
|
||||
Resolve a cut specification of the form "cuts|flowspec". The
|
||||
flowspec is optional, and if it is not specified, it is assumed to
|
||||
be @all. The cuts are a comma-separated list of cut snippets.
|
||||
|
||||
HTTP requests: q.method, q.scheme, q.host, q.port, q.path, q.url,
|
||||
q.header[key], q.content
|
||||
|
||||
HTTP responses: s.status_code, s.reason, s.header[key], s.content
|
||||
|
||||
Client connections: cc.address, cc.sni, cc.cipher_name,
|
||||
cc.alpn_proto, cc.tls_version
|
||||
|
||||
Server connections: sc.address, sc.ip, sc.cert, sc.sni,
|
||||
sc.alpn_proto, sc.tls_version
|
||||
"""
|
||||
flowspec, cuts = parse_cutspec(cutspec)
|
||||
flows = ctx.master.commands.call_args("view.resolve", [flowspec])
|
||||
ret = []
|
||||
for f in flows:
|
||||
ret.append([extract(c, f) for c in cuts])
|
||||
return ret
|
||||
|
||||
@command.command("cut.save")
|
||||
def save(self, cuts: command.Cuts, path: str) -> None:
|
||||
"""
|
||||
Save cuts to file.
|
||||
|
||||
cut.save resp.content|@focus /tmp/foo
|
||||
|
||||
cut.save req.host,resp.header[content-type]|@focus /tmp/foo
|
||||
"""
|
||||
mode = "wb"
|
||||
if path.startswith("+"):
|
||||
mode = "ab"
|
||||
path = path[1:]
|
||||
path = os.path.expanduser(path)
|
||||
with open(path, mode) as fp:
|
||||
if fp.tell() > 0:
|
||||
# We're appending to a file that already exists and has content
|
||||
fp.write(b"\n")
|
||||
for ci, c in enumerate(cuts):
|
||||
if ci > 0:
|
||||
fp.write(b"\n")
|
||||
for vi, v in enumerate(c):
|
||||
if vi > 0:
|
||||
fp.write(b", ")
|
||||
if isinstance(v, str):
|
||||
v = strutils.always_bytes(v)
|
||||
fp.write(v)
|
||||
ctx.log.alert("Saved %s cuts." % len(cuts))
|
@ -9,6 +9,11 @@ from mitmproxy import exceptions
|
||||
from mitmproxy import flow
|
||||
|
||||
|
||||
Cuts = typing.Sequence[
|
||||
typing.Sequence[typing.Union[str, bytes]]
|
||||
]
|
||||
|
||||
|
||||
def typename(t: type, ret: bool) -> str:
|
||||
"""
|
||||
Translates a type to an explanatory string. Ifl ret is True, we're
|
||||
@ -18,6 +23,8 @@ def typename(t: type, ret: bool) -> str:
|
||||
return t.__name__
|
||||
elif t == typing.Sequence[flow.Flow]:
|
||||
return "[flow]" if ret else "flowspec"
|
||||
elif t == Cuts:
|
||||
return "[cuts]" if ret else "cutspec"
|
||||
elif t == flow.Flow:
|
||||
return "flow"
|
||||
else: # pragma: no cover
|
||||
@ -125,6 +132,8 @@ def parsearg(manager: CommandManager, spec: str, argtype: type) -> typing.Any:
|
||||
raise exceptions.CommandError("Expected an integer, got %s." % spec)
|
||||
elif argtype == typing.Sequence[flow.Flow]:
|
||||
return manager.call_args("view.resolve", [spec])
|
||||
elif argtype == Cuts:
|
||||
return manager.call_args("cut", [spec])
|
||||
elif argtype == flow.Flow:
|
||||
flows = manager.call_args("view.resolve", [spec])
|
||||
if len(flows) != 1:
|
||||
|
@ -157,8 +157,6 @@ class FlowItem(urwid.WidgetWrap):
|
||||
# callback = common.export_to_clip_or_file,
|
||||
# args = (None, self.flow, common.copy_to_clipboard_or_prompt)
|
||||
# )
|
||||
elif key == "b":
|
||||
common.ask_save_body(None, self.flow)
|
||||
else:
|
||||
return key
|
||||
|
||||
|
@ -151,6 +151,7 @@ def default_keymap(km):
|
||||
|
||||
km.add("A", "flow.resume @all", context="flowlist")
|
||||
km.add("a", "flow.resume @focus", context="flowlist")
|
||||
km.add("b", "console.command 'cut.save s.content|@focus '", context="flowlist")
|
||||
km.add("d", "view.remove @focus", context="flowlist")
|
||||
km.add("D", "view.duplicate @focus", context="flowlist")
|
||||
km.add("e", "set console_eventlog=toggle", context="flowlist")
|
||||
|
@ -25,9 +25,10 @@ def always_str(str_or_bytes: Optional[AnyStr], *decode_args) -> Optional[str]:
|
||||
raise TypeError("Expected str or bytes, but got {}.".format(type(str_or_bytes).__name__))
|
||||
|
||||
|
||||
# Translate control characters to "safe" characters. This implementation initially
|
||||
# replaced them with the matching control pictures (http://unicode.org/charts/PDF/U2400.pdf),
|
||||
# but that turned out to render badly with monospace fonts. We are back to "." therefore.
|
||||
# Translate control characters to "safe" characters. This implementation
|
||||
# initially replaced them with the matching control pictures
|
||||
# (http://unicode.org/charts/PDF/U2400.pdf), but that turned out to render badly
|
||||
# with monospace fonts. We are back to "." therefore.
|
||||
_control_char_trans = {
|
||||
x: ord(".") # x + 0x2400 for unicode control group pictures
|
||||
for x in range(32)
|
||||
|
@ -19,6 +19,16 @@ def check_command_return_type(value: typing.Any, typeinfo: typing.Any) -> bool:
|
||||
for v in value:
|
||||
if not check_command_return_type(v, T):
|
||||
return False
|
||||
elif typename.startswith("typing.Union"):
|
||||
try:
|
||||
types = typeinfo.__args__ # type: ignore
|
||||
except AttributeError:
|
||||
# Python 3.5.x
|
||||
types = typeinfo.__union_params__ # type: ignore
|
||||
for T in types:
|
||||
checks = [check_command_return_type(value, T) for T in types]
|
||||
if not any(checks):
|
||||
return False
|
||||
elif value is None and typeinfo is None:
|
||||
return True
|
||||
elif not isinstance(value, typeinfo):
|
||||
|
109
test/mitmproxy/addons/test_cut.py
Normal file
109
test/mitmproxy/addons/test_cut.py
Normal file
@ -0,0 +1,109 @@
|
||||
|
||||
from mitmproxy.addons import cut
|
||||
from mitmproxy.addons import view
|
||||
from mitmproxy import exceptions
|
||||
from mitmproxy.test import taddons
|
||||
from mitmproxy.test import tflow
|
||||
import pytest
|
||||
|
||||
|
||||
def test_parse_cutspec():
|
||||
tests = [
|
||||
("", None, True),
|
||||
("req.method", ("@all", ["req.method"]), False),
|
||||
(
|
||||
"req.method,req.host",
|
||||
("@all", ["req.method", "req.host"]),
|
||||
False
|
||||
),
|
||||
(
|
||||
"req.method,req.host|~b foo",
|
||||
("~b foo", ["req.method", "req.host"]),
|
||||
False
|
||||
),
|
||||
(
|
||||
"req.method,req.host|~b foo | ~b bar",
|
||||
("~b foo | ~b bar", ["req.method", "req.host"]),
|
||||
False
|
||||
),
|
||||
(
|
||||
"req.method, req.host | ~b foo | ~b bar",
|
||||
("~b foo | ~b bar", ["req.method", "req.host"]),
|
||||
False
|
||||
),
|
||||
]
|
||||
for cutspec, output, err in tests:
|
||||
try:
|
||||
assert cut.parse_cutspec(cutspec) == output
|
||||
except exceptions.CommandError:
|
||||
if not err:
|
||||
raise
|
||||
else:
|
||||
if err:
|
||||
raise AssertionError("Expected error.")
|
||||
|
||||
|
||||
def test_headername():
|
||||
with pytest.raises(exceptions.CommandError):
|
||||
cut.headername("header[foo.")
|
||||
|
||||
|
||||
def qr(f):
|
||||
with open(f, "rb") as fp:
|
||||
return fp.read()
|
||||
|
||||
|
||||
def test_cut_file(tmpdir):
|
||||
f = str(tmpdir.join("path"))
|
||||
v = view.View()
|
||||
c = cut.Cut()
|
||||
with taddons.context() as tctx:
|
||||
tctx.master.addons.add(v, c)
|
||||
|
||||
v.add([tflow.tflow(resp=True)])
|
||||
|
||||
tctx.command(c.save, "q.method|@all", f)
|
||||
assert qr(f) == b"GET"
|
||||
tctx.command(c.save, "q.content|@all", f)
|
||||
assert qr(f) == b"content"
|
||||
tctx.command(c.save, "q.content|@all", "+" + f)
|
||||
assert qr(f) == b"content\ncontent"
|
||||
|
||||
v.add([tflow.tflow(resp=True)])
|
||||
tctx.command(c.save, "q.method|@all", f)
|
||||
assert qr(f) == b"GET\nGET"
|
||||
tctx.command(c.save, "q.method,q.path|@all", f)
|
||||
assert qr(f) == b"GET, /path\nGET, /path"
|
||||
|
||||
|
||||
def test_cut():
|
||||
v = view.View()
|
||||
c = cut.Cut()
|
||||
with taddons.context() as tctx:
|
||||
v.add([tflow.tflow(resp=True)])
|
||||
tctx.master.addons.add(v, c)
|
||||
assert c.cut("q.method|@all") == [["GET"]]
|
||||
assert c.cut("q.scheme|@all") == [["http"]]
|
||||
assert c.cut("q.host|@all") == [["address"]]
|
||||
assert c.cut("q.port|@all") == [["22"]]
|
||||
assert c.cut("q.path|@all") == [["/path"]]
|
||||
assert c.cut("q.url|@all") == [["http://address:22/path"]]
|
||||
assert c.cut("q.content|@all") == [[b"content"]]
|
||||
assert c.cut("q.header[header]|@all") == [["qvalue"]]
|
||||
assert c.cut("q.header[unknown]|@all") == [[""]]
|
||||
|
||||
assert c.cut("s.status_code|@all") == [["200"]]
|
||||
assert c.cut("s.reason|@all") == [["OK"]]
|
||||
assert c.cut("s.content|@all") == [[b"message"]]
|
||||
assert c.cut("s.header[header-response]|@all") == [["svalue"]]
|
||||
|
||||
with pytest.raises(exceptions.CommandError):
|
||||
assert c.cut("moo") == [["svalue"]]
|
||||
|
||||
v = view.View()
|
||||
c = cut.Cut()
|
||||
with taddons.context() as tctx:
|
||||
tctx.master.addons.add(v, c)
|
||||
v.add([tflow.ttcpflow()])
|
||||
assert c.cut("q.method|@all") == [[""]]
|
||||
assert c.cut("s.status|@all") == [[""]]
|
@ -60,17 +60,25 @@ def test_typename():
|
||||
assert command.typename(str, True) == "str"
|
||||
assert command.typename(typing.Sequence[flow.Flow], True) == "[flow]"
|
||||
assert command.typename(typing.Sequence[flow.Flow], False) == "flowspec"
|
||||
|
||||
assert command.typename(command.Cuts, False) == "cutspec"
|
||||
assert command.typename(command.Cuts, True) == "[cuts]"
|
||||
|
||||
assert command.typename(flow.Flow, False) == "flow"
|
||||
|
||||
|
||||
class DummyConsole:
|
||||
def load(self, l):
|
||||
l.add_command("view.resolve", self.resolve)
|
||||
l.add_command("cut", self.cut)
|
||||
|
||||
def resolve(self, spec: str) -> typing.Sequence[flow.Flow]:
|
||||
n = int(spec)
|
||||
return [tflow.tflow(resp=True)] * n
|
||||
|
||||
def cut(self, spec: str) -> command.Cuts:
|
||||
return [["test"]]
|
||||
|
||||
|
||||
def test_parsearg():
|
||||
with taddons.context() as tctx:
|
||||
@ -97,6 +105,10 @@ def test_parsearg():
|
||||
with pytest.raises(exceptions.CommandError):
|
||||
command.parsearg(tctx.master.commands, "foo", Exception)
|
||||
|
||||
assert command.parsearg(
|
||||
tctx.master.commands, "foo", command.Cuts
|
||||
) == [["test"]]
|
||||
|
||||
|
||||
class TDec:
|
||||
@command.command("cmd1")
|
||||
|
@ -4,6 +4,7 @@ from unittest import mock
|
||||
import pytest
|
||||
|
||||
from mitmproxy.utils import typecheck
|
||||
from mitmproxy import command
|
||||
|
||||
|
||||
class TBase:
|
||||
@ -93,9 +94,19 @@ def test_check_command_return_type():
|
||||
assert(typecheck.check_command_return_type(None, None))
|
||||
assert(not typecheck.check_command_return_type(["foo"], typing.Sequence[int]))
|
||||
assert(not typecheck.check_command_return_type("foo", typing.Sequence[int]))
|
||||
assert(typecheck.check_command_return_type([["foo", b"bar"]], command.Cuts))
|
||||
assert(not typecheck.check_command_return_type(["foo", b"bar"], command.Cuts))
|
||||
assert(not typecheck.check_command_return_type([["foo", 22]], command.Cuts))
|
||||
|
||||
# Python 3.5 only defines __parameters__
|
||||
m = mock.Mock()
|
||||
m.__str__ = lambda self: "typing.Sequence"
|
||||
m.__parameters__ = (int,)
|
||||
|
||||
typecheck.check_command_return_type([10], m)
|
||||
|
||||
# Python 3.5 only defines __union_params__
|
||||
m = mock.Mock()
|
||||
m.__str__ = lambda self: "typing.Union"
|
||||
m.__union_params__ = (int,)
|
||||
assert not typecheck.check_command_return_type([22], m)
|
||||
|
Loading…
Reference in New Issue
Block a user