diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py index fa194fc1a..204d61e0a 100644 --- a/mitmproxy/addons/__init__.py +++ b/mitmproxy/addons/__init__.py @@ -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(), diff --git a/mitmproxy/addons/cut.py b/mitmproxy/addons/cut.py new file mode 100644 index 000000000..b4d9b522e --- /dev/null +++ b/mitmproxy/addons/cut.py @@ -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)) diff --git a/mitmproxy/command.py b/mitmproxy/command.py index 3b58cc257..8bf2794cf 100644 --- a/mitmproxy/command.py +++ b/mitmproxy/command.py @@ -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: diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py index 8f8f54938..898c1478e 100644 --- a/mitmproxy/tools/console/flowlist.py +++ b/mitmproxy/tools/console/flowlist.py @@ -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 diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py index 29a61bb3e..b86be7e1e 100644 --- a/mitmproxy/tools/console/master.py +++ b/mitmproxy/tools/console/master.py @@ -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") diff --git a/mitmproxy/utils/strutils.py b/mitmproxy/utils/strutils.py index 1b90c2e50..db0cfd2e8 100644 --- a/mitmproxy/utils/strutils.py +++ b/mitmproxy/utils/strutils.py @@ -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) diff --git a/mitmproxy/utils/typecheck.py b/mitmproxy/utils/typecheck.py index 20791e173..c97ff5293 100644 --- a/mitmproxy/utils/typecheck.py +++ b/mitmproxy/utils/typecheck.py @@ -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): diff --git a/test/mitmproxy/addons/test_cut.py b/test/mitmproxy/addons/test_cut.py new file mode 100644 index 000000000..5eb864c09 --- /dev/null +++ b/test/mitmproxy/addons/test_cut.py @@ -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") == [[""]] diff --git a/test/mitmproxy/test_command.py b/test/mitmproxy/test_command.py index 96d79dbae..aef05adc2 100644 --- a/test/mitmproxy/test_command.py +++ b/test/mitmproxy/test_command.py @@ -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") diff --git a/test/mitmproxy/utils/test_typecheck.py b/test/mitmproxy/utils/test_typecheck.py index 22bd7c342..17f70d37b 100644 --- a/test/mitmproxy/utils/test_typecheck.py +++ b/test/mitmproxy/utils/test_typecheck.py @@ -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)