Merge remote-tracking branch 'tekii/feature-334'

4 files changed, 167 insertions, 29 deletions

diff --git a/libmproxy/console/common.py b/libmproxy/console/common.py
index 3e6e5ccc..c0929d58 100644
--- a/libmproxy/console/common.py
+++ b/libmproxy/console/common.py
@@ -1,9 +1,13 @@
 from __future__ import absolute_import
 import urwid
 import urwid.util
+import os
 from .. import utils
 from ..protocol.http import CONTENT_MISSING
-
+try:
+    import pyperclip
+except:
+    pyperclip = False
 
 VIEW_LIST = 0
 VIEW_FLOW = 1
@@ -161,6 +165,87 @@ def raw_format_flow(f, focus, extended, padding):
     pile.append(urwid.Columns(resp, dividechars=1))
     return urwid.Pile(pile)
 
+## common save body parts
+def save_body(path, master, state, content):
+    if not path:
+        return
+    state.last_saveload = path
+    path = os.path.expanduser(path)
+    try:
+        with file(path, "wb") as f:
+            f.write(content)
+    except IOError, v:
+        master.statusbar.message(v.strerror)
+
+def ask_save_body(k, master, state, content):
+    if k == "y":
+        master.path_prompt(
+            "Save message content: ",
+            state.last_saveload,
+            save_body,
+            master,
+            state,
+            content,
+        )
+
+def which_body_save(k, master, state, flow):
+    if k == "q":
+        master.path_prompt(
+            "Save request content: ",
+            state.last_saveload,
+            save_body,
+            master,
+            state,
+            flow.request.get_decoded_content(),
+        )
+    elif k == "r":
+        if flow.response:
+            master.path_prompt(
+                "Save response content: ",
+                state.last_saveload,
+                save_body,
+                master,
+                state,
+                flow.response.get_decoded_content(),
+            )
+        else:
+            master.statusbar.message("Flow has no response")
+
+## common copy_message parts 
+def copy_message( k, master, state, message):
+    if not message:
+        # only response could be None
+        master.statusbar.message("Flow has no response")                
+        return
+
+    if pyperclip:
+        if k == "c":
+            try:
+                pyperclip.copy(message.get_decoded_content())
+            except TypeError:
+                master.prompt_onekey(
+                    "Cannot copy binary data to clipboard. Save as file?",
+                    (
+                        ("yes", "y"),
+                        ("no", "n"),
+                    ),
+                    ask_save_body,
+                    master,
+                    state,
+                    message.get_decoded_content(),
+                )
+        elif k == "h":
+            try:
+                pyperclip.copy(message.headers)
+            except TypeError:
+                master.statusbar.message("Error converting headers to text")
+        elif k == "u":
+            try:
+                pyperclip.copy(message.url)
+            except TypeError:
+                master.statusbar.message("Error copying url to clipboard")
+    else:
+        master.statusbar.message("No clipboard support on your system, sorry.")
 
 class FlowCache:
     @utils.LRUCache(200)
diff --git a/libmproxy/console/flowlist.py b/libmproxy/console/flowlist.py
index 102fa7b9..fbcf1052 100644
--- a/libmproxy/console/flowlist.py
+++ b/libmproxy/console/flowlist.py
@@ -7,11 +7,14 @@ 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 response(content/headers) to clipboard"),        
+        ("G", "copy request(content/headers/url) to clipboard"),                
         ("l", "set limit filter pattern"),
         ("L", "load saved flows"),
         ("r", "replay request"),
@@ -204,6 +207,43 @@ class ConnectionItem(common.WWrap):
                 self.master.run_script_once,
                 self.flow
             )
+        elif key == "g":
+            self.master.prompt_onekey(
+                "Copy Response",
+                (
+                    ("content", "c"),
+                    ("headers", "h"),
+                ),
+                common.copy_message,
+                self.master,
+                self.state,
+                self.flow.response,
+            )
+        elif key == "G":
+            self.master.prompt_onekey(
+                "Copy Request",
+                (
+                    ("content", "c"),
+                    ("headers", "h"),
+                    ("url", "u"),                                        
+                ),
+                common.copy_message,
+                self.master,
+                self.state,
+                self.flow.request,
+            )
+        elif key == "b":
+            self.master.prompt_onekey(
+                "Save",
+                (
+                    ("request", "q"),
+                    ("response", "r"),
+                ),
+                common.which_body_save,
+                self.master,
+                self.state,
+                self.flow,
+            )
         else:
             return key
 
diff --git a/libmproxy/console/flowview.py b/libmproxy/console/flowview.py
index 2aac575d..f95cd776 100644
--- a/libmproxy/console/flowview.py
+++ b/libmproxy/console/flowview.py
@@ -19,6 +19,8 @@ def _mkhelp():
         ("D", "duplicate flow"),
         ("e", "edit request/response"),
         ("f", "load full body data"),
+        ("g", "copy response(content/headers) to clipboard"),        
+        ("G", "copy request(content/headers/url) to clipboard"),        
         ("m", "change body display mode for this entity"),
             (None,
                 common.highlight_key("automatic", "a") +
@@ -508,22 +510,6 @@ class FlowView(common.WWrap):
                     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:
@@ -691,17 +677,20 @@ class FlowView(common.WWrap):
         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
-                    )
+                    msg = "Save request body: "
+                    content = self.flow.request.content
+                else: 
+                    msg = "Save response body: "
+                    content = self.flow.response.content                    
+
+                self.master.path_prompt(
+                    msg,
+                    self.state.last_saveload,
+                    common.save_body,
+                    self.master,
+                    self.state,
+                    content,
+                )
         elif key == "d":
             if self.state.flow_count() == 1:
                 self.master.view_flowlist()
@@ -752,6 +741,31 @@ class FlowView(common.WWrap):
             )
             self.master.refresh_flow(self.flow)
             self.master.statusbar.message("")
+        elif key == "g":
+            self.master.prompt_onekey(
+                "Copy Response",
+                (
+                    ("content", "c"),
+                    ("headers", "h"),
+                ),
+                common.copy_message,
+                self.master,
+                self.state,
+                self.flow.response,
+            )
+        elif key == "G":
+            self.master.prompt_onekey(
+                "Copy Request",
+                (
+                    ("content", "c"),
+                    ("headers", "h"),
+                    ("url", "u"),                    
+                ),
+                common.copy_message,
+                self.master,
+                self.state,
+                self.flow.request,
+            )
         elif key == "m":
             p = list(contentview.view_prompts)
             p.insert(0, ("Clear", "C"))
diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index bebb4f7b..284731b0 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -953,7 +953,6 @@ class HTTPFlow(Flow):
             c += self.response.replace(pattern, repl, *args, **kwargs)
         return c
 
-
 class HttpAuthenticationError(Exception):
     def __init__(self, auth_headers=None):
         super(HttpAuthenticationError, self).__init__(