From 4b93e16e785d593eb125b29151153c424b720918 Mon Sep 17 00:00:00 2001 From: Samoilenko Roman Date: Thu, 25 Jan 2018 17:07:09 +0200 Subject: [PATCH] Fix for #2750 and #2783 (#2809) --- .gitignore | 1 + mitmproxy/addons/cut.py | 43 +++++++++++++----------- mitmproxy/addons/export.py | 14 +++++--- mitmproxy/tools/console/consoleaddons.py | 17 ++++++---- test/mitmproxy/addons/test_cut.py | 19 +++++++++++ test/mitmproxy/addons/test_export.py | 15 +++++++++ 6 files changed, 77 insertions(+), 32 deletions(-) diff --git a/.gitignore b/.gitignore index f88a29174..9fade1c30 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,7 @@ MANIFEST .cache/ .tox*/ build/ +dist/ mitmproxy/contrib/kaitaistruct/*.ksy # UI diff --git a/mitmproxy/addons/cut.py b/mitmproxy/addons/cut.py index d684b8c78..c31465c78 100644 --- a/mitmproxy/addons/cut.py +++ b/mitmproxy/addons/cut.py @@ -88,26 +88,29 @@ class Cut: if path.startswith("+"): append = True path = mitmproxy.types.Path(path[1:]) - if len(cuts) == 1 and len(flows) == 1: - with open(path, "ab" if append else "wb") as fp: - if fp.tell() > 0: - # We're appending to a file that already exists and has content - fp.write(b"\n") - v = extract(cuts[0], flows[0]) - if isinstance(v, bytes): - fp.write(v) - else: - fp.write(v.encode("utf8")) - ctx.log.alert("Saved single cut.") - else: - with open(path, "a" if append else "w", newline='', encoding="utf8") as fp: - writer = csv.writer(fp) - for f in flows: - vals = [extract(c, f) for c in cuts] - writer.writerow( - [strutils.always_str(x) or "" for x in vals] # type: ignore - ) - ctx.log.alert("Saved %s cuts over %d flows as CSV." % (len(cuts), len(flows))) + try: + if len(cuts) == 1 and len(flows) == 1: + with open(path, "ab" if append else "wb") as fp: + if fp.tell() > 0: + # We're appending to a file that already exists and has content + fp.write(b"\n") + v = extract(cuts[0], flows[0]) + if isinstance(v, bytes): + fp.write(v) + else: + fp.write(v.encode("utf8")) + ctx.log.alert("Saved single cut.") + else: + with open(path, "a" if append else "w", newline='', encoding="utf8") as fp: + writer = csv.writer(fp) + for f in flows: + vals = [extract(c, f) for c in cuts] + writer.writerow( + [strutils.always_str(x) or "" for x in vals] # type: ignore + ) + ctx.log.alert("Saved %s cuts over %d flows as CSV." % (len(cuts), len(flows))) + except IOError as e: + ctx.log.error(str(e)) @command.command("cut.clip") def clip( diff --git a/mitmproxy/addons/export.py b/mitmproxy/addons/export.py index 0169f5b1b..173296e40 100644 --- a/mitmproxy/addons/export.py +++ b/mitmproxy/addons/export.py @@ -1,5 +1,6 @@ import typing +from mitmproxy import ctx from mitmproxy import command from mitmproxy import flow from mitmproxy import exceptions @@ -58,11 +59,14 @@ class Export(): raise exceptions.CommandError("No such export format: %s" % fmt) func = formats[fmt] # type: typing.Any v = func(f) - with open(path, "wb") as fp: - if isinstance(v, bytes): - fp.write(v) - else: - fp.write(v.encode("utf-8")) + try: + with open(path, "wb") as fp: + if isinstance(v, bytes): + fp.write(v) + else: + fp.write(v.encode("utf-8")) + except IOError as e: + ctx.log.error(str(e)) @command.command("export.clip") def clip(self, fmt: str, f: flow.Flow) -> None: diff --git a/mitmproxy/tools/console/consoleaddons.py b/mitmproxy/tools/console/consoleaddons.py index b10d27e45..5907fe95a 100644 --- a/mitmproxy/tools/console/consoleaddons.py +++ b/mitmproxy/tools/console/consoleaddons.py @@ -465,13 +465,16 @@ class ConsoleAddon: Save data to file as a CSV. """ rows = self._grideditor().value - with open(path, "w", newline='', encoding="utf8") as fp: - writer = csv.writer(fp) - for row in rows: - writer.writerow( - [strutils.always_str(x) or "" for x in row] # type: ignore - ) - ctx.log.alert("Saved %s rows as CSV." % (len(rows))) + try: + with open(path, "w", newline='', encoding="utf8") as fp: + writer = csv.writer(fp) + for row in rows: + writer.writerow( + [strutils.always_str(x) or "" for x in row] # type: ignore + ) + ctx.log.alert("Saved %s rows as CSV." % (len(rows))) + except IOError as e: + ctx.log.error(str(e)) @command.command("console.grideditor.editor") def grideditor_editor(self) -> None: diff --git a/test/mitmproxy/addons/test_cut.py b/test/mitmproxy/addons/test_cut.py index c444b8ee1..266f9de70 100644 --- a/test/mitmproxy/addons/test_cut.py +++ b/test/mitmproxy/addons/test_cut.py @@ -112,6 +112,25 @@ def test_cut_save(tmpdir): assert qr(f).splitlines() == [b"GET,content", b"GET,content"] +@pytest.mark.parametrize("exception, log_message", [ + (PermissionError, "Permission denied"), + (IsADirectoryError, "Is a directory"), + (FileNotFoundError, "No such file or directory") +]) +def test_cut_save_open(exception, log_message, 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)]) + + with mock.patch("mitmproxy.addons.cut.open") as m: + m.side_effect = exception(log_message) + tctx.command(c.save, "@all", "request.method", f) + assert tctx.master.has_log(log_message, level="error") + + def test_cut(): c = cut.Cut() with taddons.context(): diff --git a/test/mitmproxy/addons/test_export.py b/test/mitmproxy/addons/test_export.py index 233c62d5e..4ceb04295 100644 --- a/test/mitmproxy/addons/test_export.py +++ b/test/mitmproxy/addons/test_export.py @@ -94,6 +94,21 @@ def test_export(tmpdir): os.unlink(f) +@pytest.mark.parametrize("exception, log_message", [ + (PermissionError, "Permission denied"), + (IsADirectoryError, "Is a directory"), + (FileNotFoundError, "No such file or directory") +]) +def test_export_open(exception, log_message, tmpdir): + f = str(tmpdir.join("path")) + e = export.Export() + with taddons.context() as tctx: + with mock.patch("mitmproxy.addons.export.open") as m: + m.side_effect = exception(log_message) + e.file("raw", tflow.tflow(resp=True), f) + assert tctx.master.has_log(log_message, level="error") + + def test_clip(tmpdir): e = export.Export() with taddons.context():