diff options
| -rw-r--r-- | mitmproxy/builtins/__init__.py | 2 | ||||
| -rw-r--r-- | mitmproxy/builtins/setheaders.py | 39 | ||||
| -rw-r--r-- | mitmproxy/console/master.py | 4 | ||||
| -rw-r--r-- | mitmproxy/console/options.py | 12 | ||||
| -rw-r--r-- | mitmproxy/console/statusbar.py | 2 | ||||
| -rw-r--r-- | mitmproxy/dump.py | 4 | ||||
| -rw-r--r-- | mitmproxy/flow/__init__.py | 8 | ||||
| -rw-r--r-- | mitmproxy/flow/master.py | 5 | ||||
| -rw-r--r-- | mitmproxy/flow/modules.py | 55 | ||||
| -rw-r--r-- | test/mitmproxy/builtins/test_setheaders.py | 64 | ||||
| -rw-r--r-- | test/mitmproxy/test_flow.py | 52 |
11 files changed, 114 insertions, 133 deletions
diff --git a/mitmproxy/builtins/__init__.py b/mitmproxy/builtins/__init__.py index 1f8ae862..3974d736 100644 --- a/mitmproxy/builtins/__init__.py +++ b/mitmproxy/builtins/__init__.py @@ -7,6 +7,7 @@ from mitmproxy.builtins import stickyauth from mitmproxy.builtins import stickycookie from mitmproxy.builtins import script from mitmproxy.builtins import replace +from mitmproxy.builtins import setheaders def default_addons(): @@ -18,4 +19,5 @@ def default_addons(): script.ScriptLoader(), filestreamer.FileStreamer(), replace.Replace(), + setheaders.SetHeaders(), ] diff --git a/mitmproxy/builtins/setheaders.py b/mitmproxy/builtins/setheaders.py new file mode 100644 index 00000000..6bda3f55 --- /dev/null +++ b/mitmproxy/builtins/setheaders.py @@ -0,0 +1,39 @@ +from mitmproxy import exceptions +from mitmproxy import filt + + +class SetHeaders: + def __init__(self): + self.lst = [] + + def configure(self, options): + """ + options.setheaders is a tuple of (fpatt, header, value) + + fpatt: String specifying a filter pattern. + header: Header name. + value: Header value string + """ + for fpatt, header, value in options.setheaders: + cpatt = filt.parse(fpatt) + if not cpatt: + raise exceptions.OptionsError( + "Invalid setheader filter pattern %s" % fpatt + ) + self.lst.append((fpatt, header, value, cpatt)) + + def run(self, f, hdrs): + for _, header, value, cpatt in self.lst: + if cpatt(f): + hdrs.pop(header, None) + for _, header, value, cpatt in self.lst: + if cpatt(f): + hdrs.add(header, value) + + def request(self, flow): + if not flow.reply.acked: + self.run(flow, flow.request.headers) + + def response(self, flow): + if not flow.reply.acked: + self.run(flow, flow.response.headers) diff --git a/mitmproxy/console/master.py b/mitmproxy/console/master.py index 0ef12001..a93eeddf 100644 --- a/mitmproxy/console/master.py +++ b/mitmproxy/console/master.py @@ -210,10 +210,6 @@ class ConsoleMaster(flow.FlowMaster): self.options = self.options # type: Options self.options.errored.connect(self.options_error) - if options.setheaders: - for i in options.setheaders: - self.setheaders.add(*i) - r = self.set_intercept(options.intercept) if r: print("Intercept error: {}".format(r), file=sys.stderr) diff --git a/mitmproxy/console/options.py b/mitmproxy/console/options.py index f0cc4ef5..afb9186d 100644 --- a/mitmproxy/console/options.py +++ b/mitmproxy/console/options.py @@ -36,7 +36,7 @@ class Options(urwid.WidgetWrap): select.Option( "Header Set Patterns", "H", - lambda: master.setheaders.count(), + lambda: len(master.options.setheaders), self.setheaders ), select.Option( @@ -156,7 +156,6 @@ class Options(urwid.WidgetWrap): self.master.showhost = False self.master.refresh_server_playback = True self.master.server.config.no_upstream_cert = False - self.master.setheaders.clear() self.master.set_ignore_filter([]) self.master.set_tcp_filter([]) @@ -165,6 +164,7 @@ class Options(urwid.WidgetWrap): anticomp = False, replacements = [], scripts = [], + setheaders = [], stickyauth = None, stickycookie = None ) @@ -197,13 +197,12 @@ class Options(urwid.WidgetWrap): signals.update_settings.send(self) def setheaders(self): - def _set(*args, **kwargs): - self.master.setheaders.set(*args, **kwargs) - signals.update_settings.send(self) + def _set(shdrs): + self.master.options.setheaders = shdrs self.master.view_grideditor( grideditor.SetHeadersEditor( self.master, - self.master.setheaders.get_specs(), + self.master.options.setheaders, _set ) ) @@ -211,7 +210,6 @@ class Options(urwid.WidgetWrap): def ignorepatterns(self): def _set(ignore): self.master.set_ignore_filter(ignore) - signals.update_settings.send(self) self.master.view_grideditor( grideditor.HostPatternEditor( self.master, diff --git a/mitmproxy/console/statusbar.py b/mitmproxy/console/statusbar.py index 1c3be19c..88e2ad0d 100644 --- a/mitmproxy/console/statusbar.py +++ b/mitmproxy/console/statusbar.py @@ -137,7 +137,7 @@ class StatusBar(urwid.WidgetWrap): def get_status(self): r = [] - if self.master.setheaders.count(): + if len(self.master.options.setheaders): r.append("[") r.append(("heading_key", "H")) r.append("eaders]") diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py index 65eb515b..eaa368a0 100644 --- a/mitmproxy/dump.py +++ b/mitmproxy/dump.py @@ -58,10 +58,6 @@ class DumpMaster(flow.FlowMaster): "HTTP/2 is disabled. Use --no-http2 to silence this warning.", file=sys.stderr) - if options.setheaders: - for i in options.setheaders: - self.setheaders.add(*i) - if options.server_replay: self.start_server_playback( self._readflow(options.server_replay), diff --git a/mitmproxy/flow/__init__.py b/mitmproxy/flow/__init__.py index 8f08bed7..b2ab74c6 100644 --- a/mitmproxy/flow/__init__.py +++ b/mitmproxy/flow/__init__.py @@ -4,8 +4,7 @@ from mitmproxy.flow import export, modules from mitmproxy.flow.io import FlowWriter, FilteredFlowWriter, FlowReader, read_flows_from_paths from mitmproxy.flow.master import FlowMaster from mitmproxy.flow.modules import ( - AppRegistry, SetHeaders, StreamLargeBodies, ClientPlaybackState, - ServerPlaybackState + AppRegistry, StreamLargeBodies, ClientPlaybackState, ServerPlaybackState ) from mitmproxy.flow.state import State, FlowView from mitmproxy.flow import options @@ -16,7 +15,6 @@ __all__ = [ "export", "modules", "FlowWriter", "FilteredFlowWriter", "FlowReader", "read_flows_from_paths", "FlowMaster", - "AppRegistry", "SetHeaders", "StreamLargeBodies", "ClientPlaybackState", - "ServerPlaybackState", "State", "FlowView", - "options", + "AppRegistry", "StreamLargeBodies", "ClientPlaybackState", + "ServerPlaybackState", "State", "FlowView", "options", ] diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py index 244f24f6..06ec18b3 100644 --- a/mitmproxy/flow/master.py +++ b/mitmproxy/flow/master.py @@ -36,7 +36,6 @@ class FlowMaster(controller.Master): self.stream_large_bodies = None # type: Optional[modules.StreamLargeBodies] self.refresh_server_playback = False - self.setheaders = modules.SetHeaders() self.replay_ignore_params = False self.replay_ignore_content = None self.replay_ignore_host = False @@ -329,8 +328,6 @@ class FlowMaster(controller.Master): if f not in self.state.flows: # don't add again on replay self.state.add_flow(f) if not f.reply.acked: - self.setheaders.run(f) - if not f.reply.acked: self.process_new_request(f) return f @@ -348,8 +345,6 @@ class FlowMaster(controller.Master): def response(self, f): self.state.update_flow(f) if not f.reply.acked: - self.setheaders.run(f) - if not f.reply.acked: if self.client_playback: self.client_playback.clear(f) return f diff --git a/mitmproxy/flow/modules.py b/mitmproxy/flow/modules.py index 0ad392d2..fb3c52da 100644 --- a/mitmproxy/flow/modules.py +++ b/mitmproxy/flow/modules.py @@ -5,7 +5,6 @@ import hashlib from six.moves import urllib from mitmproxy import controller -from mitmproxy import filt from netlib import wsgi from netlib import version from netlib import strutils @@ -39,60 +38,6 @@ class AppRegistry: return self.apps.get((host, request.port), None) -class SetHeaders: - def __init__(self): - self.lst = [] - - def set(self, r): - self.clear() - for i in r: - self.add(*i) - - def add(self, fpatt, header, value): - """ - Add a set header hook. - - fpatt: String specifying a filter pattern. - header: Header name. - value: Header value string - - Returns True if hook was added, False if the pattern could not be - parsed. - """ - cpatt = filt.parse(fpatt) - if not cpatt: - return False - self.lst.append((fpatt, header, value, cpatt)) - return True - - def get_specs(self): - """ - Retrieve the hook specifcations. Returns a list of (fpatt, rex, s) - tuples. - """ - return [i[:3] for i in self.lst] - - def count(self): - return len(self.lst) - - def clear(self): - self.lst = [] - - def run(self, f): - for _, header, value, cpatt in self.lst: - if cpatt(f): - if f.response: - f.response.headers.pop(header, None) - else: - f.request.headers.pop(header, None) - for _, header, value, cpatt in self.lst: - if cpatt(f): - if f.response: - f.response.headers.add(header, value) - else: - f.request.headers.add(header, value) - - class StreamLargeBodies(object): def __init__(self, max_size): self.max_size = max_size diff --git a/test/mitmproxy/builtins/test_setheaders.py b/test/mitmproxy/builtins/test_setheaders.py new file mode 100644 index 00000000..1a8d048c --- /dev/null +++ b/test/mitmproxy/builtins/test_setheaders.py @@ -0,0 +1,64 @@ +from .. import tutils, mastertest + +from mitmproxy.builtins import setheaders +from mitmproxy.flow import state +from mitmproxy.flow import options + + +class TestSetHeaders(mastertest.MasterTest): + def mkmaster(self, **opts): + s = state.State() + m = mastertest.RecordingMaster(options.Options(**opts), None, s) + sh = setheaders.SetHeaders() + m.addons.add(sh) + return m, sh + + def test_configure(self): + sh = setheaders.SetHeaders() + tutils.raises( + "invalid setheader filter pattern", + sh.configure, + options.Options( + setheaders = [("~b", "one", "two")] + ) + ) + + def test_setheaders(self): + m, sh = self.mkmaster( + setheaders = [ + ("~q", "one", "two"), + ("~s", "one", "three") + ] + ) + f = tutils.tflow() + f.request.headers["one"] = "xxx" + self.invoke(m, "request", f) + assert f.request.headers["one"] == "two" + + f = tutils.tflow(resp=True) + f.response.headers["one"] = "xxx" + self.invoke(m, "response", f) + assert f.response.headers["one"] == "three" + + m, sh = self.mkmaster( + setheaders = [ + ("~s", "one", "two"), + ("~s", "one", "three") + ] + ) + f = tutils.tflow(resp=True) + f.request.headers["one"] = "xxx" + f.response.headers["one"] = "xxx" + self.invoke(m, "response", f) + assert f.response.headers.get_all("one") == ["two", "three"] + + m, sh = self.mkmaster( + setheaders = [ + ("~q", "one", "two"), + ("~q", "one", "three") + ] + ) + f = tutils.tflow() + f.request.headers["one"] = "xxx" + self.invoke(m, "request", f) + assert f.request.headers.get_all("one") == ["two", "three"] diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py index 10163401..03935c16 100644 --- a/test/mitmproxy/test_flow.py +++ b/test/mitmproxy/test_flow.py @@ -961,55 +961,3 @@ class TestClientConnection: assert c3.get_state() == c.get_state() assert str(c) - - -def test_setheaders(): - h = flow.SetHeaders() - h.add("~q", "foo", "bar") - assert h.lst - - h.set( - [ - (".*", "one", "two"), - (".*", "three", "four"), - ] - ) - assert h.count() == 2 - - h.clear() - assert not h.lst - - h.add("~q", "foo", "bar") - h.add("~s", "foo", "bar") - - v = h.get_specs() - assert v == [('~q', 'foo', 'bar'), ('~s', 'foo', 'bar')] - assert h.count() == 2 - h.clear() - assert h.count() == 0 - - f = tutils.tflow() - f.request.content = b"foo" - h.add("~s", "foo", "bar") - h.run(f) - assert f.request.content == b"foo" - - h.clear() - h.add("~s", "one", "two") - h.add("~s", "one", "three") - f = tutils.tflow(resp=True) - f.request.headers["one"] = "xxx" - f.response.headers["one"] = "xxx" - h.run(f) - assert f.request.headers["one"] == "xxx" - assert f.response.headers.get_all("one") == ["two", "three"] - - h.clear() - h.add("~q", "one", "two") - h.add("~q", "one", "three") - f = tutils.tflow() - f.request.headers["one"] = "xxx" - h.run(f) - assert f.request.headers.get_all("one") == ["two", "three"] - - assert not h.add("~", "foo", "bar") |