From 0fa7c463683149a0b400108420efc09ff72f11a5 Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Wed, 30 Mar 2022 14:24:18 +0200
Subject: [PATCH] fix cut addon to work with binary content, fix #3965 (#5230)

---
 CHANGELOG.md                      |  2 ++
 mitmproxy/addons/cut.py           | 25 ++++++++++++++-----------
 test/mitmproxy/addons/test_cut.py |  8 +++++++-
 3 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b858091a3..8c3ecf459 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,8 @@
   ([#5225](https://github.com/mitmproxy/mitmproxy/issues/5225), @mhils)
 * Add example addon for domain fronting.
   ([#5217](https://github.com/mitmproxy/mitmproxy/issues/5217), @randomstuff)
+* Improve cut addon to better handle binary contents
+  ([#3965](https://github.com/mitmproxy/mitmproxy/issues/3965), @mhils)
 
 ## 19 March 2022: mitmproxy 8.0.0
 
diff --git a/mitmproxy/addons/cut.py b/mitmproxy/addons/cut.py
index c75c14ca7..49fa562ff 100644
--- a/mitmproxy/addons/cut.py
+++ b/mitmproxy/addons/cut.py
@@ -8,7 +8,6 @@ from mitmproxy import exceptions
 from mitmproxy import flow
 from mitmproxy import ctx
 from mitmproxy import certs
-from mitmproxy.utils import strutils
 import mitmproxy.types
 
 import pyperclip
@@ -54,6 +53,14 @@ def extract(cut: str, f: flow.Flow) -> typing.Union[str, bytes]:
     return str(current or "")
 
 
+def extract_str(cut: str, f: flow.Flow) -> str:
+    ret = extract(cut, f)
+    if isinstance(ret, bytes):
+        return repr(ret)
+    else:
+        return ret
+
+
 class Cut:
     @command.command("cut")
     def cut(
@@ -110,10 +117,8 @@ class Cut:
                 with open(path, "a" if append else "w", newline='', encoding="utf8") as tfp:
                     writer = csv.writer(tfp)
                     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
-                        )
+                        vals = [extract_str(c, f) for c in cuts]
+                        writer.writerow(vals)
                 ctx.log.alert("Saved %s cuts over %d flows as CSV." % (len(cuts), len(flows)))
         except OSError as e:
             ctx.log.error(str(e))
@@ -132,16 +137,14 @@ class Cut:
         v: typing.Union[str, bytes]
         fp = io.StringIO(newline="")
         if len(cuts) == 1 and len(flows) == 1:
-            v = extract(cuts[0], flows[0])
-            fp.write(strutils.always_str(v))  # type: ignore
+            v = extract_str(cuts[0], flows[0])
+            fp.write(v)
             ctx.log.alert("Clipped single cut.")
         else:
             writer = csv.writer(fp)
             for f in flows:
-                vals = [extract(c, f) for c in cuts]
-                writer.writerow(
-                    [strutils.always_str(v) for v in vals]
-                )
+                vals = [extract_str(c, f) for c in cuts]
+                writer.writerow(vals)
             ctx.log.alert("Clipped %s cuts as CSV." % len(cuts))
         try:
             pyperclip.copy(fp.getvalue())
diff --git a/test/mitmproxy/addons/test_cut.py b/test/mitmproxy/addons/test_cut.py
index e40b819ba..6ddac7de5 100644
--- a/test/mitmproxy/addons/test_cut.py
+++ b/test/mitmproxy/addons/test_cut.py
@@ -59,6 +59,12 @@ def test_extract(tdata):
     assert "CERTIFICATE" in cut.extract("server_conn.certificate_list", tf)
 
 
+def test_extract_str():
+    tf = tflow.tflow()
+    tf.request.raw_content = b"\xFF"
+    assert cut.extract_str("request.raw_content", tf) == r"b'\xff'"
+
+
 def test_headername():
     with pytest.raises(exceptions.CommandError):
         cut.headername("header[foo.")
@@ -115,7 +121,7 @@ def test_cut_save(tmpdir):
         tctx.command(c.save, "@all", "request.method", f)
         assert qr(f).splitlines() == [b"GET", b"GET"]
         tctx.command(c.save, "@all", "request.method,request.content", f)
-        assert qr(f).splitlines() == [b"GET,content", b"GET,content"]
+        assert qr(f).splitlines() == [b"GET,b'content'", b"GET,b'content'"]
 
 
 @pytest.mark.parametrize("exception, log_message", [