From e3670f9d9b1e92e8f45d03048cf314998592aa8d Mon Sep 17 00:00:00 2001
From: Shadab Zafar <dufferzafar0@gmail.com>
Date: Sat, 16 Jul 2016 17:48:02 +0530
Subject: [PATCH] Integrate ask_copy_part into exporters

This introduces some duplication but that'll be refactored later.
---
 mitmproxy/console/common.py   | 35 ++++++++++++++---------------------
 mitmproxy/console/flowlist.py |  4 ++--
 mitmproxy/flow/export.py      | 11 +++++++----
 3 files changed, 23 insertions(+), 27 deletions(-)

diff --git a/mitmproxy/console/common.py b/mitmproxy/console/common.py
index 838c719f9..4eaa7bc1e 100644
--- a/mitmproxy/console/common.py
+++ b/mitmproxy/console/common.py
@@ -194,7 +194,7 @@ def copy_to_clipboard_or_prompt(data):
         )
 
 
-def flow_format_data(key, scope, flow):
+def format_flow_data(key, scope, flow):
     if key == "u":
         data = flow.request.url
     else:
@@ -227,12 +227,12 @@ def flow_format_data(key, scope, flow):
     return data, False
 
 
-def copy_flow(key, scope, flow, writer):
+def write_flow_data(key, scope, flow, writer):
     """
     key: _c_ontent, _h_eaders+content, _u_rl
     scope: _a_ll, re_q_uest, re_s_ponse
     """
-    data, err = flow_format_data(key, scope, flow)
+    data, err = format_flow_data(key, scope, flow)
 
     if err:
         signals.status_message.send(message=err)
@@ -250,22 +250,6 @@ def copy_flow(key, scope, flow, writer):
     writer(data)
 
 
-def ask_copy_part(scope, flow):
-    choices = [
-        ("content", "c"),
-        ("headers+content", "h")
-    ]
-    if scope != "s":
-        choices.append(("url", "u"))
-
-    signals.status_prompt_onekey.send(
-        prompt = "Copy",
-        keys = choices,
-        callback = copy_flow,
-        args = (scope, flow)
-    )
-
-
 def ask_save_body(scope, flow):
     """
     Save either the request or the response body to disk.
@@ -315,9 +299,18 @@ def export_to_clip_or_file(key, scope, flow, writer):
     '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))
+        if key in ["c", "h", "u"]:
+            write_flow_data(key, scope, flow, writer)
+        if key == "r":
+            writer(export.curl_command(flow))
+        if key == "p":
+            writer(export.python_code(flow))
+        if key == "l":
+            writer(export.locust_code(flow))
+        if key == "t":
+            writer(export.locust_task(flow))
 
 flowcache = utils.LRUCache(800)
 
diff --git a/mitmproxy/console/flowlist.py b/mitmproxy/console/flowlist.py
index 697550ccf..1ca9d0ae0 100644
--- a/mitmproxy/console/flowlist.py
+++ b/mitmproxy/console/flowlist.py
@@ -268,7 +268,7 @@ class ConnectionItem(urwid.WidgetWrap):
             signals.status_prompt_onekey.send(
                 self,
                 prompt = "Export to file",
-                keys = [(e[0], e[1]) for e in export.EXPORTERS],
+                keys = export.EXPORTERS,
                 callback = common.export_to_clip_or_file,
                 args = ("a", self.flow, common.ask_save_path)
             )
@@ -276,7 +276,7 @@ class ConnectionItem(urwid.WidgetWrap):
             signals.status_prompt_onekey.send(
                 self,
                 prompt = "Export to clipboard",
-                keys = [(e[0], e[1]) for e in export.EXPORTERS],
+                keys = export.EXPORTERS,
                 callback = common.export_to_clip_or_file,
                 args = ("a", self.flow, common.copy_to_clipboard_or_prompt)
             )
diff --git a/mitmproxy/flow/export.py b/mitmproxy/flow/export.py
index 5ef358ca8..476b18980 100644
--- a/mitmproxy/flow/export.py
+++ b/mitmproxy/flow/export.py
@@ -195,8 +195,11 @@ def locust_task(flow):
 
 
 EXPORTERS = [
-    ("as curl command", "c", curl_command),
-    ("as python code", "p", python_code),
-    ("as locust code", "l", locust_code),
-    ("as locust task", "t", locust_task),
+    ("content", "c"),
+    ("headers+content", "h"),
+    ("url", "u"),
+    ("as curl command", "r"),
+    ("as python code", "p"),
+    ("as locust code", "l"),
+    ("as locust task", "t"),
 ]