Merge branch 'flow-export' of https://github.com/dufferzafar/mitmproxy into dufferzafar-flow-export

5 files changed, 248 insertions, 260 deletions

diff --git a/mitmproxy/console/common.py b/mitmproxy/console/common.py
index 5d15e0cd..281fd658 100644
--- a/mitmproxy/console/common.py
+++ b/mitmproxy/console/common.py
@@ -7,9 +7,9 @@ import urwid.util
 import six
 
 import netlib
-from mitmproxy import flow
 from mitmproxy import utils
 from mitmproxy.console import signals
+from mitmproxy.flow import export
 from netlib import human
 
 try:
@@ -129,88 +129,6 @@ else:
     SYMBOL_MARK = "[m]"
 
 
-def raw_format_flow(f, focus, extended):
-    f = dict(f)
-    pile = []
-    req = []
-    if extended:
-        req.append(
-            fcol(
-                human.format_timestamp(f["req_timestamp"]),
-                "highlight"
-            )
-        )
-    else:
-        req.append(fcol(">>" if focus else "  ", "focus"))
-
-    if f["marked"]:
-        req.append(fcol(SYMBOL_MARK, "mark"))
-
-    if f["req_is_replay"]:
-        req.append(fcol(SYMBOL_REPLAY, "replay"))
-    req.append(fcol(f["req_method"], "method"))
-
-    preamble = sum(i[1] for i in req) + len(req) - 1
-
-    if f["intercepted"] and not f["acked"]:
-        uc = "intercept"
-    elif "resp_code" in f or "err_msg" in f:
-        uc = "text"
-    else:
-        uc = "title"
-
-    url = f["req_url"]
-    if f["req_http_version"] not in ("HTTP/1.0", "HTTP/1.1"):
-        url += " " + f["req_http_version"]
-    req.append(
-        urwid.Text([(uc, url)])
-    )
-
-    pile.append(urwid.Columns(req, dividechars=1))
-
-    resp = []
-    resp.append(
-        ("fixed", preamble, urwid.Text(""))
-    )
-
-    if "resp_code" in f:
-        codes = {
-            2: "code_200",
-            3: "code_300",
-            4: "code_400",
-            5: "code_500",
-        }
-        ccol = codes.get(f["resp_code"] / 100, "code_other")
-        resp.append(fcol(SYMBOL_RETURN, ccol))
-        if f["resp_is_replay"]:
-            resp.append(fcol(SYMBOL_REPLAY, "replay"))
-        resp.append(fcol(f["resp_code"], ccol))
-        if extended:
-            resp.append(fcol(f["resp_reason"], ccol))
-        if f["intercepted"] and f["resp_code"] and not f["acked"]:
-            rc = "intercept"
-        else:
-            rc = "text"
-
-        if f["resp_ctype"]:
-            resp.append(fcol(f["resp_ctype"], rc))
-        resp.append(fcol(f["resp_clen"], rc))
-        resp.append(fcol(f["roundtrip"], rc))
-
-    elif f["err_msg"]:
-        resp.append(fcol(SYMBOL_RETURN, "error"))
-        resp.append(
-            urwid.Text([
-                (
-                    "error",
-                    f["err_msg"]
-                )
-            ])
-        )
-    pile.append(urwid.Columns(resp, dividechars=1))
-    return urwid.Pile(pile)
-
-
 # Save file to disk
 def save_data(path, data):
     if not path:
@@ -243,7 +161,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,
@@ -251,49 +169,25 @@ def ask_save_path(prompt, data):
     )
 
 
-def copy_flow_format_data(part, scope, flow):
-    if part == "u":
-        data = flow.request.url
-    else:
-        data = ""
-        if scope in ("q", "a"):
-            request = flow.request.copy()
-            request.decode(strict=False)
-            if request.content is None:
-                return None, "Request content is missing"
-            if part == "h":
-                data += netlib.http.http1.assemble_request(request)
-            elif part == "c":
-                data += request.content
-            else:
-                raise ValueError("Unknown part: {}".format(part))
-        if scope == "a" and flow.request.raw_content and flow.response:
-            # Add padding between request and response
-            data += "\r\n" * 2
-        if scope in ("s", "a") and flow.response:
-            response = flow.response.copy()
-            response.decode(strict=False)
-            if response.content is None:
-                return None, "Response content is missing"
-            if part == "h":
-                data += netlib.http.http1.assemble_response(response)
-            elif part == "c":
-                data += response.content
-            else:
-                raise ValueError("Unknown part: {}".format(part))
-    return data, False
-
+def ask_scope_and_callback(flow, cb, *args):
+    request_has_content = flow.request and flow.request.raw_content
+    response_has_content = flow.response and flow.response.raw_content
 
-def export_prompt(k, f):
-    exporters = {
-        "c": flow.export.curl_command,
-        "p": flow.export.python_code,
-        "r": flow.export.raw_request,
-        "l": flow.export.locust_code,
-        "t": flow.export.locust_task,
-    }
-    if k in exporters:
-        copy_to_clipboard_or_prompt(exporters[k](f))
+    if request_has_content and response_has_content:
+        signals.status_prompt_onekey.send(
+            prompt = "Save",
+            keys = (
+                ("request", "q"),
+                ("response", "s"),
+                ("both", "b"),
+            ),
+            callback = cb,
+            args = (flow,) + args
+        )
+    elif response_has_content:
+        cb("s", flow, *args)
+    else:
+        cb("q", flow, *args)
 
 
 def copy_to_clipboard_or_prompt(data):
@@ -310,7 +204,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 = (
@@ -321,12 +215,43 @@ def copy_to_clipboard_or_prompt(data):
         )
 
 
-def copy_flow(part, scope, flow, master, state):
+def format_flow_data(key, scope, flow):
+    data = ""
+    if scope in ("q", "b"):
+        request = flow.request.copy()
+        request.decode(strict=False)
+        if request.content is None:
+            return None, "Request content is missing"
+        if key == "h":
+            data += netlib.http.http1.assemble_request(request)
+        elif key == "c":
+            data += request.get_content(strict=False)
+        else:
+            raise ValueError("Unknown key: {}".format(key))
+    if scope == "b" and flow.request.raw_content and flow.response:
+        # Add padding between request and response
+        data += "\r\n" * 2
+    if scope in ("s", "b") and flow.response:
+        response = flow.response.copy()
+        response.decode(strict=False)
+        if response.content is None:
+            return None, "Response content is missing"
+        if key == "h":
+            data += netlib.http.http1.assemble_response(response)
+        elif key == "c":
+            data += response.get_content(strict=False)
+        else:
+            raise ValueError("Unknown key: {}".format(key))
+    return data, False
+
+
+def handle_flow_data(scope, flow, key, writer):
     """
-    part: _c_ontent, _h_eaders+content, _u_rl
-    scope: _a_ll, re_q_uest, re_s_ponse
+    key: _c_ontent, _h_eaders+content, _u_rl
+    scope: re_q_uest, re_s_ponse, _b_oth
+    writer: copy_to_clipboard_or_prompt, ask_save_path
     """
-    data, err = copy_flow_format_data(part, scope, flow)
+    data, err = format_flow_data(key, scope, flow)
 
     if err:
         signals.status_message.send(message=err)
@@ -334,76 +259,154 @@ def copy_flow(part, scope, flow, master, state):
 
     if not data:
         if scope == "q":
-            signals.status_message.send(message="No request content to copy.")
+            signals.status_message.send(message="No request content.")
         elif scope == "s":
-            signals.status_message.send(message="No response content to copy.")
+            signals.status_message.send(message="No response content.")
         else:
-            signals.status_message.send(message="No contents to copy.")
+            signals.status_message.send(message="No content.")
         return
 
-    copy_to_clipboard_or_prompt(data)
+    writer(data)
 
 
-def ask_copy_part(scope, flow, master, state):
-    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, master, state)
-    )
-
-
-def ask_save_body(part, master, state, flow):
+def ask_save_body(scope, flow):
     """
-    Save either the request or the response body to disk. part can either be
-    "q" (request), "s" (response) or None (ask user if necessary).
+    Save either the request or the response body to disk.
+
+    scope: re_q_uest, re_s_ponse, _b_oth, None (ask user if necessary)
     """
 
     request_has_content = flow.request and flow.request.raw_content
     response_has_content = flow.response and flow.response.raw_content
 
-    if part is None:
-        # We first need to determine whether we want to save the request or the
-        # response content.
-        if request_has_content and response_has_content:
-            signals.status_prompt_onekey.send(
-                prompt = "Save",
-                keys = (
-                    ("request", "q"),
-                    ("response", "s"),
-                ),
-                callback = ask_save_body,
-                args = (master, state, flow)
-            )
-        elif response_has_content:
-            ask_save_body("s", master, state, flow)
-        else:
-            ask_save_body("q", master, state, flow)
-
-    elif part == "q" and request_has_content:
+    if scope is None:
+        ask_scope_and_callback(flow, ask_save_body)
+    elif scope == "q" and request_has_content:
         ask_save_path(
-            "Save request content",
             flow.request.get_content(strict=False),
+            "Save request content to"
         )
-    elif part == "s" and response_has_content:
+    elif scope == "s" and response_has_content:
         ask_save_path(
-            "Save response content",
             flow.response.get_content(strict=False),
+            "Save response content to"
+        )
+    elif scope == "b" and request_has_content and response_has_content:
+        ask_save_path(
+            (flow.request.get_content(strict=False) + "\n" +
+             flow.response.get_content(strict=False)),
+            "Save request & response content to"
         )
     else:
-        signals.status_message.send(message="No content to save.")
+        signals.status_message.send(message="No content.")
 
 
+def export_to_clip_or_file(key, scope, flow, writer):
+    """
+    Export selected flow to clipboard or a file.
+
+    key:    _c_ontent, _h_eaders+content, _u_rl,
+            cu_r_l_command, _p_ython_code,
+            _l_ocust_code, locust_t_ask
+    scope:  None, _a_ll, re_q_uest, re_s_ponse
+    writer: copy_to_clipboard_or_prompt, ask_save_path
+    """
+
+    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)
 
 
+def raw_format_flow(f, focus, extended):
+    f = dict(f)
+    pile = []
+    req = []
+    if extended:
+        req.append(
+            fcol(
+                human.format_timestamp(f["req_timestamp"]),
+                "highlight"
+            )
+        )
+    else:
+        req.append(fcol(">>" if focus else "  ", "focus"))
+
+    if f["marked"]:
+        req.append(fcol(SYMBOL_MARK, "mark"))
+
+    if f["req_is_replay"]:
+        req.append(fcol(SYMBOL_REPLAY, "replay"))
+    req.append(fcol(f["req_method"], "method"))
+
+    preamble = sum(i[1] for i in req) + len(req) - 1
+
+    if f["intercepted"] and not f["acked"]:
+        uc = "intercept"
+    elif "resp_code" in f or "err_msg" in f:
+        uc = "text"
+    else:
+        uc = "title"
+
+    url = f["req_url"]
+    if f["req_http_version"] not in ("HTTP/1.0", "HTTP/1.1"):
+        url += " " + f["req_http_version"]
+    req.append(
+        urwid.Text([(uc, url)])
+    )
+
+    pile.append(urwid.Columns(req, dividechars=1))
+
+    resp = []
+    resp.append(
+        ("fixed", preamble, urwid.Text(""))
+    )
+
+    if "resp_code" in f:
+        codes = {
+            2: "code_200",
+            3: "code_300",
+            4: "code_400",
+            5: "code_500",
+        }
+        ccol = codes.get(f["resp_code"] / 100, "code_other")
+        resp.append(fcol(SYMBOL_RETURN, ccol))
+        if f["resp_is_replay"]:
+            resp.append(fcol(SYMBOL_REPLAY, "replay"))
+        resp.append(fcol(f["resp_code"], ccol))
+        if extended:
+            resp.append(fcol(f["resp_reason"], ccol))
+        if f["intercepted"] and f["resp_code"] and not f["acked"]:
+            rc = "intercept"
+        else:
+            rc = "text"
+
+        if f["resp_ctype"]:
+            resp.append(fcol(f["resp_ctype"], rc))
+        resp.append(fcol(f["resp_clen"], rc))
+        resp.append(fcol(f["roundtrip"], rc))
+
+    elif f["err_msg"]:
+        resp.append(fcol(SYMBOL_RETURN, "error"))
+        resp.append(
+            urwid.Text([
+                (
+                    "error",
+                    f["err_msg"]
+                )
+            ])
+        )
+    pile.append(urwid.Columns(resp, dividechars=1))
+    return urwid.Pile(pile)
+
+
 def format_flow(f, focus, extended=False, hostheader=False, marked=False):
     d = dict(
         intercepted = f.intercepted,
diff --git a/mitmproxy/console/flowlist.py b/mitmproxy/console/flowlist.py
index bc523874..1af7d00f 100644
--- a/mitmproxy/console/flowlist.py
+++ b/mitmproxy/console/flowlist.py
@@ -5,6 +5,7 @@ import urwid
 import netlib.http.url
 from mitmproxy.console import common
 from mitmproxy.console import signals
+from mitmproxy.flow import export
 
 
 def _mkhelp():
@@ -13,10 +14,9 @@ def _mkhelp():
         ("A", "accept all intercepted flows"),
         ("a", "accept this intercepted flow"),
         ("b", "save request/response body"),
-        ("C", "clear flow list or eventlog"),
+        ("C", "export flow to clipboard"),
         ("d", "delete flow"),
         ("D", "duplicate flow"),
-        ("E", "export"),
         ("e", "toggle eventlog"),
         ("F", "toggle follow flow list"),
         ("l", "set limit filter pattern"),
@@ -24,13 +24,14 @@ def _mkhelp():
         ("m", "toggle flow mark"),
         ("M", "toggle marked flow view"),
         ("n", "create a new request"),
-        ("P", "copy flow to clipboard"),
+        ("E", "export flow to file"),
         ("r", "replay request"),
         ("U", "unmark all marked flows"),
         ("V", "revert changes to request"),
         ("w", "save flows "),
         ("W", "stream flows to file"),
         ("X", "kill and delete flow, even if it's mid-intercept"),
+        ("z", "clear flow list or eventlog"),
         ("tab", "tab between eventlog and flow list"),
         ("enter", "view flow"),
         ("|", "run script on this flow"),
@@ -52,7 +53,7 @@ class LogBufferBox(urwid.ListBox):
 
     def keypress(self, size, key):
         key = common.shortcuts(key)
-        if key == "C":
+        if key == "z":
             self.master.clear_events()
             key = None
         elif key == "G":
@@ -263,24 +264,24 @@ class ConnectionItem(urwid.WidgetWrap):
                 callback = self.master.run_script_once,
                 args = (self.flow,)
             )
-        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"),
-                    ("as locust code", "l"),
-                    ("as locust task", "t"),
-                ),
-                callback = common.export_prompt,
-                args = (self.flow,)
+                prompt = "Export to file",
+                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)
+            )
+        elif key == "C":
+            signals.status_prompt_onekey.send(
+                self,
+                prompt = "Export to clipboard",
+                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)
             )
         elif key == "b":
-            common.ask_save_body(None, self.master, self.state, self.flow)
+            common.ask_save_body(None, self.flow)
         else:
             return key
 
@@ -362,7 +363,7 @@ class FlowListBox(urwid.ListBox):
         if key == "A":
             self.master.accept_all()
             signals.flowlist_change.send(self)
-        elif key == "C":
+        elif key == "z":
             self.master.clear_flows()
         elif key == "e":
             self.master.toggle_eventlog()
diff --git a/mitmproxy/console/flowview.py b/mitmproxy/console/flowview.py
index c85a9f73..938c8e86 100644
--- a/mitmproxy/console/flowview.py
+++ b/mitmproxy/console/flowview.py
@@ -18,6 +18,7 @@ from mitmproxy.console import grideditor
 from mitmproxy.console import searchable
 from mitmproxy.console import signals
 from mitmproxy.console import tabs
+from mitmproxy.flow import export
 from netlib.http import Headers
 from netlib.http import status_codes
 
@@ -32,9 +33,9 @@ def _mkhelp():
         ("A", "accept all intercepted flows"),
         ("a", "accept this intercepted flow"),
         ("b", "save request/response body"),
+        ("C", "export flow to clipboard"),
         ("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"),
@@ -75,8 +76,7 @@ def _mkhelp():
          [("text", ": XML")]
          ),
         ("M", "change default body display mode"),
-        ("p", "previous flow"),
-        ("P", "copy request/response (content/headers) to clipboard"),
+        ("E", "export flow to file"),
         ("r", "replay request"),
         ("V", "revert changes to request"),
         ("v", "view body in external viewer"),
@@ -589,20 +589,6 @@ 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"),
-                    ("as locust code", "l"),
-                    ("as locust task", "t"),
-                ),
-                callback = common.export_prompt,
-                args = (self.flow,)
-            )
         elif key == "|":
             signals.status_prompt_path.send(
                 prompt = "Send flow to script",
@@ -610,7 +596,7 @@ class FlowView(tabs.Tabs):
                 args = (self.flow,)
             )
 
-        if not conn and key in set(list("befgmxvz")):
+        if not conn and key in set(list("befgmxvzEC")):
             signals.status_message.send(
                 message = "Tab to the request or response",
                 expire = 1
@@ -663,12 +649,6 @@ class FlowView(tabs.Tabs):
                 )
                 signals.flow_change.send(self, flow = self.flow)
                 signals.status_message.send(message="")
-            elif key == "P":
-                if self.tab_offset == TAB_REQ:
-                    scope = "q"
-                else:
-                    scope = "s"
-                common.ask_copy_part(scope, self.flow, self.master, self.state)
             elif key == "m":
                 p = list(contentviews.view_prompts)
                 p.insert(0, ("Clear", "C"))
@@ -679,6 +659,30 @@ class FlowView(tabs.Tabs):
                     callback = self.change_this_display_mode
                 )
                 key = None
+            elif key == "E":
+                if self.tab_offset == TAB_REQ:
+                    scope = "q"
+                else:
+                    scope = "s"
+                signals.status_prompt_onekey.send(
+                    self,
+                    prompt = "Export to file",
+                    keys = [(e[0], e[1]) for e in export.EXPORTERS],
+                    callback = common.export_to_clip_or_file,
+                    args = (scope, self.flow, common.ask_save_path)
+                )
+            elif key == "C":
+                if self.tab_offset == TAB_REQ:
+                    scope = "q"
+                else:
+                    scope = "s"
+                signals.status_prompt_onekey.send(
+                    self,
+                    prompt = "Export to clipboard",
+                    keys = [(e[0], e[1]) for e in export.EXPORTERS],
+                    callback = common.export_to_clip_or_file,
+                    args = (scope, self.flow, common.copy_to_clipboard_or_prompt)
+                )
             elif key == "x":
                 signals.status_prompt_onekey.send(
                     prompt = "Delete body",
diff --git a/mitmproxy/flow/export.py b/mitmproxy/flow/export.py
index deeeb998..731aaf0e 100644
--- a/mitmproxy/flow/export.py
+++ b/mitmproxy/flow/export.py
@@ -97,11 +97,6 @@ def python_code(flow):
     return code
 
 
-def raw_request(flow):
-    data = netlib.http.http1.assemble_request(flow.request)
-    return _native(data)
-
-
 def is_json(headers, content):
     # type: (netlib.http.Headers, bytes) -> bool
     if headers:
@@ -197,3 +192,18 @@ def locust_task(flow):
     task_code = code[start_task:end_task]
 
     return task_code
+
+
+def url(flow):
+    return flow.request.url
+
+
+EXPORTERS = [
+    ("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),
+]
diff --git a/test/mitmproxy/test_flow_export.py b/test/mitmproxy/test_flow_export.py
index e6d65e40..86ff937d 100644
--- a/test/mitmproxy/test_flow_export.py
+++ b/test/mitmproxy/test_flow_export.py
@@ -1,4 +1,3 @@
-from textwrap import dedent
 import re
 
 import netlib.tutils
@@ -70,41 +69,6 @@ class TestExportPythonCode():
         python_equals("data/test_flow_export/python_patch.py", export.python_code(flow))
 
 
-class TestRawRequest():
-    def test_get(self):
-        flow = tutils.tflow(req=req_get())
-        result = dedent("""
-            GET /path?a=foo&a=bar&b=baz HTTP/1.1\r
-            header: qvalue\r
-            content-length: 7\r
-            host: address:22\r
-            \r
-        """).strip(" ").lstrip()
-        assert export.raw_request(flow) == result
-
-    def test_post(self):
-        flow = tutils.tflow(req=req_post())
-        result = dedent("""
-            POST /path HTTP/1.1\r
-            host: address:22\r
-            \r
-            content
-        """).strip()
-        assert export.raw_request(flow) == result
-
-    def test_patch(self):
-        flow = tutils.tflow(req=req_patch())
-        result = dedent("""
-            PATCH /path?query=param HTTP/1.1\r
-            header: qvalue\r
-            content-length: 7\r
-            host: address:22\r
-            \r
-            content
-        """).strip()
-        assert export.raw_request(flow) == result
-
-
 class TestExportLocustCode():
     def test_get(self):
         flow = tutils.tflow(req=req_get())
@@ -153,3 +117,9 @@ class TestIsJson():
         headers = Headers(content_type="application/json")
         j = export.is_json(headers, b'{"name": "example", "email": "example@example.com"}')
         assert isinstance(j, dict)
+
+
+class TestURL():
+    def test_url(self):
+        flow = tutils.tflow()
+        assert export.url(flow) == "http://address:22/path"