From 0ee0cf566817fe842d3997445f6d891a5a1a90f6 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Sat, 4 Jul 2020 00:19:43 +0200 Subject: [PATCH] minor addon improvements, fix tests --- mitmproxy/addons/mapremote.py | 53 ++++------ mitmproxy/addons/modifybody.py | 51 +++------- mitmproxy/addons/modifyheaders.py | 106 +++++++++++--------- test/mitmproxy/addons/test_mapremote.py | 18 +++- test/mitmproxy/addons/test_modifybody.py | 25 +---- test/mitmproxy/addons/test_modifyheaders.py | 80 ++++++++------- 6 files changed, 151 insertions(+), 182 deletions(-) diff --git a/mitmproxy/addons/mapremote.py b/mitmproxy/addons/mapremote.py index 1c1a7b900..03f303da4 100644 --- a/mitmproxy/addons/mapremote.py +++ b/mitmproxy/addons/mapremote.py @@ -1,8 +1,7 @@ -import os import re import typing -from mitmproxy import exceptions +from mitmproxy import exceptions, http from mitmproxy import ctx from mitmproxy.addons.modifyheaders import parse_modify_spec, ModifySpec @@ -26,40 +25,26 @@ class MapRemote: self.replacements = [] for option in ctx.options.map_remote: try: - spec = parse_modify_spec(option) - try: - re.compile(spec.subject) - except re.error: - raise ValueError(f"Invalid regular expression: {spec.subject}") + spec = parse_modify_spec(option, True) except ValueError as e: - raise exceptions.OptionsError( - f"Cannot parse map_remote option {option}: {e}" - ) from e + raise exceptions.OptionsError(f"Cannot parse map_remote option {option}: {e}") from e self.replacements.append(spec) - def request(self, flow): - if not flow.reply.has_message: - for spec in self.replacements: - if spec.matches(flow): - self.replace(flow.request, spec.subject, spec.replacement) + def request(self, flow: http.HTTPFlow) -> None: + if flow.reply and flow.reply.has_message: + return + for spec in self.replacements: + if spec.matches(flow): + try: + replacement = spec.read_replacement() + except IOError as e: + ctx.log.warn(f"Could not read replacement file: {e}") + continue - def replace(self, obj, search, repl): - """ - Replaces all matches of the regex search in the url of the request with repl. - - Returns: - The number of replacements made. - """ - if repl.startswith(b"@"): - path = os.path.expanduser(repl[1:]) - try: - with open(path, "rb") as f: - repl = f.read() - except IOError: - ctx.log.warn("Could not read replacement file: %s" % repl) - return - - replacements = 0 - obj.url, replacements = re.subn(search, repl, obj.pretty_url.encode("utf8", "surrogateescape"), flags=re.DOTALL) - return replacements + url = flow.request.pretty_url.encode("utf8", "surrogateescape") + new_url = re.sub(spec.subject, replacement, url) + # this is a bit messy: setting .url also updates the host header, + # so we really only do that if the replacement affected the URL. + if url != new_url: + flow.request.url = new_url diff --git a/mitmproxy/addons/modifybody.py b/mitmproxy/addons/modifybody.py index 3fb728753..78d09f90b 100644 --- a/mitmproxy/addons/modifybody.py +++ b/mitmproxy/addons/modifybody.py @@ -1,4 +1,3 @@ -import os import re import typing @@ -26,26 +25,12 @@ class ModifyBody: self.replacements = [] for option in ctx.options.modify_body: try: - spec = parse_modify_spec(option) - try: - re.compile(spec.subject) - except re.error: - raise ValueError(f"Invalid regular expression: {spec.subject}") + spec = parse_modify_spec(option, True) except ValueError as e: - raise exceptions.OptionsError( - f"Cannot parse modify_body option {option}: {e}" - ) from e + raise exceptions.OptionsError(f"Cannot parse modify_body option {option}: {e}") from e self.replacements.append(spec) - def run(self, flow): - for spec in self.replacements: - if spec.matches(flow): - if flow.response: - self.replace(flow.response, spec.subject, spec.replacement) - else: - self.replace(flow.request, spec.subject, spec.replacement) - def request(self, flow): if not flow.reply.has_message: self.run(flow) @@ -54,23 +39,15 @@ class ModifyBody: if not flow.reply.has_message: self.run(flow) - def replace(self, obj, search, repl): - """ - Replaces all matches of the regex search in the body of the message with repl. - - Returns: - The number of replacements made. - """ - if repl.startswith(b"@"): - repl = os.path.expanduser(repl[1:]) - try: - with open(repl, "rb") as f: - repl = f.read() - except IOError: - ctx.log.warn("Could not read replacement file: %s" % repl) - return - - replacements = 0 - if obj.content: - obj.content, replacements = re.subn(search, repl, obj.content, flags=re.DOTALL) - return replacements + def run(self, flow): + for spec in self.replacements: + if spec.matches(flow): + try: + replacement = spec.read_replacement() + except IOError as e: + ctx.log.warn(f"Could not read replacement file: {e}") + continue + if flow.response: + flow.response.content = re.sub(spec.subject, replacement, flow.response.content, flags=re.DOTALL) + else: + flow.request.content = re.sub(spec.subject, replacement, flow.request.content, flags=re.DOTALL) diff --git a/mitmproxy/addons/modifyheaders.py b/mitmproxy/addons/modifyheaders.py index d560f5986..addac643a 100644 --- a/mitmproxy/addons/modifyheaders.py +++ b/mitmproxy/addons/modifyheaders.py @@ -1,26 +1,39 @@ -import os +import re import typing +from pathlib import Path -from mitmproxy import exceptions +from mitmproxy import exceptions, http from mitmproxy import flowfilter +from mitmproxy.net.http import Headers from mitmproxy.utils import strutils from mitmproxy import ctx class ModifySpec(typing.NamedTuple): - """ - match_str: a string specifying a flow filter pattern. - matches: the parsed match_str as a flowfilter.TFilter object - subject: a header name for ModifyHeaders and a regex pattern for ModifyBody - replacement: the replacement string - """ - match_str: str matches: flowfilter.TFilter subject: bytes - replacement: bytes + replacement_str: str + + def read_replacement(self) -> bytes: + """ + Process the replacement str. This usually just involves converting it to bytes. + However, if it starts with `@`, we interpret the rest as a file path to read from. + + Raises: + - IOError if the file cannot be read. + """ + if self.replacement_str.startswith("@"): + return Path(self.replacement_str[1:]).expanduser().read_bytes() + else: + # We could cache this at some point, but unlikely to be a problem. + return strutils.escaped_str_to_bytes(self.replacement_str) -def parse_modify_spec(option) -> ModifySpec: +def _match_all(flow) -> bool: + return True + + +def parse_modify_spec(option, subject_is_regex: bool) -> ModifySpec: """ The form for the modify_* options is as follows: @@ -48,24 +61,31 @@ def parse_modify_spec(option) -> ModifySpec: sep, rem = option[0], option[1:] parts = rem.split(sep, 2) if len(parts) == 2: - flow_filter_pattern = ".*" + flow_filter = _match_all subject, replacement = parts elif len(parts) == 3: flow_filter_pattern, subject, replacement = parts + flow_filter = flowfilter.parse(flow_filter_pattern) # type: ignore + if not flow_filter: + raise ValueError(f"Invalid filter pattern: {flow_filter_pattern}") else: raise ValueError("Invalid number of parameters (2 or 3 are expected)") - flow_filter = flowfilter.parse(flow_filter_pattern) - if not flow_filter: - raise ValueError(f"Invalid filter pattern: {flow_filter_pattern}") - subject = strutils.escaped_str_to_bytes(subject) - replacement = strutils.escaped_str_to_bytes(replacement) + if subject_is_regex: + try: + re.compile(subject) + except re.error as e: + raise ValueError(f"Invalid regular expression {subject!r} ({e})") - if replacement.startswith(b"@") and not os.path.isfile(os.path.expanduser(replacement[1:])): - raise ValueError(f"Invalid file path: {replacement[1:]}") + spec = ModifySpec(flow_filter, subject, replacement) - return ModifySpec(flow_filter_pattern, flow_filter, subject, replacement) + try: + spec.read_replacement() + except IOError as e: + raise ValueError(f"Invalid file path: {replacement[1:]} ({e})") + + return spec class ModifyHeaders: @@ -87,35 +107,11 @@ class ModifyHeaders: if "modify_headers" in updated: for option in ctx.options.modify_headers: try: - spec = parse_modify_spec(option) + spec = parse_modify_spec(option, False) except ValueError as e: - raise exceptions.OptionsError( - f"Cannot parse modify_headers option {option}: {e}" - ) from e + raise exceptions.OptionsError(f"Cannot parse modify_headers option {option}: {e}") from e self.replacements.append(spec) - def run(self, flow, hdrs): - # unset all specified headers - for spec in self.replacements: - if spec.matches(flow): - hdrs.pop(spec.subject, None) - - # set all specified headers if the replacement string is not empty - for spec in self.replacements: - if spec.replacement.startswith(b"@"): - path = os.path.expanduser(spec.replacement[1:]) - try: - with open(path, "rb") as file: - replacement = file.read() - except IOError: - ctx.log.warn(f"Could not read replacement file {path}") - return - else: - replacement = spec.replacement - - if spec.matches(flow) and replacement: - hdrs.add(spec.subject, replacement) - def request(self, flow): if not flow.reply.has_message: self.run(flow, flow.request.headers) @@ -123,3 +119,21 @@ class ModifyHeaders: def response(self, flow): if not flow.reply.has_message: self.run(flow, flow.response.headers) + + def run(self, flow: http.HTTPFlow, hdrs: Headers) -> None: + # unset all specified headers + for spec in self.replacements: + if spec.matches(flow): + hdrs.pop(spec.subject, None) + + # set all specified headers if the replacement string is not empty + for spec in self.replacements: + if spec.matches(flow): + try: + replacement = spec.read_replacement() + except IOError as e: + ctx.log.warn(f"Could not read replacement file: {e}") + continue + else: + if replacement: + hdrs.add(spec.subject, replacement) diff --git a/test/mitmproxy/addons/test_mapremote.py b/test/mitmproxy/addons/test_mapremote.py index b03af7817..970682b0f 100644 --- a/test/mitmproxy/addons/test_mapremote.py +++ b/test/mitmproxy/addons/test_mapremote.py @@ -12,7 +12,7 @@ class TestMapRemote: with taddons.context(mr) as tctx: tctx.configure(mr, map_remote=["one/two/three"]) with pytest.raises(Exception, match="Cannot parse map_remote .* Invalid number"): - tctx.configure(mr, map_remote = ["/"]) + tctx.configure(mr, map_remote=["/"]) with pytest.raises(Exception, match="Cannot parse map_remote .* Invalid filter"): tctx.configure(mr, map_remote=["/~b/two/three"]) with pytest.raises(Exception, match="Cannot parse map_remote .* Invalid regular expression"): @@ -33,6 +33,16 @@ class TestMapRemote: mr.request(f) assert f.request.url == "https://mitmproxy.org/img/test.jpg" + def test_has_reply(self): + mr = mapremote.MapRemote() + with taddons.context(mr) as tctx: + tctx.configure(mr, map_remote=[":example.org:mitmproxy.org"]) + f = tflow.tflow() + f.request.url = b"https://example.org/images/test.jpg" + f.kill() + mr.request(f) + assert f.request.url == "https://example.org/images/test.jpg" + class TestMapRemoteFile: def test_simple(self, tmpdir): @@ -42,7 +52,7 @@ class TestMapRemoteFile: tmpfile.write("mitmproxy.org") tctx.configure( mr, - map_remote=[":example.org:@" + str(tmpfile)] + map_remote=["|example.org|@" + str(tmpfile)] ) f = tflow.tflow() f.request.url = b"https://example.org/test" @@ -63,10 +73,10 @@ class TestMapRemoteFile: tmpfile.write("mitmproxy.org") tctx.configure( mr, - map_remote=[":example.org:@" + str(tmpfile)] + map_remote=["|example.org|@" + str(tmpfile)] ) tmpfile.remove() f = tflow.tflow() f.request.url = b"https://example.org/test" mr.request(f) - assert await tctx.master.await_log("could not read") \ No newline at end of file + assert await tctx.master.await_log("could not read") diff --git a/test/mitmproxy/addons/test_modifybody.py b/test/mitmproxy/addons/test_modifybody.py index 6431e4f84..f78932e52 100644 --- a/test/mitmproxy/addons/test_modifybody.py +++ b/test/mitmproxy/addons/test_modifybody.py @@ -1,38 +1,17 @@ import pytest from mitmproxy.addons import modifybody -from mitmproxy.addons.modifyheaders import parse_modify_spec from mitmproxy.test import taddons from mitmproxy.test import tflow class TestModifyBody: - def test_parse_modify_spec(self): - x = parse_modify_spec("/foo/bar/voing") - assert [x[0], x[2], x[3]] == ["foo", b"bar", b"voing"] - - x = parse_modify_spec("/foo/bar/vo/ing/") - assert [x[0], x[2], x[3]] == ["foo", b"bar", b"vo/ing/"] - - x = parse_modify_spec("/bar/voing") - assert [x[0], x[2], x[3]] == [".*", b"bar", b"voing"] - - with pytest.raises(Exception, match="Invalid number of parameters"): - parse_modify_spec("/") - - with pytest.raises(Exception, match="Invalid filter pattern"): - parse_modify_spec("/~b/one/two") - def test_configure(self): mb = modifybody.ModifyBody() with taddons.context(mb) as tctx: tctx.configure(mb, modify_body=["one/two/three"]) - with pytest.raises(Exception, match="Cannot parse modify_body .* Invalid number"): - tctx.configure(mb, modify_body = ["/"]) - with pytest.raises(Exception, match="Cannot parse modify_body .* Invalid filter"): - tctx.configure(mb, modify_body=["/~b/two/three"]) - with pytest.raises(Exception, match="Cannot parse modify_body .* Invalid regular expression"): - tctx.configure(mb, modify_body=["/foo/+/three"]) + with pytest.raises(Exception, match="Cannot parse modify_body"): + tctx.configure(mb, modify_body=["/"]) tctx.configure(mb, modify_body=["/a/b/c/"]) def test_simple(self): diff --git a/test/mitmproxy/addons/test_modifyheaders.py b/test/mitmproxy/addons/test_modifyheaders.py index bf75d4315..f3ebd279f 100644 --- a/test/mitmproxy/addons/test_modifyheaders.py +++ b/test/mitmproxy/addons/test_modifyheaders.py @@ -3,49 +3,53 @@ import pytest from mitmproxy.test import tflow from mitmproxy.test import taddons -from mitmproxy.addons import modifyheaders +from mitmproxy.addons.modifyheaders import parse_modify_spec, ModifyHeaders + + +def test_parse_modify_spec(): + spec = parse_modify_spec("/foo/bar/voing", True) + assert spec.matches.pattern == "foo" + assert spec.subject == b"bar" + assert spec.read_replacement() == b"voing" + + spec = parse_modify_spec("/foo/bar/vo/ing/", False) + assert spec.matches.pattern == "foo" + assert spec.subject == b"bar" + assert spec.read_replacement() == b"vo/ing/" + + spec = parse_modify_spec("/bar/voing", False) + assert spec.matches(tflow.tflow()) + assert spec.subject == b"bar" + assert spec.read_replacement() == b"voing" + + with pytest.raises(ValueError, match="Invalid number of parameters"): + parse_modify_spec("/", False) + + with pytest.raises(ValueError, match="Invalid filter pattern"): + parse_modify_spec("/~b/one/two", False) + + with pytest.raises(ValueError, match="Invalid filter pattern"): + parse_modify_spec("/~b/one/two", False) + + with pytest.raises(ValueError, match="Invalid regular expression"): + parse_modify_spec("/[/two", True) class TestModifyHeaders: - def test_parse_modify_spec(self): - x = modifyheaders.parse_modify_spec("/foo/bar/voing") - assert [x[0], x[2], x[3]] == ["foo", b"bar", b"voing"] - - x = modifyheaders.parse_modify_spec("/foo/bar/vo/ing/") - assert [x[0], x[2], x[3]] == ["foo", b"bar", b"vo/ing/"] - - x = modifyheaders.parse_modify_spec("/bar/voing") - assert [x[0], x[2], x[3]] == [".*", b"bar", b"voing"] - - with pytest.raises(Exception, match="Invalid number of parameters"): - modifyheaders.parse_modify_spec("/") - - with pytest.raises(Exception, match="Invalid filter pattern"): - modifyheaders.parse_modify_spec("/~b/one/two") - - with pytest.raises(Exception, match="Invalid file path"): - modifyheaders.parse_modify_spec("/~q/foo/@nonexistent") def test_configure(self): - mh = modifyheaders.ModifyHeaders() + mh = ModifyHeaders() with taddons.context(mh) as tctx: - with pytest.raises(Exception, match="Cannot parse modify_headers .* Invalid number"): - tctx.configure(mh, modify_headers = ["/"]) - - with pytest.raises(Exception, match="Cannot parse modify_headers .* Invalid filter"): - tctx.configure(mh, modify_headers = ["/~b/one/two"]) - - with pytest.raises(Exception, match="Cannot parse modify_headers .* Invalid file"): - tctx.configure(mh, modify_headers = ["/~q/foo/@nonexistent"]) - - tctx.configure(mh, modify_headers = ["/foo/bar/voing"]) + with pytest.raises(Exception, match="Cannot parse modify_headers"): + tctx.configure(mh, modify_headers=["/"]) + tctx.configure(mh, modify_headers=["/foo/bar/voing"]) def test_modify_headers(self): - mh = modifyheaders.ModifyHeaders() + mh = ModifyHeaders() with taddons.context(mh) as tctx: tctx.configure( mh, - modify_headers = [ + modify_headers=[ "/~q/one/two", "/~s/one/three" ] @@ -62,7 +66,7 @@ class TestModifyHeaders: tctx.configure( mh, - modify_headers = [ + modify_headers=[ "/~s/one/two", "/~s/one/three" ] @@ -75,7 +79,7 @@ class TestModifyHeaders: tctx.configure( mh, - modify_headers = [ + modify_headers=[ "/~q/one/two", "/~q/one/three" ] @@ -88,7 +92,7 @@ class TestModifyHeaders: # test removal of existing headers tctx.configure( mh, - modify_headers = [ + modify_headers=[ "/~q/one/", "/~s/one/" ] @@ -105,7 +109,7 @@ class TestModifyHeaders: tctx.configure( mh, - modify_headers = [ + modify_headers=[ "/one/" ] ) @@ -122,7 +126,7 @@ class TestModifyHeaders: class TestModifyHeadersFile: def test_simple(self, tmpdir): - mh = modifyheaders.ModifyHeaders() + mh = ModifyHeaders() with taddons.context(mh) as tctx: tmpfile = tmpdir.join("replacement") tmpfile.write("two") @@ -137,7 +141,7 @@ class TestModifyHeadersFile: @pytest.mark.asyncio async def test_nonexistent(self, tmpdir): - mh = modifyheaders.ModifyHeaders() + mh = ModifyHeaders() with taddons.context(mh) as tctx: with pytest.raises(Exception, match="Cannot parse modify_headers .* Invalid file path"): tctx.configure(