Common function for both file and clipboard

2 files changed, 22 insertions, 18 deletions

diff --git a/mitmproxy/console/common.py b/mitmproxy/console/common.py
index c143f0c9..d84c4620 100644
--- a/mitmproxy/console/common.py
+++ b/mitmproxy/console/common.py
@@ -7,9 +7,14 @@ import urwid.util
 import six
 
 import netlib
+<<<<<<< a7bc7d29a484c87bc1576d33bf5ab23aa5031c43
 from mitmproxy import flow
+=======
+from mitmproxy import models
+>>>>>>> Common function for both file and clipboard
 from mitmproxy import utils
 from mitmproxy.console import signals
+from mitmproxy.flow import export
 from netlib import human
 
 try:
@@ -243,7 +248,7 @@ def ask_save_overwrite(path, data):
         save_data(path, data)
 
 
-def ask_save_path(prompt, data):
+def ask_save_path(data, prompt="File path"):
     signals.status_prompt_path.send(
         prompt = prompt,
         callback = ask_save_overwrite,
@@ -298,7 +303,7 @@ def copy_to_clipboard_or_prompt(data):
     except (RuntimeError, UnicodeDecodeError, AttributeError, TypeError):
         def save(k):
             if k == "y":
-                ask_save_path("Save data", data)
+                ask_save_path(data, "Save data")
         signals.status_prompt_onekey.send(
             prompt = "Cannot copy data to clipboard. Save as file?",
             keys = (
@@ -377,29 +382,28 @@ def ask_save_body(part, master, state, flow):
 
     elif part == "q" and request_has_content:
         ask_save_path(
-            "Save request content",
             flow.request.get_content(strict=False),
+            "Save request content",
         )
     elif part == "s" and response_has_content:
         ask_save_path(
-            "Save response content",
             flow.response.get_content(strict=False),
+            "Save response content",
         )
     else:
         signals.status_message.send(message="No content to save.")
 
 
-def export_to_clipboard(k, f):
-    for exporter in flow.export.EXPORTERS:
-        if k == exporter[1]:
-            copy_to_clipboard_or_prompt(exporter[2](f))
-
-
-def export_to_file(k, f):
-    for exporter in flow.export.EXPORTERS:
-        if k == exporter[1]:
-            ask_save_path("File path", exporter[2](f))
+def export_to_clip_or_file(key, flow, writer):
+    """
+    Export selected flow to clipboard or a file.
 
+    'writer' is a function that handles the data
+    can be: copy_to_clipboard_or_prompt or ask_save_path
+    """
+    for exporter in export.EXPORTERS:
+        if key == exporter[1]:
+            writer(exporter[2](flow))
 
 flowcache = utils.LRUCache(800)
 
diff --git a/mitmproxy/console/flowlist.py b/mitmproxy/console/flowlist.py
index 5b1701d8..68d63938 100644
--- a/mitmproxy/console/flowlist.py
+++ b/mitmproxy/console/flowlist.py
@@ -269,16 +269,16 @@ class ConnectionItem(urwid.WidgetWrap):
                 self,
                 prompt = "Export to file",
                 keys = [(e[0], e[1]) for e in export.EXPORTERS],
-                callback = common.export_to_file,
-                args = (self.flow,)
+                callback = common.export_to_clip_or_file,
+                args = (self.flow, common.ask_save_path)
             )
         elif key == "P":
             signals.status_prompt_onekey.send(
                 self,
                 prompt = "Export to clipboard",
                 keys = [(e[0], e[1]) for e in export.EXPORTERS],
-                callback = common.export_to_clipboard,
-                args = (self.flow,)
+                callback = common.export_to_clip_or_file,
+                args = (self.flow, common.copy_to_clipboard_or_prompt)
             )
         elif key == "b":
             common.ask_save_body(None, self.master, self.state, self.flow)