diff --git a/mitmproxy/addons/core.py b/mitmproxy/addons/core.py index b482edbb0..46cff8b50 100644 --- a/mitmproxy/addons/core.py +++ b/mitmproxy/addons/core.py @@ -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))) diff --git a/mitmproxy/tools/console/flowview.py b/mitmproxy/tools/console/flowview.py index ea61f65f5..1916ef6ba 100644 --- a/mitmproxy/tools/console/flowview.py +++ b/mitmproxy/tools/console/flowview.py @@ -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", diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py index b2208c05a..4ff568a4e 100644 --- a/mitmproxy/tools/console/master.py +++ b/mitmproxy/tools/console/master.py @@ -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") diff --git a/test/mitmproxy/addons/test_core.py b/test/mitmproxy/addons/test_core.py index 64d0fa19d..302b78ae0 100644 --- a/test/mitmproxy/addons/test_core.py +++ b/test/mitmproxy/addons/test_core.py @@ -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"