First pass of playback function for mitmdump.

This commit is contained in:
Aldo Cortesi 2011-02-21 09:54:39 +13:00
parent deb79a9c5a
commit fd4dd8cb6b
7 changed files with 85 additions and 2 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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