added support for 'b' key in flowlist

2 files changed, 41 insertions, 0 deletions

diff --git a/libmproxy/console/common.py b/libmproxy/console/common.py
index 33738c0a..2f6c909c 100644
--- a/libmproxy/console/common.py
+++ b/libmproxy/console/common.py
@@ -189,8 +189,36 @@ def ask_save_body(k, master, state, content):
             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:
diff --git a/libmproxy/console/flowlist.py b/libmproxy/console/flowlist.py
index 6939be23..4acb17dc 100644
--- a/libmproxy/console/flowlist.py
+++ b/libmproxy/console/flowlist.py
@@ -7,6 +7,7 @@ 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"),
@@ -230,6 +231,18 @@ class ConnectionItem(common.WWrap):
                 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