Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-27 02:24:18 +00:00
Code Issues Projects Releases Wiki Activity

Merge pull request #1372 from cortesi/setheaders

Browse Source 
setheaders -> addon
This commit is contained in:
Aldo Cortesi 2016-07-17 18:49:41 +12:00 committed by GitHub
commit 6c86d7bd4f
11 changed files with 114 additions and 133 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mitmproxy/builtins/__init__.py
View File

@ -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(),
]

39
mitmproxy/builtins/setheaders.py Normal file
View File

@ -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)

4
mitmproxy/console/master.py
View File

@ -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)

12
mitmproxy/console/options.py
View File

@ -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,

2
mitmproxy/console/statusbar.py
View File

@ -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]")

4
mitmproxy/dump.py
View File

@ -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),

8
mitmproxy/flow/__init__.py
View File

@ -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",
]

5
mitmproxy/flow/master.py
View File

@ -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
@ -328,8 +327,6 @@ class FlowMaster(controller.Master):
return
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
@ -347,8 +344,6 @@ class FlowMaster(controller.Master):
@controller.handler
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)

55
mitmproxy/flow/modules.py
View File

@ -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

64
test/mitmproxy/builtins/test_setheaders.py Normal file
View File

@ -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"]

52
test/mitmproxy/test_flow.py
View File

@ -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")