mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-23 00:01:36 +00:00
First pass of playback function for mitmdump.
This commit is contained in:
parent
deb79a9c5a
commit
fd4dd8cb6b
@ -10,6 +10,7 @@ class Options(object):
|
||||
"wfile",
|
||||
"request_script",
|
||||
"response_script",
|
||||
"replay",
|
||||
]
|
||||
def __init__(self, **kwargs):
|
||||
for k, v in kwargs.items():
|
||||
@ -38,6 +39,15 @@ class DumpMaster(flow.FlowMaster):
|
||||
except IOError, v:
|
||||
raise DumpError(v.strerror)
|
||||
|
||||
if options.replay:
|
||||
path = os.path.expanduser(options.replay)
|
||||
try:
|
||||
f = file(path, "r")
|
||||
flows = list(flow.FlowReader(f).stream())
|
||||
except IOError, v:
|
||||
raise DumpError(v.strerror)
|
||||
self.start_playback(flows)
|
||||
|
||||
def _runscript(self, f, script):
|
||||
try:
|
||||
ret = f.run_script(script)
|
||||
@ -56,6 +66,7 @@ class DumpMaster(flow.FlowMaster):
|
||||
f = flow.FlowMaster.handle_request(self, r)
|
||||
if self.o.request_script:
|
||||
self._runscript(f, self.o.request_script)
|
||||
if not self.playback(f):
|
||||
r.ack()
|
||||
|
||||
def indent(self, n, t):
|
||||
|
@ -328,6 +328,27 @@ class FlowMaster(controller.Master):
|
||||
def __init__(self, server, state):
|
||||
controller.Master.__init__(self, server)
|
||||
self.state = state
|
||||
self._playback_state = None
|
||||
|
||||
def start_playback(self, flows):
|
||||
self._playback_state = ServerPlaybackState()
|
||||
self._playback_state.load(flows)
|
||||
|
||||
def playback(self, flow):
|
||||
"""
|
||||
This method should be called by child classes in the handle_request
|
||||
handler. Returns True if playback has taken place, None if not.
|
||||
"""
|
||||
if self._playback_state:
|
||||
rflow = self._playback_state.next_flow(flow)
|
||||
if not rflow:
|
||||
return None
|
||||
response = proxy.Response.from_state(flow.request, rflow.response.get_state())
|
||||
response.set_replay()
|
||||
flow.response = response
|
||||
flow.request.ack(response)
|
||||
return True
|
||||
return None
|
||||
|
||||
def handle_clientconnect(self, r):
|
||||
self.state.clientconnect(r)
|
||||
|
@ -265,6 +265,13 @@ class Response(controller.Msg):
|
||||
self.timestamp = timestamp or time.time()
|
||||
self.cached = False
|
||||
controller.Msg.__init__(self)
|
||||
self.replay = False
|
||||
|
||||
def set_replay(self):
|
||||
self.replay = True
|
||||
|
||||
def is_replay(self):
|
||||
return self.replay
|
||||
|
||||
def load_state(self, state):
|
||||
self.code = state["code"]
|
||||
@ -308,7 +315,10 @@ class Response(controller.Msg):
|
||||
return self.cached
|
||||
|
||||
def short(self):
|
||||
return "%s %s"%(self.code, self.msg)
|
||||
r = "%s %s"%(self.code, self.msg)
|
||||
if self.is_replay():
|
||||
r = "[replay] " + r
|
||||
return r
|
||||
|
||||
def assemble(self):
|
||||
"""
|
||||
|
4
mitmdump
4
mitmdump
@ -47,6 +47,9 @@ if __name__ == '__main__':
|
||||
parser.add_option("", "--respscript",
|
||||
action="store", dest="response_script", default=None,
|
||||
help="Script to run when a response is recieved.")
|
||||
parser.add_option("-r", "--replay",
|
||||
action="store", dest="replay", default=None,
|
||||
help="Replay server responses from a saved file.")
|
||||
|
||||
options, args = parser.parse_args()
|
||||
|
||||
@ -61,6 +64,7 @@ if __name__ == '__main__':
|
||||
wfile = options.wfile,
|
||||
request_script = options.request_script,
|
||||
response_script = options.response_script,
|
||||
replay = options.replay,
|
||||
)
|
||||
if args:
|
||||
filt = " ".join(args)
|
||||
|
@ -23,6 +23,23 @@ class uDumpMaster(libpry.AutoTree):
|
||||
self._cycle(m, content)
|
||||
return cs.getvalue()
|
||||
|
||||
def test_replay(self):
|
||||
cs = StringIO()
|
||||
|
||||
o = dump.Options(replay="nonexistent")
|
||||
libpry.raises(dump.DumpError, dump.DumpMaster, None, o, None, outfile=cs)
|
||||
|
||||
t = self.tmpdir()
|
||||
p = os.path.join(t, "rep")
|
||||
f = open(p, "w")
|
||||
fw = flow.FlowWriter(f)
|
||||
t = utils.tflow()
|
||||
fw.add(t)
|
||||
f.close()
|
||||
|
||||
o = dump.Options(replay=p)
|
||||
m = dump.DumpMaster(None, o, None, outfile=cs)
|
||||
|
||||
def test_options(self):
|
||||
o = dump.Options(verbosity = 2)
|
||||
assert o.verbosity == 2
|
||||
|
@ -309,6 +309,24 @@ class uFlowMaster(libpry.AutoTree):
|
||||
err = proxy.Error(f.request, "msg")
|
||||
fm.handle_error(err)
|
||||
|
||||
def test_replay(self):
|
||||
s = flow.State()
|
||||
|
||||
f = utils.tflow()
|
||||
f.response = utils.tresp(f.request)
|
||||
pb = [f]
|
||||
|
||||
fm = flow.FlowMaster(None, s)
|
||||
assert not fm.playback(utils.tflow())
|
||||
|
||||
fm.start_playback(pb)
|
||||
assert fm.playback(utils.tflow())
|
||||
|
||||
fm.start_playback(pb)
|
||||
r = utils.tflow()
|
||||
r.request.content = "gibble"
|
||||
assert not fm.playback(r)
|
||||
|
||||
|
||||
|
||||
tests = [
|
||||
|
@ -1,3 +1,4 @@
|
||||
import os.path
|
||||
from libmproxy import proxy, utils, filt, flow
|
||||
|
||||
def treq(conn=None):
|
||||
@ -20,3 +21,4 @@ def tflow():
|
||||
r = treq()
|
||||
return flow.Flow(r)
|
||||
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user