Merge branch 'master' of github.com:mitmproxy/mitmproxy

5 files changed, 233 insertions, 45 deletions

diff --git a/libmproxy/console/common.py b/libmproxy/console/common.py
index 3e6e5ccc..a2cfd57b 100644
--- a/libmproxy/console/common.py
+++ b/libmproxy/console/common.py
@@ -1,9 +1,14 @@
 from __future__ import absolute_import
 import urwid
 import urwid.util
+import os
 from .. import utils
-from ..protocol.http import CONTENT_MISSING
+from ..protocol.http import CONTENT_MISSING, decoded
 
+try:
+    import pyperclip
+except:
+    pyperclip = False
 
 VIEW_LIST = 0
 VIEW_FLOW = 1
@@ -12,6 +17,17 @@ VIEW_FLOW = 1
 VIEW_FLOW_REQUEST = 0
 VIEW_FLOW_RESPONSE = 1
 
+METHOD_OPTIONS = [
+    ("get", "g"),
+    ("post", "p"),
+    ("put", "u"),
+    ("head", "h"),
+    ("trace", "t"),
+    ("delete", "d"),
+    ("options", "o"),
+    ("edit raw", "e"),
+]
+
 
 def highlight_key(s, k):
     l = []
@@ -162,6 +178,146 @@ def raw_format_flow(f, focus, extended, padding):
     return urwid.Pile(pile)
 
 
+# Save file to disk
+def save_data(path, data, master, state):
+    if not path:
+        return
+    state.last_saveload = path
+    path = os.path.expanduser(path)
+    try:
+        with file(path, "wb") as f:
+            f.write(data)
+    except IOError, v:
+        master.statusbar.message(v.strerror)
+
+
+def ask_save_path(prompt, data, master, state):
+    master.path_prompt(
+        prompt,
+        state.last_saveload,
+        save_data,
+        data,
+        master,
+        state
+    )
+
+
+def copy_flow_format_data(part, scope, flow):
+    if part == "u":
+        data = flow.request.url
+    else:
+        data = ""
+        if scope in ("q", "a"):
+            with decoded(flow.request):
+                if part == "h":
+                    data += flow.request.assemble()
+                elif part == "c":
+                    data += flow.request.content
+                else:
+                    raise ValueError("Unknown part: {}".format(part))
+        if scope == "a" and flow.request.content and flow.response:
+            # Add padding between request and response
+            data += "\r\n" * 2
+        if scope in ("s", "a") and flow.response:
+            with decoded(flow.response):
+                if part == "h":
+                    data += flow.response.assemble()
+                elif part == "c":
+                    data += flow.response.content
+                else:
+                    raise ValueError("Unknown part: {}".format(part))
+    return data
+
+
+def copy_flow(part, scope, flow, master, state):
+    """
+    part: _c_ontent, _a_ll, _u_rl
+    scope: _a_ll, re_q_uest, re_s_ponse
+    """
+    data = copy_flow_format_data(part, scope, flow)
+
+    if not data:
+        if scope == "q":
+            master.statusbar.message("No request content to copy.")
+        elif scope == "s":
+            master.statusbar.message("No response content to copy.")
+        else:
+            master.statusbar.message("No contents to copy.")
+        return
+
+    try:
+        master.add_event(str(len(data)))
+        pyperclip.copy(data)
+    except RuntimeError:
+        def save(k):
+            if k == "y":
+                ask_save_path("Save data: ", data, master, state)
+
+        master.prompt_onekey(
+            "Cannot copy binary data to clipboard. Save as file?",
+            (
+                ("yes", "y"),
+                ("no", "n"),
+            ),
+            save
+        )
+
+
+def ask_copy_part(scope, flow, master, state):
+    choices = [
+        ("content", "c"),
+        ("headers+content", "h")
+    ]
+    if scope != "s":
+        choices.append(("url", "u"))
+
+    master.prompt_onekey(
+        "Copy",
+        choices,
+        copy_flow,
+        scope,
+        flow,
+        master,
+        state
+    )
+
+
+def ask_save_body(part, master, state, 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).
+    """
+
+    request_has_content = flow.request and flow.request.content
+    response_has_content = flow.response and flow.response.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:
+            master.prompt_onekey(
+                "Save",
+                (
+                    ("request", "q"),
+                    ("response", "s"),
+                ),
+                ask_save_body,
+                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:
+        ask_save_path("Save request content: ", flow.request.get_decoded_content(), master, state)
+    elif part == "s" and response_has_content:
+        ask_save_path("Save response content: ", flow.response.get_decoded_content(), master, state)
+    else:
+        master.statusbar.message("No content to save.")
+
+
 class FlowCache:
     @utils.LRUCache(200)
     def format_flow(self, *args):
@@ -211,7 +367,6 @@ def format_flow(f, focus, extended=False, hostheader=False, padding=2):
     return flowcache.format_flow(tuple(sorted(d.items())), focus, extended, padding)
 
 
-
 def int_version(v):
     SIG = 3
     v = urwid.__version__.split("-")[0].split(".")
diff --git a/libmproxy/console/flowlist.py b/libmproxy/console/flowlist.py
index 102fa7b9..9e7c6d69 100644
--- a/libmproxy/console/flowlist.py
+++ b/libmproxy/console/flowlist.py
@@ -1,5 +1,6 @@
 from __future__ import absolute_import
 import urwid
+from netlib import http
 from . import common
 
 def _mkhelp():
@@ -7,13 +8,16 @@ def _mkhelp():
     keys = [
         ("A", "accept all intercepted flows"),
         ("a", "accept this intercepted flow"),
+        ("b", "save request/response body"),
         ("C", "clear flow list or eventlog"),
         ("d", "delete flow"),
         ("D", "duplicate flow"),
         ("e", "toggle eventlog"),
         ("F", "toggle follow flow list"),
+        ("g", "copy flow to clipboard"),
         ("l", "set limit filter pattern"),
         ("L", "load saved flows"),
+        ("n", "create a new request"),
         ("r", "replay request"),
         ("V", "revert changes to request"),
         ("w", "save flows "),
@@ -204,6 +208,10 @@ class ConnectionItem(common.WWrap):
                 self.master.run_script_once,
                 self.flow
             )
+        elif key == "g":
+            common.ask_copy_part("a", self.flow, self.master, self.state)
+        elif key == "b":
+            common.ask_save_body(None, self.master, self.state, self.flow)
         else:
             return key
 
@@ -239,6 +247,32 @@ class FlowListBox(urwid.ListBox):
         self.master = master
         urwid.ListBox.__init__(self, master.flow_list_walker)
 
+    def get_method_raw(self, k):
+        if k:
+            self.get_url(k)    
+
+    def get_method(self, k):
+        if k == "e":
+            self.master.prompt("Method:", "", self.get_method_raw)
+        else:
+            method = ""
+            for i in common.METHOD_OPTIONS:
+                if i[1] == k:
+                    method = i[0].upper()
+            self.get_url(method)
+
+    def get_url(self,method):
+        self.master.prompt("URL:", "http://www.example.com/", self.new_request, method)
+
+    def new_request(self, url, method):
+        parts = http.parse_url(str(url))
+        if not parts:
+            self.master.statusbar.message("Invalid Url")
+            return
+        scheme, host, port, path = parts
+        f = self.master.create_request(method, scheme, host, port, path)
+        self.master.view_flow(f)
+
     def keypress(self, size, key):
         key = common.shortcuts(key)
         if key == "A":
@@ -256,6 +290,8 @@ class FlowListBox(urwid.ListBox):
                 self.master.state.last_saveload,
                 self.master.load_flows_callback
             )
+        elif key == "n":
+            self.master.prompt_onekey("Method", common.METHOD_OPTIONS, self.get_method)
         elif key == "F":
             self.master.toggle_follow_flows()
         elif key == "W":
diff --git a/libmproxy/console/flowview.py b/libmproxy/console/flowview.py
index 2aac575d..5c91512c 100644
--- a/libmproxy/console/flowview.py
+++ b/libmproxy/console/flowview.py
@@ -19,6 +19,7 @@ def _mkhelp():
         ("D", "duplicate flow"),
         ("e", "edit request/response"),
         ("f", "load full body data"),
+        ("g", "copy response(content/headers) to clipboard"),        
         ("m", "change body display mode for this entity"),
             (None,
                 common.highlight_key("automatic", "a") +
@@ -108,16 +109,6 @@ cache = CallbackCache()
 class FlowView(common.WWrap):
     REQ = 0
     RESP = 1
-    method_options = [
-        ("get", "g"),
-        ("post", "p"),
-        ("put", "u"),
-        ("head", "h"),
-        ("trace", "t"),
-        ("delete", "d"),
-        ("options", "o"),
-        ("edit raw", "e"),
-    ]
 
     highlight_color = "focusfield"
 
@@ -503,27 +494,11 @@ class FlowView(common.WWrap):
         if m == "e":
             self.master.prompt_edit("Method", self.flow.request.method, self.set_method_raw)
         else:
-            for i in self.method_options:
+            for i in common.METHOD_OPTIONS:
                 if i[1] == m:
                     self.flow.request.method = i[0].upper()
             self.master.refresh_flow(self.flow)
 
-    def save_body(self, path):
-        if not path:
-            return
-        self.state.last_saveload = path
-        if self.state.view_flow_mode == common.VIEW_FLOW_REQUEST:
-            c = self.flow.request
-        else:
-            c = self.flow.response
-        path = os.path.expanduser(path)
-        try:
-            f = file(path, "wb")
-            f.write(str(c.content))
-            f.close()
-        except IOError, v:
-            self.master.statusbar.message(v.strerror)
-
     def set_url(self, url):
         request = self.flow.request
         try:
@@ -614,7 +589,7 @@ class FlowView(common.WWrap):
         elif part == "u" and self.state.view_flow_mode == common.VIEW_FLOW_REQUEST:
             self.master.prompt_edit("URL", message.url, self.set_url)
         elif part == "m" and self.state.view_flow_mode == common.VIEW_FLOW_REQUEST:
-            self.master.prompt_onekey("Method", self.method_options, self.edit_method)
+            self.master.prompt_onekey("Method", common.METHOD_OPTIONS, self.edit_method)
         elif part == "c" and self.state.view_flow_mode == common.VIEW_FLOW_RESPONSE:
             self.master.prompt_edit("Code", str(message.code), self.set_resp_code)
         elif part == "m" and self.state.view_flow_mode == common.VIEW_FLOW_RESPONSE:
@@ -689,19 +664,10 @@ class FlowView(common.WWrap):
             self.master.accept_all()
             self.master.view_flow(self.flow)
         elif key == "b":
-            if conn:
-                if self.state.view_flow_mode == common.VIEW_FLOW_REQUEST:
-                    self.master.path_prompt(
-                        "Save request body: ",
-                        self.state.last_saveload,
-                        self.save_body
-                    )
-                else:
-                    self.master.path_prompt(
-                        "Save response body: ",
-                        self.state.last_saveload,
-                        self.save_body
-                    )
+            if self.state.view_flow_mode == common.VIEW_FLOW_REQUEST:
+                common.ask_save_body("q", self.master, self.state, self.flow)
+            else:
+                common.ask_save_body("s", self.master, self.state, self.flow)
         elif key == "d":
             if self.state.flow_count() == 1:
                 self.master.view_flowlist()
@@ -752,6 +718,12 @@ class FlowView(common.WWrap):
             )
             self.master.refresh_flow(self.flow)
             self.master.statusbar.message("")
+        elif key == "g":
+            if self.state.view_flow_mode == common.VIEW_FLOW_REQUEST:
+                scope = "q"
+            else:
+                scope = "s"
+            common.ask_copy_part(scope, self.flow, self.master, self.state)
         elif key == "m":
             p = list(contentview.view_prompts)
             p.insert(0, ("Clear", "C"))
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 49ec5a0b..14497964 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -8,12 +8,13 @@ import Cookie
 import cookielib
 import os
 import re
-from netlib import odict, wsgi
+from netlib import odict, wsgi, tcp
 import netlib.http
 from . import controller, protocol, tnetstring, filt, script, version
 from .onboarding import app
 from .protocol import http, handle
 from .proxy.config import HostMatcher
+from .proxy.connection import ClientConnection, ServerConnection
 import urlparse
 
 ODict = odict.ODict
@@ -763,6 +764,31 @@ class FlowMaster(controller.Master):
     def duplicate_flow(self, f):
         return self.load_flow(f.copy())
 
+    def create_request(self, method, scheme, host, port, path):
+        """
+            this method creates a new artificial and minimalist request also adds it to flowlist
+        """        
+        c = ClientConnection.from_state(dict(
+                address=dict(address=(host, port), use_ipv6=False),
+                clientcert=None
+            ))
+
+        s = ServerConnection.from_state(dict(
+                address=dict(address=(host, port), use_ipv6=False),
+                state=[],
+                source_address=None, #source_address=dict(address=(host, port), use_ipv6=False),
+                cert=None,
+                sni=host,
+                ssl_established=True
+            ))
+        f = http.HTTPFlow(c,s);
+        headers = ODictCaseless()
+        
+        req = http.HTTPRequest("absolute", method, scheme, host, port, path, (1, 1), headers, None,
+                                 None, None, None)
+        f.request = req
+        return self.load_flow(f)
+
     def load_flow(self, f):
         """
             Loads a flow, and returns a new flow object.
diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index 2f858a7c..046d0b42 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -1459,7 +1459,6 @@ class RequestReplayThread(threading.Thread):
                     if r.scheme == "https":
                         server.establish_ssl(self.config.clientcerts, sni=self.flow.server_conn.sni)
                     r.form_out = "relative"
-
                 server.send(r.assemble())
                 self.flow.server_conn = server
                 self.flow.response = HTTPResponse.from_stream(server.rfile, r.method,