command cuts: add completion

- Remove shortcuts for request, response, etc. - we don't need them if we have completion
- Restrict cuts specification to a set of prefixes
- Extend cuts to add a few more items
This commit is contained in:
Aldo Cortesi 2017-12-16 10:30:08 +13:00
parent 50a94db2cc
commit cd913d598d
10 changed files with 141 additions and 88 deletions

View File

@ -17,14 +17,6 @@ def headername(spec: str):
return spec[len("header["):-1].strip()
flow_shortcuts = {
"q": "request",
"s": "response",
"cc": "client_conn",
"sc": "server_conn",
}
def is_addr(v):
return isinstance(v, tuple) and len(v) > 1
@ -35,8 +27,6 @@ def extract(cut: str, f: flow.Flow) -> typing.Union[str, bytes]:
for i, spec in enumerate(path):
if spec.startswith("_"):
raise exceptions.CommandError("Can't access internal attribute %s" % spec)
if isinstance(current, flow.Flow):
spec = flow_shortcuts.get(spec, spec)
part = getattr(current, spec, None)
if i == len(path) - 1:
@ -65,13 +55,12 @@ class Cut:
) -> command.Cuts:
"""
Cut data from a set of flows. Cut specifications are attribute paths
from the base of the flow object, with a few conveniences - "q",
"s", "cc" and "sc" are shortcuts for request, response, client_conn
and server_conn, "port" and "host" retrieve parts of an address
tuple, ".header[key]" retrieves a header value. Return values
converted to strings or bytes: SSL certicates are converted to PEM
format, bools are "true" or "false", "bytes" are preserved, and all
other values are converted to strings.
from the base of the flow object, with a few conveniences - "port"
and "host" retrieve parts of an address tuple, ".header[key]"
retrieves a header value. Return values converted to strings or
bytes: SSL certicates are converted to PEM format, bools are "true"
or "false", "bytes" are preserved, and all other values are
converted to strings.
"""
ret = []
for f in flows:
@ -101,7 +90,7 @@ class Cut:
if fp.tell() > 0:
# We're appending to a file that already exists and has content
fp.write(b"\n")
for v in [extract(cuts[0], f) for f in flows]:
v = extract(cuts[0], flows[0])
if isinstance(v, bytes):
fp.write(v)
else:
@ -129,8 +118,8 @@ class Cut:
column, the data is written to file as-is, with raw bytes preserved.
"""
fp = io.StringIO(newline="")
if len(cuts) == 1 and len(cuts[0]) == 1:
v = cuts[0][0]
if len(cuts) == 1 and len(flows) == 1:
v = extract(cuts[0], flows[0])
if isinstance(v, bytes):
fp.write(strutils.always_str(v))
else:
@ -138,9 +127,10 @@ class Cut:
ctx.log.alert("Clipped single cut.")
else:
writer = csv.writer(fp)
for r in cuts:
for f in flows:
vals = [extract(c, f) for c in cuts]
writer.writerow(
[strutils.always_str(c) or "" for c in r] # type: ignore
[strutils.always_str(v) or "" for v in vals] # type: ignore
)
ctx.log.alert("Clipped %s cuts as CSV." % len(cuts))
pyperclip.copy(fp.getvalue())

View File

@ -30,7 +30,46 @@ Cuts = typing.Sequence[
class Cut(str):
pass
# This is an awkward location for these values, but it's better than having
# the console core import and depend on an addon. FIXME: Add a way for
# addons to add custom types and manage their completion and validation.
valid_prefixes = [
"request.method",
"request.scheme",
"request.host",
"request.http_version",
"request.port",
"request.path",
"request.url",
"request.text",
"request.content",
"request.raw_content",
"request.timestamp_start",
"request.timestamp_end",
"request.header[",
"response.status_code",
"response.reason",
"response.text",
"response.content",
"response.timestamp_start",
"response.timestamp_end",
"response.raw_content",
"response.header[",
"client_conn.address.port",
"client_conn.address.host",
"client_conn.tls_version",
"client_conn.sni",
"client_conn.ssl_established",
"server_conn.address.port",
"server_conn.address.host",
"server_conn.ip_address.host",
"server_conn.tls_version",
"server_conn.sni",
"server_conn.ssl_established",
]
class Path(str):

View File

@ -53,6 +53,8 @@ def twebsocketflow(client_conn=True, server_conn=True, messages=True, err=None,
sec_websocket_version="13",
sec_websocket_key="1234",
),
timestamp_start=1,
timestamp_end=2,
content=b''
)
resp = http.HTTPResponse(
@ -64,6 +66,8 @@ def twebsocketflow(client_conn=True, server_conn=True, messages=True, err=None,
upgrade='websocket',
sec_websocket_accept=b'',
),
timestamp_start=1,
timestamp_end=2,
content=b'',
)
handshake_flow = http.HTTPFlow(client_conn, server_conn)

View File

@ -31,7 +31,9 @@ def treq(**kwargs):
path=b"/path",
http_version=b"HTTP/1.1",
headers=http.Headers(((b"header", b"qvalue"), (b"content-length", b"7"))),
content=b"content"
content=b"content",
timestamp_start=1,
timestamp_end=2,
)
default.update(kwargs)
return http.Request(**default)
@ -48,8 +50,8 @@ def tresp(**kwargs):
reason=b"OK",
headers=http.Headers(((b"header-response", b"svalue"), (b"content-length", b"7"))),
content=b"message",
timestamp_start=time.time(),
timestamp_end=time.time(),
timestamp_start=1,
timestamp_end=2,
)
default.update(kwargs)
return http.Response(**default)

View File

@ -113,6 +113,19 @@ class CommandBuffer():
),
parse = parts,
)
if last.type == typing.Sequence[mitmproxy.command.Cut]:
spec = parts[-1].value.split(",")
opts = []
for pref in mitmproxy.command.Cut.valid_prefixes:
spec[-1] = pref
opts.append(",".join(spec))
self.completion = CompletionState(
completer = ListCompleter(
parts[-1].value,
opts,
),
parse = parts,
)
elif isinstance(last.type, mitmproxy.command.Choice):
self.completion = CompletionState(
completer = ListCompleter(

View File

@ -31,7 +31,7 @@ def map(km):
km.add("A", "flow.resume @all", ["flowlist", "flowview"], "Resume all intercepted flows")
km.add("a", "flow.resume @focus", ["flowlist", "flowview"], "Resume this intercepted flow")
km.add(
"b", "console.command cut.save @focus s.content ",
"b", "console.command cut.save @focus response.content ",
["flowlist", "flowview"],
"Save response body to file"
)

View File

@ -13,36 +13,41 @@ from unittest import mock
def test_extract():
tf = tflow.tflow(resp=True)
tests = [
["q.method", "GET"],
["q.scheme", "http"],
["q.host", "address"],
["q.port", "22"],
["q.path", "/path"],
["q.url", "http://address:22/path"],
["q.text", "content"],
["q.content", b"content"],
["q.raw_content", b"content"],
["q.header[header]", "qvalue"],
["request.method", "GET"],
["request.scheme", "http"],
["request.host", "address"],
["request.http_version", "HTTP/1.1"],
["request.port", "22"],
["request.path", "/path"],
["request.url", "http://address:22/path"],
["request.text", "content"],
["request.content", b"content"],
["request.raw_content", b"content"],
["request.timestamp_start", "1"],
["request.timestamp_end", "2"],
["request.header[header]", "qvalue"],
["s.status_code", "200"],
["s.reason", "OK"],
["s.text", "message"],
["s.content", b"message"],
["s.raw_content", b"message"],
["s.header[header-response]", "svalue"],
["response.status_code", "200"],
["response.reason", "OK"],
["response.text", "message"],
["response.content", b"message"],
["response.raw_content", b"message"],
["response.header[header-response]", "svalue"],
["response.timestamp_start", "1"],
["response.timestamp_end", "2"],
["cc.address.port", "22"],
["cc.address.host", "127.0.0.1"],
["cc.tls_version", "TLSv1.2"],
["cc.sni", "address"],
["cc.ssl_established", "false"],
["client_conn.address.port", "22"],
["client_conn.address.host", "127.0.0.1"],
["client_conn.tls_version", "TLSv1.2"],
["client_conn.sni", "address"],
["client_conn.ssl_established", "false"],
["sc.address.port", "22"],
["sc.address.host", "address"],
["sc.ip_address.host", "192.168.0.1"],
["sc.tls_version", "TLSv1.2"],
["sc.sni", "address"],
["sc.ssl_established", "false"],
["server_conn.address.port", "22"],
["server_conn.address.host", "address"],
["server_conn.ip_address.host", "192.168.0.1"],
["server_conn.tls_version", "TLSv1.2"],
["server_conn.sni", "address"],
["server_conn.ssl_established", "false"],
]
for t in tests:
ret = cut.extract(t[0], tf)
@ -53,7 +58,7 @@ def test_extract():
d = f.read()
c1 = certs.SSLCert.from_pem(d)
tf.server_conn.cert = c1
assert "CERTIFICATE" in cut.extract("sc.cert", tf)
assert "CERTIFICATE" in cut.extract("server_conn.cert", tf)
def test_headername():
@ -74,15 +79,15 @@ def test_cut_clip():
v.add([tflow.tflow(resp=True)])
with mock.patch('pyperclip.copy') as pc:
tctx.command(c.clip, "@all", "q.method")
tctx.command(c.clip, "@all", "request.method")
assert pc.called
with mock.patch('pyperclip.copy') as pc:
tctx.command(c.clip, "@all", "q.content")
tctx.command(c.clip, "@all", "request.content")
assert pc.called
with mock.patch('pyperclip.copy') as pc:
tctx.command(c.clip, "@all", "q.method,q.content")
tctx.command(c.clip, "@all", "request.method,request.content")
assert pc.called
@ -94,17 +99,17 @@ def test_cut_save(tmpdir):
tctx.master.addons.add(v, c)
v.add([tflow.tflow(resp=True)])
tctx.command(c.save, "@all", "q.method", f)
tctx.command(c.save, "@all", "request.method", f)
assert qr(f) == b"GET"
tctx.command(c.save, "@all", "q.content", f)
tctx.command(c.save, "@all", "request.content", f)
assert qr(f) == b"content"
tctx.command(c.save, "@all", "q.content", "+" + f)
tctx.command(c.save, "@all", "request.content", "+" + f)
assert qr(f) == b"content\ncontent"
v.add([tflow.tflow(resp=True)])
tctx.command(c.save, "@all", "q.method", f)
tctx.command(c.save, "@all", "request.method", f)
assert qr(f).splitlines() == [b"GET", b"GET"]
tctx.command(c.save, "@all", "q.method,q.content", f)
tctx.command(c.save, "@all", "request.method,request.content", f)
assert qr(f).splitlines() == [b"GET,content", b"GET,content"]
@ -112,20 +117,20 @@ def test_cut():
c = cut.Cut()
with taddons.context():
tflows = [tflow.tflow(resp=True)]
assert c.cut(tflows, ["q.method"]) == [["GET"]]
assert c.cut(tflows, ["q.scheme"]) == [["http"]]
assert c.cut(tflows, ["q.host"]) == [["address"]]
assert c.cut(tflows, ["q.port"]) == [["22"]]
assert c.cut(tflows, ["q.path"]) == [["/path"]]
assert c.cut(tflows, ["q.url"]) == [["http://address:22/path"]]
assert c.cut(tflows, ["q.content"]) == [[b"content"]]
assert c.cut(tflows, ["q.header[header]"]) == [["qvalue"]]
assert c.cut(tflows, ["q.header[unknown]"]) == [[""]]
assert c.cut(tflows, ["request.method"]) == [["GET"]]
assert c.cut(tflows, ["request.scheme"]) == [["http"]]
assert c.cut(tflows, ["request.host"]) == [["address"]]
assert c.cut(tflows, ["request.port"]) == [["22"]]
assert c.cut(tflows, ["request.path"]) == [["/path"]]
assert c.cut(tflows, ["request.url"]) == [["http://address:22/path"]]
assert c.cut(tflows, ["request.content"]) == [[b"content"]]
assert c.cut(tflows, ["request.header[header]"]) == [["qvalue"]]
assert c.cut(tflows, ["request.header[unknown]"]) == [[""]]
assert c.cut(tflows, ["s.status_code"]) == [["200"]]
assert c.cut(tflows, ["s.reason"]) == [["OK"]]
assert c.cut(tflows, ["s.content"]) == [[b"message"]]
assert c.cut(tflows, ["s.header[header-response]"]) == [["svalue"]]
assert c.cut(tflows, ["response.status_code"]) == [["200"]]
assert c.cut(tflows, ["response.reason"]) == [["OK"]]
assert c.cut(tflows, ["response.content"]) == [[b"message"]]
assert c.cut(tflows, ["response.header[header-response]"]) == [["svalue"]]
assert c.cut(tflows, ["moo"]) == [[""]]
with pytest.raises(exceptions.CommandError):
assert c.cut(tflows, ["__dict__"]) == [[""]]
@ -133,5 +138,5 @@ def test_cut():
c = cut.Cut()
with taddons.context():
tflows = [tflow.ttcpflow()]
assert c.cut(tflows, ["q.method"]) == [[""]]
assert c.cut(tflows, ["s.status"]) == [[""]]
assert c.cut(tflows, ["request.method"]) == [[""]]
assert c.cut(tflows, ["response.status"]) == [[""]]

View File

@ -30,7 +30,7 @@ def test_order_refresh():
with taddons.context() as tctx:
tctx.configure(v, view_order="time")
v.add([tf])
tf.request.timestamp_start = 1
tf.request.timestamp_start = 10
assert not sargs
v.update([tf])
assert sargs
@ -41,7 +41,7 @@ def test_order_generators():
tf = tflow.tflow(resp=True)
rs = view.OrderRequestStart(v)
assert rs.generate(tf) == 0
assert rs.generate(tf) == 1
rm = view.OrderRequestMethod(v)
assert rm.generate(tf) == tf.request.method

View File

@ -150,10 +150,10 @@ class TestResponseUtils:
n = time.time()
r.headers["date"] = email.utils.formatdate(n)
pre = r.headers["date"]
r.refresh(n)
r.refresh(1)
assert pre == r.headers["date"]
r.refresh(n + 60)
r.refresh(61)
d = email.utils.parsedate_tz(r.headers["date"])
d = email.utils.mktime_tz(d)
# Weird that this is not exact...

View File

@ -45,8 +45,8 @@ export default function(){
"port": 22,
"pretty_host": "address",
"scheme": "http",
"timestamp_end": null,
"timestamp_start": null
"timestamp_end": 2,
"timestamp_start": 1
},
"response": {
"contentHash": "ab530a13e45914982b79f9b7e3fba994cfd1f3fb22f71cea1afbf02b460c6d1d",