Refactor exporters

3 files changed, 18 insertions, 25 deletions

diff --git a/mitmproxy/console/common.py b/mitmproxy/console/common.py
index e1107f02..281fd658 100644
--- a/mitmproxy/console/common.py
+++ b/mitmproxy/console/common.py
@@ -312,22 +312,15 @@ def export_to_clip_or_file(key, scope, flow, writer):
     writer: copy_to_clipboard_or_prompt, ask_save_path
     """
 
-    for exporter in export.EXPORTERS:
-        if key in ["c", "h"]:
-            if scope is None:
-                ask_scope_and_callback(flow, handle_flow_data, key, writer)
-            else:
-                handle_flow_data(scope, flow, key, writer)
-        elif key == "u":
-            writer(flow.request.url)
-        elif key == "r":
-            writer(export.curl_command(flow))
-        elif key == "p":
-            writer(export.python_code(flow))
-        elif key == "l":
-            writer(export.locust_code(flow))
-        elif key == "t":
-            writer(export.locust_task(flow))
+    for _, exp_key, exporter in export.EXPORTERS:
+        if key == exp_key:
+            if exporter is None:  # 'c' & 'h'
+                if scope is None:
+                    ask_scope_and_callback(flow, handle_flow_data, key, writer)
+                else:
+                    handle_flow_data(scope, flow, key, writer)
+            else:  # other keys
+                writer(exporter(flow))
 
 flowcache = utils.LRUCache(800)
 
diff --git a/mitmproxy/console/flowlist.py b/mitmproxy/console/flowlist.py
index 9953e452..a33002a9 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 = export.EXPORTERS,
+                keys = [(e[0], e[1]) for e in export.EXPORTERS],
                 callback = common.export_to_clip_or_file,
                 args = (None, 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 = export.EXPORTERS,
+                keys = [(e[0], e[1]) for e in export.EXPORTERS],
                 callback = common.export_to_clip_or_file,
                 args = (None, self.flow, common.copy_to_clipboard_or_prompt)
             )
diff --git a/mitmproxy/flow/export.py b/mitmproxy/flow/export.py
index dbb0aad9..731aaf0e 100644
--- a/mitmproxy/flow/export.py
+++ b/mitmproxy/flow/export.py
@@ -199,11 +199,11 @@ def url(flow):
 
 
 EXPORTERS = [
-    ("content", "c"),
-    ("headers+content", "h"),
-    ("url", "u"),
-    ("as curl command", "r"),
-    ("as python code", "p"),
-    ("as locust code", "l"),
-    ("as locust task", "t"),
+    ("content", "c", None),
+    ("headers+content", "h", None),
+    ("url", "u", url),
+    ("as curl command", "r", curl_command),
+    ("as python code", "p", python_code),
+    ("as locust code", "l", locust_code),
+    ("as locust task", "t", locust_task),
 ]