From 5377d5a61dfdad9790c7b64118622c40ad405e73 Mon Sep 17 00:00:00 2001 From: Shadab Zafar Date: Sun, 10 Jul 2016 20:23:58 +0530 Subject: [PATCH] Refactor --- mitmproxy/console/common.py | 385 ++++++++++++++++++------------------ 1 file changed, 193 insertions(+), 192 deletions(-) diff --git a/mitmproxy/console/common.py b/mitmproxy/console/common.py index d84c4620d..033439143 100644 --- a/mitmproxy/console/common.py +++ b/mitmproxy/console/common.py @@ -134,6 +134,199 @@ else: SYMBOL_MARK = "[m]" +# Save file to disk +def save_data(path, data): + if not path: + return + try: + with open(path, "wb") as f: + f.write(data) + except IOError as v: + signals.status_message.send(message=v.strerror) + + +def ask_save_overwrite(path, data): + if not path: + return + path = os.path.expanduser(path) + if os.path.exists(path): + def save_overwrite(k): + if k == "y": + save_data(path, data) + + signals.status_prompt_onekey.send( + prompt = "'" + path + "' already exists. Overwrite?", + keys = ( + ("yes", "y"), + ("no", "n"), + ), + callback = save_overwrite + ) + else: + save_data(path, data) + + +def ask_save_path(data, prompt="File path"): + signals.status_prompt_path.send( + prompt = prompt, + callback = ask_save_overwrite, + args = (data, ) + ) + + +def copy_to_clipboard_or_prompt(data): + # pyperclip calls encode('utf-8') on data to be copied without checking. + # if data are already encoded that way UnicodeDecodeError is thrown. + toclip = "" + try: + toclip = data.decode('utf-8') + except (UnicodeDecodeError): + toclip = data + + try: + pyperclip.copy(toclip) + except (RuntimeError, UnicodeDecodeError, AttributeError, TypeError): + def save(k): + if k == "y": + ask_save_path(data, "Save data") + signals.status_prompt_onekey.send( + prompt = "Cannot copy data to clipboard. Save as file?", + keys = ( + ("yes", "y"), + ("no", "n"), + ), + callback = save + ) + + +def flow_format_data(part, scope, flow): + if part == "u": + data = flow.request.url + else: + data = "" + if scope in ("q", "a"): + request = flow.request.copy() + request.decode(strict=False) + if request.content is None: + return None, "Request content is missing" + if part == "h": + data += netlib.http.http1.assemble_request(request) + elif part == "c": + data += request.content + else: + raise ValueError("Unknown part: {}".format(part)) + if scope == "a" and flow.request.raw_content and flow.response: + # Add padding between request and response + data += "\r\n" * 2 + if scope in ("s", "a") and flow.response: + response = flow.response.copy() + response.decode(strict=False) + if response.content is None: + return None, "Response content is missing" + if part == "h": + data += netlib.http.http1.assemble_response(response) + elif part == "c": + data += response.content + else: + raise ValueError("Unknown part: {}".format(part)) + return data, False + + +def copy_flow(part, scope, flow, master, state): + """ + part: _c_ontent, _h_eaders+content, _u_rl + scope: _a_ll, re_q_uest, re_s_ponse + """ + data, err = flow_format_data(part, scope, flow) + + if err: + signals.status_message.send(message=err) + return + + if not data: + if scope == "q": + signals.status_message.send(message="No request content to copy.") + elif scope == "s": + signals.status_message.send(message="No response content to copy.") + else: + signals.status_message.send(message="No contents to copy.") + return + + copy_to_clipboard_or_prompt(data) + + +def ask_copy_part(scope, flow, master, state): + choices = [ + ("content", "c"), + ("headers+content", "h") + ] + if scope != "s": + choices.append(("url", "u")) + + signals.status_prompt_onekey.send( + prompt = "Copy", + keys = choices, + callback = copy_flow, + args = (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.raw_content + response_has_content = flow.response and flow.response.raw_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: + signals.status_prompt_onekey.send( + prompt = "Save", + keys = ( + ("request", "q"), + ("response", "s"), + ), + callback = ask_save_body, + args = (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( + flow.request.get_content(strict=False), + "Save request content", + ) + elif part == "s" and response_has_content: + ask_save_path( + flow.response.get_content(strict=False), + "Save response content", + ) + else: + signals.status_message.send(message="No content to save.") + + +def export_to_clip_or_file(key, flow, writer): + """ + Export selected flow to clipboard or a file. + + 'writer' is a function that handles the data + can be: copy_to_clipboard_or_prompt or ask_save_path + """ + for exporter in export.EXPORTERS: + if key == exporter[1]: + writer(exporter[2](flow)) + +flowcache = utils.LRUCache(800) + + def raw_format_flow(f, focus, extended): f = dict(f) pile = [] @@ -216,198 +409,6 @@ def raw_format_flow(f, focus, extended): return urwid.Pile(pile) -# Save file to disk -def save_data(path, data): - if not path: - return - try: - with open(path, "wb") as f: - f.write(data) - except IOError as v: - signals.status_message.send(message=v.strerror) - - -def ask_save_overwrite(path, data): - if not path: - return - path = os.path.expanduser(path) - if os.path.exists(path): - def save_overwrite(k): - if k == "y": - save_data(path, data) - - signals.status_prompt_onekey.send( - prompt = "'" + path + "' already exists. Overwrite?", - keys = ( - ("yes", "y"), - ("no", "n"), - ), - callback = save_overwrite - ) - else: - save_data(path, data) - - -def ask_save_path(data, prompt="File path"): - signals.status_prompt_path.send( - prompt = prompt, - callback = ask_save_overwrite, - args = (data, ) - ) - - -def copy_flow_format_data(part, scope, flow): - if part == "u": - data = flow.request.url - else: - data = "" - if scope in ("q", "a"): - request = flow.request.copy() - request.decode(strict=False) - if request.content is None: - return None, "Request content is missing" - if part == "h": - data += netlib.http.http1.assemble_request(request) - elif part == "c": - data += request.content - else: - raise ValueError("Unknown part: {}".format(part)) - if scope == "a" and flow.request.raw_content and flow.response: - # Add padding between request and response - data += "\r\n" * 2 - if scope in ("s", "a") and flow.response: - response = flow.response.copy() - response.decode(strict=False) - if response.content is None: - return None, "Response content is missing" - if part == "h": - data += netlib.http.http1.assemble_response(response) - elif part == "c": - data += response.content - else: - raise ValueError("Unknown part: {}".format(part)) - return data, False - - -def copy_to_clipboard_or_prompt(data): - # pyperclip calls encode('utf-8') on data to be copied without checking. - # if data are already encoded that way UnicodeDecodeError is thrown. - toclip = "" - try: - toclip = data.decode('utf-8') - except (UnicodeDecodeError): - toclip = data - - try: - pyperclip.copy(toclip) - except (RuntimeError, UnicodeDecodeError, AttributeError, TypeError): - def save(k): - if k == "y": - ask_save_path(data, "Save data") - signals.status_prompt_onekey.send( - prompt = "Cannot copy data to clipboard. Save as file?", - keys = ( - ("yes", "y"), - ("no", "n"), - ), - callback = save - ) - - -def copy_flow(part, scope, flow, master, state): - """ - part: _c_ontent, _h_eaders+content, _u_rl - scope: _a_ll, re_q_uest, re_s_ponse - """ - data, err = copy_flow_format_data(part, scope, flow) - - if err: - signals.status_message.send(message=err) - return - - if not data: - if scope == "q": - signals.status_message.send(message="No request content to copy.") - elif scope == "s": - signals.status_message.send(message="No response content to copy.") - else: - signals.status_message.send(message="No contents to copy.") - return - - copy_to_clipboard_or_prompt(data) - - -def ask_copy_part(scope, flow, master, state): - choices = [ - ("content", "c"), - ("headers+content", "h") - ] - if scope != "s": - choices.append(("url", "u")) - - signals.status_prompt_onekey.send( - prompt = "Copy", - keys = choices, - callback = copy_flow, - args = (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.raw_content - response_has_content = flow.response and flow.response.raw_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: - signals.status_prompt_onekey.send( - prompt = "Save", - keys = ( - ("request", "q"), - ("response", "s"), - ), - callback = ask_save_body, - args = (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( - flow.request.get_content(strict=False), - "Save request content", - ) - elif part == "s" and response_has_content: - ask_save_path( - flow.response.get_content(strict=False), - "Save response content", - ) - else: - signals.status_message.send(message="No content to save.") - - -def export_to_clip_or_file(key, flow, writer): - """ - Export selected flow to clipboard or a file. - - 'writer' is a function that handles the data - can be: copy_to_clipboard_or_prompt or ask_save_path - """ - for exporter in export.EXPORTERS: - if key == exporter[1]: - writer(exporter[2](flow)) - -flowcache = utils.LRUCache(800) - - def format_flow(f, focus, extended=False, hostheader=False, marked=False): d = dict( intercepted = f.intercepted,