s/part/key

This commit is contained in:
Shadab Zafar 2016-07-10 21:11:25 +05:30
parent e56df1b706
commit 3ae28f1c31
2 changed files with 22 additions and 27 deletions

View File

@ -7,11 +7,6 @@ import urwid.util
import six import six
import netlib import netlib
<<<<<<< a7bc7d29a484c87bc1576d33bf5ab23aa5031c43
from mitmproxy import flow
=======
from mitmproxy import models
>>>>>>> Common function for both file and clipboard
from mitmproxy import utils from mitmproxy import utils
from mitmproxy.console import signals from mitmproxy.console import signals
from mitmproxy.flow import export from mitmproxy.flow import export
@ -199,8 +194,8 @@ def copy_to_clipboard_or_prompt(data):
) )
def flow_format_data(part, scope, flow): def flow_format_data(key, scope, flow):
if part == "u": if key == "u":
data = flow.request.url data = flow.request.url
else: else:
data = "" data = ""
@ -209,12 +204,12 @@ def flow_format_data(part, scope, flow):
request.decode(strict=False) request.decode(strict=False)
if request.content is None: if request.content is None:
return None, "Request content is missing" return None, "Request content is missing"
if part == "h": if key == "h":
data += netlib.http.http1.assemble_request(request) data += netlib.http.http1.assemble_request(request)
elif part == "c": elif key == "c":
data += request.content data += request.content
else: else:
raise ValueError("Unknown part: {}".format(part)) raise ValueError("Unknown key: {}".format(key))
if scope == "a" and flow.request.raw_content and flow.response: if scope == "a" and flow.request.raw_content and flow.response:
# Add padding between request and response # Add padding between request and response
data += "\r\n" * 2 data += "\r\n" * 2
@ -223,21 +218,21 @@ def flow_format_data(part, scope, flow):
response.decode(strict=False) response.decode(strict=False)
if response.content is None: if response.content is None:
return None, "Response content is missing" return None, "Response content is missing"
if part == "h": if key == "h":
data += netlib.http.http1.assemble_response(response) data += netlib.http.http1.assemble_response(response)
elif part == "c": elif key == "c":
data += response.content data += response.content
else: else:
raise ValueError("Unknown part: {}".format(part)) raise ValueError("Unknown key: {}".format(key))
return data, False return data, False
def copy_flow(part, scope, flow): def copy_flow(key, scope, flow, writer):
""" """
part: _c_ontent, _h_eaders+content, _u_rl key: _c_ontent, _h_eaders+content, _u_rl
scope: _a_ll, re_q_uest, re_s_ponse scope: _a_ll, re_q_uest, re_s_ponse
""" """
data, err = flow_format_data(part, scope, flow) data, err = flow_format_data(key, scope, flow)
if err: if err:
signals.status_message.send(message=err) signals.status_message.send(message=err)
@ -252,7 +247,7 @@ def copy_flow(part, scope, flow):
signals.status_message.send(message="No contents to copy.") signals.status_message.send(message="No contents to copy.")
return return
copy_to_clipboard_or_prompt(data) writer(data)
def ask_copy_part(scope, flow): def ask_copy_part(scope, flow):
@ -271,17 +266,17 @@ def ask_copy_part(scope, flow):
) )
def ask_save_body(part, flow): def ask_save_body(scope, flow):
""" """
Save either the request or the response body to disk. Save either the request or the response body to disk.
'part' can either be "q" (request), "s" (response) or None (ask user if necessary). 'scope' can either be "q" (request), "s" (response) or None (ask user if necessary).
""" """
request_has_content = flow.request and flow.request.raw_content request_has_content = flow.request and flow.request.raw_content
response_has_content = flow.response and flow.response.raw_content response_has_content = flow.response and flow.response.raw_content
if part is None: if scope is None:
# We first need to determine whether we want to save the request or the # We first need to determine whether we want to save the request or the
# response content. # response content.
if request_has_content and response_has_content: if request_has_content and response_has_content:
@ -299,21 +294,21 @@ def ask_save_body(part, flow):
else: else:
ask_save_body("q", flow) ask_save_body("q", flow)
elif part == "q" and request_has_content: elif scope == "q" and request_has_content:
ask_save_path( ask_save_path(
flow.request.get_content(strict=False), flow.request.get_content(strict=False),
"Save request content", "Save request content to"
) )
elif part == "s" and response_has_content: elif scope == "s" and response_has_content:
ask_save_path( ask_save_path(
flow.response.get_content(strict=False), flow.response.get_content(strict=False),
"Save response content", "Save response content to"
) )
else: else:
signals.status_message.send(message="No content to save.") signals.status_message.send(message="No content to save.")
def export_to_clip_or_file(key, flow, writer): def export_to_clip_or_file(key, scope, flow, writer):
""" """
Export selected flow to clipboard or a file. Export selected flow to clipboard or a file.

View File

@ -270,7 +270,7 @@ class ConnectionItem(urwid.WidgetWrap):
prompt = "Export to file", prompt = "Export to file",
keys = [(e[0], e[1]) for e in export.EXPORTERS], keys = [(e[0], e[1]) for e in export.EXPORTERS],
callback = common.export_to_clip_or_file, callback = common.export_to_clip_or_file,
args = (self.flow, common.ask_save_path) args = ("a", self.flow, common.ask_save_path)
) )
elif key == "P": elif key == "P":
signals.status_prompt_onekey.send( signals.status_prompt_onekey.send(
@ -278,7 +278,7 @@ class ConnectionItem(urwid.WidgetWrap):
prompt = "Export to clipboard", prompt = "Export to clipboard",
keys = [(e[0], e[1]) for e in export.EXPORTERS], keys = [(e[0], e[1]) for e in export.EXPORTERS],
callback = common.export_to_clip_or_file, callback = common.export_to_clip_or_file,
args = (self.flow, common.copy_to_clipboard_or_prompt) args = ("a", self.flow, common.copy_to_clipboard_or_prompt)
) )
elif key == "b": elif key == "b":
common.ask_save_body(None, self.flow) common.ask_save_body(None, self.flow)