Add flow.comment command and keybinding to add a comment to a flow. (#4608)

* Add `flow.comment` command and keybinding to add a comment to a flow.

* Store comment in Flow().comment. Add ~comment flowfilter syntax.

* resolve: Pythonic flow.comment

* Be consistent and use comment variable.

CHANGELOG.md

@@ -66,6 +66,7 @@ If you depend on these features, please raise your voice in
 * Fix parsing of certificate issuer/subject with escaped special characters (@Prinzhorn)
 * Customize markers with emoji, and filters: The `flow.mark` command may be used to mark a flow with either the default
   "red ball" marker, a single character, or an emoji like `:grapes:`. Use the `~marker` filter to filter on marker characters. (@rbdixon)
+* New `flow.comment` command to add a comment to the flow. Add `~comment <regex>` filter syntax to search flow comments. (@rbdixon)
 * --- TODO: add new PRs above this line ---
 * ... and various other fixes, documentation improvements, dependency version bumps, etc.
 

mitmproxy/addons/__init__.py

@@ -5,6 +5,7 @@ from mitmproxy.addons import blocklist
 from mitmproxy.addons import browser
 from mitmproxy.addons import clientplayback
 from mitmproxy.addons import command_history
+from mitmproxy.addons import comment
 from mitmproxy.addons import core
 from mitmproxy.addons import cut
 from mitmproxy.addons import disable_h2c
@@ -37,6 +38,7 @@ def default_addons():
         anticomp.AntiComp(),
         clientplayback.ClientPlayback(),
         command_history.CommandHistory(),
+        comment.Comment(),
         cut.Cut(),
         disable_h2c.DisableH2C(),
         export.Export(),

mitmproxy/addons/comment.py

@@ -0,0 +1,16 @@
+import typing
+from mitmproxy import command, flow, ctx
+from mitmproxy.hooks import UpdateHook
+
+
+class Comment:
+    @command.command("flow.comment")
+    def comment(self, flow: typing.Sequence[flow.Flow], comment: str) -> None:
+        "Add a comment to a flow"
+
+        updated = []
+        for f in flow:
+            f.comment = comment
+            updated.append(f)
+
+        ctx.master.addons.trigger(UpdateHook(updated))

mitmproxy/flow.py

@@ -123,6 +123,7 @@ class Flow(stateobject.StateObject):
         self.marked: str = ""
         self.is_replay: typing.Optional[str] = None
         self.metadata: typing.Dict[str, typing.Any] = dict()
+        self.comment: str = ""
 
     _stateobject_attributes = dict(
         id=str,
@@ -134,6 +135,7 @@ class Flow(stateobject.StateObject):
         is_replay=str,
         marked=str,
         metadata=typing.Dict[str, typing.Any],
+        comment=str,
     )
 
     def get_state(self):

mitmproxy/flowfilter.py

@@ -425,6 +425,16 @@ class FMarker(_Rex):
         return self.re.search(f.marked)
 
 
+class FComment(_Rex):
+    code = "comment"
+    help = "Flow comment"
+    flags = re.MULTILINE
+    is_binary = False
+
+    def __call__(self, f):
+        return self.re.search(f.comment)
+
+
 class _Int(_Action):
 
     def __init__(self, num):
@@ -512,6 +522,7 @@ filter_rex: Sequence[Type[_Rex]] = [
     FUrl,
     FMeta,
     FMarker,
+    FComment,
 ]
 filter_int = [
     FCode

mitmproxy/io/compat.py

@@ -309,6 +309,12 @@ def convert_12_13(data):
     return data
 
 
+def convert_13_14(data):
+    data["version"] = 14
+    data["comment"] = ""
+    return data
+
+
 def _convert_dict_keys(o: Any) -> Any:
     if isinstance(o, dict):
         return {strutils.always_str(k): _convert_dict_keys(v) for k, v in o.items()}
@@ -369,6 +375,7 @@ converters = {
     10: convert_10_11,
     11: convert_11_12,
     12: convert_12_13,
+    13: convert_13_14,
 }
 
 

mitmproxy/tools/console/defaultkeys.py

@@ -1,5 +1,6 @@
 def map(km):
     km.add(":", "console.command ", ["commonkey", "global"], "Command prompt")
+    km.add(";", "console.command flow.comment @focus ''", ["flowlist", "flowview"], "Add comment to flow")
     km.add("?", "console.view.help", ["global"], "View help")
     km.add("B", "browser.start", ["global"], "Start an attached browser")
     km.add("C", "console.view.commands", ["global"], "View commands")

mitmproxy/tools/console/flowdetailview.py

@@ -33,6 +33,10 @@ def flowdetails(state, flow: mitmproxy.flow.Flow):
         req = None
         resp = None
     metadata = flow.metadata
+    comment = flow.comment
+
+    if comment:
+        text.append(urwid.Text([("head", "Comment: "), ("text", comment)]))
 
     if metadata is not None and len(metadata) > 0:
         parts = [(str(k), repr(v)) for k, v in metadata.items()]

mitmproxy/version.py

@@ -7,7 +7,7 @@ MITMPROXY = "mitmproxy " + VERSION
 
 # Serialization format version. This is displayed nowhere, it just needs to be incremented by one
 # for each change in the file format.
-FLOW_FORMAT_VERSION = 13
+FLOW_FORMAT_VERSION = 14
 
 
 def get_dev_version() -> str:

test/mitmproxy/addons/test_comment.py

@@ -0,0 +1,12 @@
+from mitmproxy.test import tflow, taddons
+from mitmproxy.addons.comment import Comment
+
+
+def test_comment():
+    c = Comment()
+    f = tflow.tflow()
+
+    with taddons.context():
+        c.comment([f], "foo")
+
+    assert f.comment == "foo"

test/mitmproxy/test_flow.py

@@ -35,6 +35,8 @@ class TestSerialize:
         sio = io.BytesIO()
         f = tflow.tflow()
         f.marked = ":default:"
+        f.marked = True
+        f.comment = "test comment"
         f.request.content = bytes(range(256))
         w = mitmproxy.io.FlowWriter(sio)
         w.add(f)

test/mitmproxy/test_flowfilter.py

@@ -25,6 +25,7 @@ class TestParsing:
         assert flowfilter.parse("~replay")
         assert flowfilter.parse("~replayq")
         assert flowfilter.parse("~replays")
+        assert flowfilter.parse("~comment .")
         p = flowfilter.parse("~q ~c 10")
         self._dump(p)
         assert len(p.lst) == 2
@@ -615,6 +616,10 @@ class TestMatchingDummyFlow:
 
         assert not self.q("~q", f)
 
+        assert not self.q("~comment .", f)
+        f.comment = "comment"
+        assert self.q("~comment .", f)
+
 
 @patch('traceback.extract_tb')
 def test_pyparsing_bug(extract_tb):