command: flow.set

Use this to replace the flow edit components of flowview entirely.

4 files changed, 123 insertions, 134 deletions

diff --git a/mitmproxy/addons/core.py b/mitmproxy/addons/core.py
index b482edbb..46cff8b5 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 ea61f65f..1916ef6b 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 b2208c05..4ff568a4 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 64d0fa19..302b78ae 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"