make copy context-sensitive

3 files changed, 89 insertions, 90 deletions

diff --git a/libmproxy/console/common.py b/libmproxy/console/common.py
index aaec81b5..e4a4acba 100644
--- a/libmproxy/console/common.py
+++ b/libmproxy/console/common.py
@@ -3,7 +3,8 @@ 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:
@@ -190,6 +191,86 @@ def ask_save_path(prompt, 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.
@@ -226,44 +307,6 @@ def ask_save_body(part, master, state, flow):
         master.statusbar.message("No content to save.")
 
 
-# common copy_message parts
-def copy_message(k, master, state, message):
-    if not pyperclip:
-        master.statusbar.message("No clipboard support on your system, sorry.")
-        return None
-    if not message:
-        # only response could be None
-        master.statusbar.message("Flow has no response")
-        return
-
-    data = None
-    if k == "c":
-        data = message.get_decoded_content()
-    elif k == "h":
-        data = message.headers
-    elif k == "u":
-        data = message.url
-
-    if not data:
-        master.statusbar.message("No content to copy.")
-        return
-
-    try:
-        pyperclip.copy(data)
-    except TypeError:
-        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
-        )
-
-
 class FlowCache:
     @utils.LRUCache(200)
     def format_flow(self, *args):
@@ -313,7 +356,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 e5468be1..c5cef061 100644
--- a/libmproxy/console/flowlist.py
+++ b/libmproxy/console/flowlist.py
@@ -13,8 +13,7 @@ def _mkhelp():
         ("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"),                
+        ("g", "copy flow to clipboard"),
         ("l", "set limit filter pattern"),
         ("L", "load saved flows"),
         ("r", "replay request"),
@@ -208,30 +207,7 @@ class ConnectionItem(common.WWrap):
                 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
-            )
+            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:
diff --git a/libmproxy/console/flowview.py b/libmproxy/console/flowview.py
index 4e604b00..e04d6bcc 100644
--- a/libmproxy/console/flowview.py
+++ b/libmproxy/console/flowview.py
@@ -730,30 +730,11 @@ 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,
-            )
+            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"))