From 6d2b823a5438672ece9b537bc6e3c2aefac6e7ac Mon Sep 17 00:00:00 2001 From: Brad Dixon Date: Fri, 28 May 2021 14:38:27 -0400 Subject: [PATCH] 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 | 1 + mitmproxy/addons/__init__.py | 2 ++ mitmproxy/addons/comment.py | 16 ++++++++++++++++ mitmproxy/flow.py | 2 ++ mitmproxy/flowfilter.py | 11 +++++++++++ mitmproxy/io/compat.py | 7 +++++++ mitmproxy/tools/console/defaultkeys.py | 1 + mitmproxy/tools/console/flowdetailview.py | 4 ++++ mitmproxy/version.py | 2 +- test/mitmproxy/addons/test_comment.py | 12 ++++++++++++ test/mitmproxy/test_flow.py | 2 ++ test/mitmproxy/test_flowfilter.py | 5 +++++ 12 files changed, 64 insertions(+), 1 deletion(-) create mode 100644 mitmproxy/addons/comment.py create mode 100644 test/mitmproxy/addons/test_comment.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 198e2a5b0..f73196422 100644 --- a/CHANGELOG.md +++ b/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 ` filter syntax to search flow comments. (@rbdixon) * --- TODO: add new PRs above this line --- * ... and various other fixes, documentation improvements, dependency version bumps, etc. diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py index 5615a64bc..5af458171 100644 --- a/mitmproxy/addons/__init__.py +++ b/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(), diff --git a/mitmproxy/addons/comment.py b/mitmproxy/addons/comment.py new file mode 100644 index 000000000..2239cca7b --- /dev/null +++ b/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)) diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py index 61dec3f8f..43e21d8a5 100644 --- a/mitmproxy/flow.py +++ b/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): diff --git a/mitmproxy/flowfilter.py b/mitmproxy/flowfilter.py index 00641fa9b..fea8e9e42 100644 --- a/mitmproxy/flowfilter.py +++ b/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 diff --git a/mitmproxy/io/compat.py b/mitmproxy/io/compat.py index c7e7b76ae..8afbbf5cc 100644 --- a/mitmproxy/io/compat.py +++ b/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, } diff --git a/mitmproxy/tools/console/defaultkeys.py b/mitmproxy/tools/console/defaultkeys.py index 0fbced832..6fcb5f5de 100644 --- a/mitmproxy/tools/console/defaultkeys.py +++ b/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") diff --git a/mitmproxy/tools/console/flowdetailview.py b/mitmproxy/tools/console/flowdetailview.py index 9ae16a5b2..2a0271b13 100644 --- a/mitmproxy/tools/console/flowdetailview.py +++ b/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()] diff --git a/mitmproxy/version.py b/mitmproxy/version.py index 726c7fcc7..bd09b2d5a 100644 --- a/mitmproxy/version.py +++ b/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: diff --git a/test/mitmproxy/addons/test_comment.py b/test/mitmproxy/addons/test_comment.py new file mode 100644 index 000000000..ba628cd49 --- /dev/null +++ b/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" diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index c6c8d20d6..cb1e351c6 100644 --- a/test/mitmproxy/test_flow.py +++ b/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) diff --git a/test/mitmproxy/test_flowfilter.py b/test/mitmproxy/test_flowfilter.py index 03c2eff61..5ab2e16c4 100644 --- a/test/mitmproxy/test_flowfilter.py +++ b/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):