diff --git a/libmproxy/console/common.py b/libmproxy/console/common.py index c0929d582..aaec81b52 100644 --- a/libmproxy/console/common.py +++ b/libmproxy/console/common.py @@ -165,87 +165,104 @@ 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): + +# 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(content) + f.write(data) 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, +def ask_save_path(prompt, data, master, state): + master.path_prompt( + prompt, + state.last_saveload, + save_data, + data, + 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.response.get_decoded_content(), + flow ) + elif response_has_content: + ask_save_body("s", master, state, flow) else: - master.statusbar.message("Flow has no response") + ask_save_body("q", master, state, flow) -## common copy_message parts -def copy_message( k, master, state, message): + 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.") + + +# 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") + 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.") + 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) diff --git a/libmproxy/console/flowlist.py b/libmproxy/console/flowlist.py index fbcf10521..e5468be19 100644 --- a/libmproxy/console/flowlist.py +++ b/libmproxy/console/flowlist.py @@ -217,7 +217,7 @@ class ConnectionItem(common.WWrap): common.copy_message, self.master, self.state, - self.flow.response, + self.flow.response ) elif key == "G": self.master.prompt_onekey( @@ -230,20 +230,10 @@ class ConnectionItem(common.WWrap): common.copy_message, self.master, self.state, - self.flow.request, + 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, - ) + common.ask_save_body(None, self.master, self.state, self.flow) else: return key diff --git a/libmproxy/console/flowview.py b/libmproxy/console/flowview.py index f95cd7766..4e604b00c 100644 --- a/libmproxy/console/flowview.py +++ b/libmproxy/console/flowview.py @@ -675,22 +675,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: - 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, - ) + 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()