Merge pull request #884 from dufferzafar/export

Export/Copy flow as...

3 files changed, 113 insertions, 34 deletions

diff --git a/libmproxy/console/common.py b/libmproxy/console/common.py
index 12fdfe27..a9bd846b 100644
--- a/libmproxy/console/common.py
+++ b/libmproxy/console/common.py
@@ -207,7 +207,7 @@ def raw_format_flow(f, focus, extended, padding):
 
 
 # Save file to disk
-def save_data(path, data, master, state):
+def save_data(path, data):
     if not path:
         return
     try:
@@ -217,32 +217,32 @@ def save_data(path, data, master, state):
         signals.status_message.send(message=v.strerror)
 
 
-def ask_save_overwite(path, data, master, state):
+def ask_save_overwrite(path, data):
     if not path:
         return
     path = os.path.expanduser(path)
     if os.path.exists(path):
-        def save_overwite(k):
+        def save_overwrite(k):
             if k == "y":
-                save_data(path, data, master, state)
+                save_data(path, data)
 
         signals.status_prompt_onekey.send(
-            prompt = "'" + path + "' already exists. Overwite?",
+            prompt = "'" + path + "' already exists. Overwrite?",
             keys = (
                 ("yes", "y"),
                 ("no", "n"),
             ),
-            callback = save_overwite
+            callback = save_overwrite
         )
     else:
-        save_data(path, data, master, state)
+        save_data(path, data)
 
 
-def ask_save_path(prompt, data, master, state):
+def ask_save_path(prompt, data):
     signals.status_prompt_path.send(
         prompt = prompt,
-        callback = ask_save_overwite,
-        args = (data, master, state)
+        callback = ask_save_overwrite,
+        args =  (data, )
     )
 
 
@@ -277,26 +277,60 @@ def copy_flow_format_data(part, scope, flow):
     return data, False
 
 
-def copy_flow(part, scope, flow, master, state):
-    """
-    part: _c_ontent, _h_eaders+content, _u_rl
-    scope: _a_ll, re_q_uest, re_s_ponse
-    """
-    data, err = copy_flow_format_data(part, scope, flow)
+def export_prompt(k, flow):
+    if k == "c":
+        copy_as_curl_command(flow)
+    elif k == "p":
+        copy_as_python_code(flow)
+    elif k == "r":
+        copy_as_raw_request(flow)
 
-    if err:
-        signals.status_message.send(message=err)
-        return
 
-    if not data:
-        if scope == "q":
-            signals.status_message.send(message="No request content to copy.")
-        elif scope == "s":
-            signals.status_message.send(message="No response content to copy.")
-        else:
-            signals.status_message.send(message="No contents to copy.")
+def copy_as_curl_command(flow):
+    data = "curl "
+
+    for k, v in flow.request.headers.fields:
+      data += "-H '%s:%s' " % (k, v)
+
+    if flow.request.method != "GET":
+      data += "-X %s " % flow.request.method
+
+    full_url = flow.request.scheme + "://" + flow.request.host + flow.request.path
+    data += "'%s'" % full_url
+
+    if flow.request.content:
+      data += " --data-binary '%s'" % flow.request.content
+
+    copy_to_clipboard_or_prompt(data)
+
+
+def copy_as_python_code(flow):
+    if flow.request.method != "GET":
+        signals.status_message.send(message="Currently, only GET methods are supported")
         return
 
+    data = ("import requests\n"
+            "headers = {%s}\n"
+            "url = '%s'\n"
+            "resp = requests.get(url, headers=headers)")
+
+    headers = "\n"
+    for k, v in flow.request.headers.fields:
+      headers += "    '%s': '%s',\n" % (k, v)
+
+    full_url = flow.request.scheme + "://" + flow.request.host + flow.request.path
+
+    data = data % (headers, full_url)
+
+    copy_to_clipboard_or_prompt(data)
+
+
+def copy_as_raw_request(flow):
+    data = netlib.http.http1.assemble_request(flow.request)
+    copy_to_clipboard_or_prompt(data)
+
+
+def copy_to_clipboard_or_prompt(data):
     # pyperclip calls encode('utf-8') on data to be copied without checking.
     # if data are already encoded that way UnicodeDecodeError is thrown.
     toclip = ""
@@ -310,7 +344,7 @@ def copy_flow(part, scope, flow, master, state):
     except (RuntimeError, UnicodeDecodeError, AttributeError):
         def save(k):
             if k == "y":
-                ask_save_path("Save data", data, master, state)
+                ask_save_path("Save data", data)
         signals.status_prompt_onekey.send(
             prompt = "Cannot copy data to clipboard. Save as file?",
             keys = (
@@ -321,6 +355,29 @@ def copy_flow(part, scope, flow, master, state):
         )
 
 
+def copy_flow(part, scope, flow, master, state):
+    """
+    part: _c_ontent, _h_eaders+content, _u_rl
+    scope: _a_ll, re_q_uest, re_s_ponse
+    """
+    data, err = copy_flow_format_data(part, scope, flow)
+
+    if err:
+        signals.status_message.send(message=err)
+        return
+
+    if not data:
+        if scope == "q":
+            signals.status_message.send(message="No request content to copy.")
+        elif scope == "s":
+            signals.status_message.send(message="No response content to copy.")
+        else:
+            signals.status_message.send(message="No contents to copy.")
+        return
+
+    copy_to_clipboard_or_prompt(data)
+
+
 def ask_copy_part(scope, flow, master, state):
     choices = [
         ("content", "c"),
@@ -367,16 +424,12 @@ def ask_save_body(part, master, state, flow):
     elif part == "q" and request_has_content:
         ask_save_path(
             "Save request content",
-            flow.request.get_decoded_content(),
-            master,
-            state
+            flow.request.get_decoded_content()
         )
     elif part == "s" and response_has_content:
         ask_save_path(
             "Save response content",
-            flow.response.get_decoded_content(),
-            master,
-            state
+            flow.response.get_decoded_content()
         )
     else:
         signals.status_message.send(message="No content to save.")
diff --git a/libmproxy/console/flowlist.py b/libmproxy/console/flowlist.py
index 2b77f4a3..2c18d197 100644
--- a/libmproxy/console/flowlist.py
+++ b/libmproxy/console/flowlist.py
@@ -16,6 +16,7 @@ def _mkhelp():
         ("C", "clear flow list or eventlog"),
         ("d", "delete flow"),
         ("D", "duplicate flow"),
+        ("E", "export"),
         ("e", "toggle eventlog"),
         ("F", "toggle follow flow list"),
         ("l", "set limit filter pattern"),
@@ -254,6 +255,18 @@ class ConnectionItem(urwid.WidgetWrap):
             )
         elif key == "P":
             common.ask_copy_part("a", self.flow, self.master, self.state)
+        elif key == "E":
+            signals.status_prompt_onekey.send(
+                self,
+                prompt = "Export",
+                keys = (
+                    ("as curl command", "c"),
+                    ("as python code", "p"),
+                    ("as raw request", "r"),
+                ),
+                callback = common.export_prompt,
+                args = (self.flow,)
+            )
         elif key == "b":
             common.ask_save_body(None, self.master, self.state, self.flow)
         else:
diff --git a/libmproxy/console/flowview.py b/libmproxy/console/flowview.py
index 0038558b..7962c69c 100644
--- a/libmproxy/console/flowview.py
+++ b/libmproxy/console/flowview.py
@@ -25,8 +25,9 @@ def _mkhelp():
         ("A", "accept all intercepted flows"),
         ("a", "accept this intercepted flow"),
         ("b", "save request/response body"),
-        ("d", "delete flow"),
         ("D", "duplicate flow"),
+        ("d", "delete flow"),
+        ("E", "export"),
         ("e", "edit request/response"),
         ("f", "load full body data"),
         ("m", "change body display mode for this entity"),
@@ -574,6 +575,18 @@ class FlowView(tabs.Tabs):
                 callback = self.master.save_one_flow,
                 args = (self.flow,)
             )
+        elif key == "E":
+            signals.status_prompt_onekey.send(
+                self,
+                prompt = "Export",
+                keys = (
+                    ("as curl command", "c"),
+                    ("as python code", "p"),
+                    ("as raw request", "r"),
+                ),
+                callback = common.export_prompt,
+                args = (self.flow,)
+            )
         elif key == "|":
             signals.status_prompt_path.send(
                 prompt = "Send flow to script",