command: flow.set
Use this to replace the flow edit components of flowview entirely.
This commit is contained in:
Aldo Cortesi
parent
46373977e2
commit
670d1e408b
|
|
@ -4,6 +4,7 @@ from mitmproxy import ctx
|
|||
from mitmproxy import exceptions
|
||||
from mitmproxy import command
|
||||
from mitmproxy import flow
|
||||
from mitmproxy.net.http import status_codes
|
||||
|
||||
|
||||
class Core:
|
||||
|
|
@ -79,3 +80,76 @@ class Core:
|
|||
updated.append(f)
|
||||
ctx.log.alert("Reverted %s flows." % len(updated))
|
||||
ctx.master.addons.trigger("update", updated)
|
||||
|
||||
@command.command("flow.set.options")
|
||||
def flow_set_options(self) -> typing.Sequence[str]:
|
||||
return [
|
||||
"host",
|
||||
"status_code",
|
||||
"method",
|
||||
"path",
|
||||
"url",
|
||||
"reason",
|
||||
]
|
||||
|
||||
@command.command("flow.set")
|
||||
def flow_set(
|
||||
self,
|
||||
flows: typing.Sequence[flow.Flow], spec: str, sval: str
|
||||
) -> None:
|
||||
"""
|
||||
Quickly set a number of common values on flows.
|
||||
"""
|
||||
opts = self.flow_set_options()
|
||||
if spec not in opts:
|
||||
raise exceptions.CommandError(
|
||||
"Set spec must be one of: %s." % ", ".join(opts)
|
||||
)
|
||||
|
||||
val = sval # type: typing.Union[int, str]
|
||||
if spec == "status_code":
|
||||
try:
|
||||
val = int(val)
|
||||
except ValueError as v:
|
||||
raise exceptions.CommandError(
|
||||
"Status code is not an integer: %s" % val
|
||||
) from v
|
||||
|
||||
updated = []
|
||||
for f in flows:
|
||||
req = getattr(f, "request", None)
|
||||
rupdate = True
|
||||
if req:
|
||||
if spec == "method":
|
||||
req.method = val
|
||||
elif spec == "host":
|
||||
req.host = val
|
||||
elif spec == "path":
|
||||
req.path = val
|
||||
elif spec == "url":
|
||||
try:
|
||||
req.url = val
|
||||
except ValueError as e:
|
||||
raise exceptions.CommandError(
|
||||
"URL %s is invalid: %s" % (repr(val), e)
|
||||
) from e
|
||||
else:
|
||||
self.rupdate = False
|
||||
|
||||
resp = getattr(f, "response", None)
|
||||
supdate = True
|
||||
if resp:
|
||||
if spec == "status_code":
|
||||
resp.status_code = val
|
||||
if val in status_codes.RESPONSES:
|
||||
resp.reason = status_codes.RESPONSES[int(val)]
|
||||
elif spec == "reason":
|
||||
resp.reason = val
|
||||
else:
|
||||
supdate = False
|
||||
|
||||
if rupdate or supdate:
|
||||
updated.append(f)
|
||||
|
||||
ctx.master.addons.trigger("update", updated)
|
||||
ctx.log.alert("Set %s on %s flows." % (spec, len(updated)))
|
||||
|
|
|
|||
|
|
@ -9,7 +9,6 @@ import urwid
|
|||
from mitmproxy import contentviews
|
||||
from mitmproxy import exceptions
|
||||
from mitmproxy import http
|
||||
from mitmproxy.net.http import status_codes
|
||||
from mitmproxy.tools.console import common
|
||||
from mitmproxy.tools.console import flowdetailview
|
||||
from mitmproxy.tools.console import overlay
|
||||
|
|
@ -286,90 +285,6 @@ class FlowDetails(tabs.Tabs):
|
|||
]
|
||||
return searchable.Searchable(txt)
|
||||
|
||||
def set_method_raw(self, m):
|
||||
if m:
|
||||
self.flow.request.method = m
|
||||
signals.flow_change.send(self, flow = self.flow)
|
||||
|
||||
def edit_method(self, m):
|
||||
if m == "e":
|
||||
signals.status_prompt.send(
|
||||
prompt = "Method",
|
||||
text = self.flow.request.method,
|
||||
callback = self.set_method_raw
|
||||
)
|
||||
else:
|
||||
for i in common.METHOD_OPTIONS:
|
||||
if i[1] == m:
|
||||
self.flow.request.method = i[0].upper()
|
||||
signals.flow_change.send(self, flow = self.flow)
|
||||
|
||||
def set_url(self, url):
|
||||
request = self.flow.request
|
||||
try:
|
||||
request.url = str(url)
|
||||
except ValueError:
|
||||
return "Invalid URL."
|
||||
signals.flow_change.send(self, flow = self.flow)
|
||||
|
||||
def set_resp_status_code(self, status_code):
|
||||
try:
|
||||
status_code = int(status_code)
|
||||
except ValueError:
|
||||
return None
|
||||
self.flow.response.status_code = status_code
|
||||
if status_code in status_codes.RESPONSES:
|
||||
self.flow.response.reason = status_codes.RESPONSES[status_code]
|
||||
signals.flow_change.send(self, flow = self.flow)
|
||||
|
||||
def set_resp_reason(self, reason):
|
||||
self.flow.response.reason = reason
|
||||
signals.flow_change.send(self, flow = self.flow)
|
||||
|
||||
def edit(self, part):
|
||||
if self.tab_offset == TAB_REQ:
|
||||
message = self.flow.request
|
||||
else:
|
||||
if not self.flow.response:
|
||||
self.flow.response = http.HTTPResponse.make(200, b"")
|
||||
message = self.flow.response
|
||||
|
||||
self.flow.backup()
|
||||
if part == "r":
|
||||
# Fix an issue caused by some editors when editing a
|
||||
# request/response body. Many editors make it hard to save a
|
||||
# file without a terminating newline on the last line. When
|
||||
# editing message bodies, this can cause problems. For now, I
|
||||
# just strip the newlines off the end of the body when we return
|
||||
# from an editor.
|
||||
c = self.master.spawn_editor(message.get_content(strict=False) or b"")
|
||||
message.content = c.rstrip(b"\n")
|
||||
elif part == "u":
|
||||
signals.status_prompt.send(
|
||||
prompt = "URL",
|
||||
text = message.url,
|
||||
callback = self.set_url
|
||||
)
|
||||
elif part == "m" and message == self.flow.request:
|
||||
signals.status_prompt_onekey.send(
|
||||
prompt = "Method",
|
||||
keys = common.METHOD_OPTIONS,
|
||||
callback = self.edit_method
|
||||
)
|
||||
elif part == "o":
|
||||
signals.status_prompt.send(
|
||||
prompt = "Code",
|
||||
text = str(message.status_code),
|
||||
callback = self.set_resp_status_code
|
||||
)
|
||||
elif part == "m" and message == self.flow.response:
|
||||
signals.status_prompt.send(
|
||||
prompt = "Message",
|
||||
text = message.reason,
|
||||
callback = self.set_resp_reason
|
||||
)
|
||||
signals.flow_change.send(self, flow = self.flow)
|
||||
|
||||
def view_flow(self, flow):
|
||||
signals.pop_view_state.send(self)
|
||||
self.master.view_flow(flow, self.tab_offset)
|
||||
|
|
@ -453,39 +368,6 @@ class FlowDetails(tabs.Tabs):
|
|||
callback = self.master.run_script_once,
|
||||
args = (self.flow,)
|
||||
)
|
||||
elif key == "e":
|
||||
if self.tab_offset == TAB_REQ:
|
||||
signals.status_prompt_onekey.send(
|
||||
prompt="Edit request",
|
||||
keys=(
|
||||
("cookies", "c"),
|
||||
("query", "q"),
|
||||
("path", "p"),
|
||||
("url", "u"),
|
||||
("header", "h"),
|
||||
("form", "f"),
|
||||
("raw body", "r"),
|
||||
("method", "m"),
|
||||
),
|
||||
callback=self.edit
|
||||
)
|
||||
elif self.tab_offset == TAB_RESP:
|
||||
signals.status_prompt_onekey.send(
|
||||
prompt="Edit response",
|
||||
keys=(
|
||||
("cookies", "c"),
|
||||
("code", "o"),
|
||||
("message", "m"),
|
||||
("header", "h"),
|
||||
("raw body", "r"),
|
||||
),
|
||||
callback=self.edit
|
||||
)
|
||||
else:
|
||||
signals.status_message.send(
|
||||
message="Tab to the request or response",
|
||||
expire=1
|
||||
)
|
||||
elif key in set("bfgmxvzEC") and not conn:
|
||||
signals.status_message.send(
|
||||
message = "Tab to the request or response",
|
||||
|
|
|
|||
|
|
@ -22,7 +22,6 @@ from mitmproxy import flow
|
|||
from mitmproxy.addons import intercept
|
||||
from mitmproxy.addons import readfile
|
||||
from mitmproxy.addons import view
|
||||
from mitmproxy.tools.console import grideditor
|
||||
from mitmproxy.tools.console import keymap
|
||||
from mitmproxy.tools.console import overlay
|
||||
from mitmproxy.tools.console import palettes
|
||||
|
|
@ -148,10 +147,14 @@ class ConsoleAddon:
|
|||
"cookies",
|
||||
"form",
|
||||
"path",
|
||||
"method",
|
||||
"query",
|
||||
"reason",
|
||||
"request-headers",
|
||||
"response-headers",
|
||||
"status_code",
|
||||
"set-cookies",
|
||||
"url",
|
||||
]
|
||||
|
||||
@command.command("console.edit.focus")
|
||||
|
|
@ -173,6 +176,10 @@ class ConsoleAddon:
|
|||
self.master.switch_view("edit_focus_response_headers")
|
||||
elif part == "set-cookies":
|
||||
self.master.switch_view("edit_focus_setcookies")
|
||||
elif part in ["url", "method", "status_code", "reason"]:
|
||||
self.master.commands.call(
|
||||
"console.command flow.set @focus %s " % part
|
||||
)
|
||||
|
||||
def running(self):
|
||||
self.started = True
|
||||
|
|
@ -241,10 +248,10 @@ def default_keymap(km):
|
|||
km.add("enter", "console.view.flow @focus", context="flowlist")
|
||||
|
||||
km.add(
|
||||
"t",
|
||||
"e",
|
||||
"console.choose Part console.edit.focus.options "
|
||||
"console.edit.focus {choice}",
|
||||
context="flowlist"
|
||||
context="flowview"
|
||||
)
|
||||
|
||||
km.add(" ", "view.focus.next", context="flowview")
|
||||
|
|
@ -468,19 +475,6 @@ class ConsoleMaster(master.Master):
|
|||
def view_commands(self):
|
||||
self.window.push("commands")
|
||||
|
||||
def view_grideditor(self, ge):
|
||||
signals.push_view_state.send(
|
||||
self,
|
||||
window = window.Window(
|
||||
self,
|
||||
ge,
|
||||
None,
|
||||
statusbar.StatusBar(self, grideditor.base.FOOTER),
|
||||
ge.make_help(),
|
||||
"grideditor"
|
||||
)
|
||||
)
|
||||
|
||||
def view_flowlist(self):
|
||||
self.window.push("flowlist")
|
||||
|
||||
|
|
|
|||
|
|
@ -61,3 +61,42 @@ def test_revert():
|
|||
assert f.modified()
|
||||
sa.revert([f])
|
||||
assert not f.modified()
|
||||
|
||||
|
||||
def test_flow_set():
|
||||
sa = core.Core()
|
||||
with taddons.context():
|
||||
f = tflow.tflow(resp=True)
|
||||
assert sa.flow_set_options()
|
||||
|
||||
with pytest.raises(exceptions.CommandError):
|
||||
sa.flow_set([f], "flibble", "post")
|
||||
|
||||
assert f.request.method != "post"
|
||||
sa.flow_set([f], "method", "post")
|
||||
assert f.request.method == "POST"
|
||||
|
||||
assert f.request.host != "testhost"
|
||||
sa.flow_set([f], "host", "testhost")
|
||||
assert f.request.host == "testhost"
|
||||
|
||||
assert f.request.path != "/test/path"
|
||||
sa.flow_set([f], "path", "/test/path")
|
||||
assert f.request.path == "/test/path"
|
||||
|
||||
assert f.request.url != "http://foo.com/bar"
|
||||
sa.flow_set([f], "url", "http://foo.com/bar")
|
||||
assert f.request.url == "http://foo.com/bar"
|
||||
with pytest.raises(exceptions.CommandError):
|
||||
sa.flow_set([f], "url", "oink")
|
||||
|
||||
assert f.response.status_code != 404
|
||||
sa.flow_set([f], "status_code", "404")
|
||||
assert f.response.status_code == 404
|
||||
assert f.response.reason == "Not Found"
|
||||
with pytest.raises(exceptions.CommandError):
|
||||
sa.flow_set([f], "status_code", "oink")
|
||||
|
||||
assert f.response.reason != "foo"
|
||||
sa.flow_set([f], "reason", "foo")
|
||||
assert f.response.reason == "foo"
|
||||
|
|
|
|||
Loading…
Reference in New Issue