diff --git a/mitmproxy/addons/static_viewer.py b/mitmproxy/addons/static_viewer.py index cf1ae9aeb..a46292304 100644 --- a/mitmproxy/addons/static_viewer.py +++ b/mitmproxy/addons/static_viewer.py @@ -42,24 +42,27 @@ def save_flows(path: pathlib.Path, flows: typing.Iterable[flow.Flow]) -> None: def save_flows_content(path: pathlib.Path, flows: typing.Iterable[flow.Flow]) -> None: - for flow in flows: + for f in flows: for m in ('request', 'response'): - message = getattr(flow, m) - message_path = path / "flows" / flow.id / m + message = getattr(f, m) + message_path = path / "flows" / f.id / m os.makedirs(str(message_path / "content"), exist_ok=True) - with open(str(message_path / '_content'), 'wb') as f: + with open(str(message_path / '_content'), 'wb') as content_file: # don't use raw_content here as this is served with a default content type - if not message: - # skip missing message - continue - f.write(message.content) + if message: + content_file.write(message.content) + else: + content_file.write(b'No content.') # content_view t = time.time() - description, lines, error = contentviews.get_message_content_view( - 'Auto', message - ) + if message: + description, lines, error = contentviews.get_message_content_view( + 'Auto', message + ) + else: + description, lines = 'No content.', [] if time.time() - t > 0.1: ctx.log( "Slow content view: {} took {}s".format( @@ -68,10 +71,10 @@ def save_flows_content(path: pathlib.Path, flows: typing.Iterable[flow.Flow]) -> ), "info" ) - with open(str(message_path / "content" / "Auto.json"), "w") as f: + with open(str(message_path / "content" / "Auto.json"), "w") as content_view_file: json.dump( dict(lines=list(lines), description=description), - f + content_view_file ) diff --git a/test/mitmproxy/addons/test_static_viewer.py b/test/mitmproxy/addons/test_static_viewer.py new file mode 100644 index 000000000..7e34a9fc0 --- /dev/null +++ b/test/mitmproxy/addons/test_static_viewer.py @@ -0,0 +1,60 @@ +import json + +from mitmproxy.test import taddons +from mitmproxy.test import tflow + +from mitmproxy import flowfilter +from mitmproxy.tools.web.app import flow_to_json + +from mitmproxy.addons import static_viewer +from mitmproxy.addons import save + + +def test_save_static(tmpdir): + tmpdir.mkdir('static') + static_viewer.save_static(tmpdir) + assert len(tmpdir.listdir()) == 2 + assert tmpdir.join('index.html').check(file=1) + assert tmpdir.join('static/static.js').read() == 'MITMWEB_STATIC = true;' + + +def test_save_filter_help(tmpdir): + static_viewer.save_filter_help(tmpdir) + f = tmpdir.join('/filter-help.json') + assert f.check(file=1) + assert f.read() == json.dumps(dict(commands=flowfilter.help)) + + +def test_save_flows(tmpdir): + flows = [tflow.tflow(req=True, resp=None), tflow.tflow(req=True, resp=True)] + static_viewer.save_flows(tmpdir, flows) + assert tmpdir.join('flows.json').check(file=1) + assert tmpdir.join('flows.json').read() == json.dumps([flow_to_json(f) for f in flows]) + + +def test_save_flows_content(tmpdir): + flows = [tflow.tflow(req=True, resp=None), tflow.tflow(req=True, resp=True)] + static_viewer.save_flows_content(tmpdir, flows) + flows_path = tmpdir.join('flows') + assert len(flows_path.listdir()) == len(flows) + for p in flows_path.listdir(): + assert p.join('request').check(dir=1) + assert p.join('response').check(dir=1) + assert p.join('request/_content').check(file=1) + assert p.join('request/content').check(dir=1) + assert p.join('response/_content').check(file=1) + assert p.join('response/content').check(dir=1) + assert p.join('request/content/Auto.json').check(file=1) + assert p.join('response/content/Auto.json').check(file=1) + + +def test_static_viewer(tmpdir): + s = static_viewer.StaticViewer() + sa = save.Save() + with taddons.context() as tctx: + sa.save([tflow.tflow(resp=True)], str(tmpdir.join('foo'))) + tctx.master.addons.add(s) + tctx.configure(s, web_static_viewer=str(tmpdir), rfile=str(tmpdir.join('foo'))) + assert tmpdir.join('index.html').check(file=1) + assert tmpdir.join('static').check(dir=1) + assert tmpdir.join('flows').check(dir=1)