Merge pull request #2686 from cortesi/flowspecopts

commander: add completion for flowspecs

2 files changed, 35 insertions, 1 deletions

diff --git a/mitmproxy/command.py b/mitmproxy/command.py
index 2d51317c..22c2376b 100644
--- a/mitmproxy/command.py
+++ b/mitmproxy/command.py
@@ -24,6 +24,32 @@ def lexer(s):
     return lex
 
 
+# This is an awkward location for these values, but it's better than having
+# the console core import and depend on an addon. FIXME: Add a way for
+# addons to add custom types and manage their completion and validation.
+valid_flow_prefixes = [
+    "@all",
+    "@focus",
+    "@shown",
+    "@hidden",
+    "@marked",
+    "@unmarked",
+    "~q",
+    "~s",
+    "~a",
+    "~hq",
+    "~hs",
+    "~b",
+    "~bq",
+    "~bs",
+    "~t",
+    "~d",
+    "~m",
+    "~u",
+    "~c",
+]
+
+
 Cuts = typing.Sequence[
     typing.Sequence[typing.Union[str, bytes]]
 ]
diff --git a/mitmproxy/tools/console/commander/commander.py b/mitmproxy/tools/console/commander/commander.py
index b94d6f69..ef32b953 100644
--- a/mitmproxy/tools/console/commander/commander.py
+++ b/mitmproxy/tools/console/commander/commander.py
@@ -6,6 +6,7 @@ import typing
 import urwid
 from urwid.text_layout import calc_coords
 
+import mitmproxy.flow
 import mitmproxy.master
 import mitmproxy.command
 
@@ -142,7 +143,14 @@ class CommandBuffer():
                     ),
                     parse = parts,
                 )
-
+            elif last.type in (typing.Sequence[mitmproxy.flow.Flow], mitmproxy.flow.Flow):
+                self.completion = CompletionState(
+                    completer = ListCompleter(
+                        "",
+                        mitmproxy.command.valid_flow_prefixes,
+                    ),
+                    parse = parts,
+                )
         if self.completion:
             nxt = self.completion.completer.cycle()
             buf = " ".join([i.value for i in self.completion.parse[:-1]]) + " " + nxt