Fix for #2750 and #2783 (#2809)

6 files changed, 77 insertions, 32 deletions

diff --git a/.gitignore b/.gitignore
index f88a2917..9fade1c3 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 d684b8c7..c31465c7 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 0169f5b1..173296e4 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 b10d27e4..5907fe95 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 c444b8ee..266f9de7 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 233c62d5..4ceb0429 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():