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:
Aldo Cortesi 2017-04-29 18:33:13 +12:00
parent d439b34511
commit 4b568f99d6
10 changed files with 270 additions and 5 deletions

View File

@ -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
View 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))

View File

@ -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:

View File

@ -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

View File

@ -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")

View File

@ -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)

View File

@ -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):

View 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") == [[""]]

View File

@ -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")

View File

@ -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)