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",
|
"wfile",
|
||||||
"request_script",
|
"request_script",
|
||||||
"response_script",
|
"response_script",
|
||||||
|
"replay",
|
||||||
]
|
]
|
||||||
def __init__(self, **kwargs):
|
def __init__(self, **kwargs):
|
||||||
for k, v in kwargs.items():
|
for k, v in kwargs.items():
|
||||||
@ -38,6 +39,15 @@ class DumpMaster(flow.FlowMaster):
|
|||||||
except IOError, v:
|
except IOError, v:
|
||||||
raise DumpError(v.strerror)
|
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):
|
def _runscript(self, f, script):
|
||||||
try:
|
try:
|
||||||
ret = f.run_script(script)
|
ret = f.run_script(script)
|
||||||
@ -56,6 +66,7 @@ class DumpMaster(flow.FlowMaster):
|
|||||||
f = flow.FlowMaster.handle_request(self, r)
|
f = flow.FlowMaster.handle_request(self, r)
|
||||||
if self.o.request_script:
|
if self.o.request_script:
|
||||||
self._runscript(f, self.o.request_script)
|
self._runscript(f, self.o.request_script)
|
||||||
|
if not self.playback(f):
|
||||||
r.ack()
|
r.ack()
|
||||||
|
|
||||||
def indent(self, n, t):
|
def indent(self, n, t):
|
||||||
|
@ -328,6 +328,27 @@ class FlowMaster(controller.Master):
|
|||||||
def __init__(self, server, state):
|
def __init__(self, server, state):
|
||||||
controller.Master.__init__(self, server)
|
controller.Master.__init__(self, server)
|
||||||
self.state = state
|
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):
|
def handle_clientconnect(self, r):
|
||||||
self.state.clientconnect(r)
|
self.state.clientconnect(r)
|
||||||
|
@ -265,6 +265,13 @@ class Response(controller.Msg):
|
|||||||
self.timestamp = timestamp or time.time()
|
self.timestamp = timestamp or time.time()
|
||||||
self.cached = False
|
self.cached = False
|
||||||
controller.Msg.__init__(self)
|
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):
|
def load_state(self, state):
|
||||||
self.code = state["code"]
|
self.code = state["code"]
|
||||||
@ -308,7 +315,10 @@ class Response(controller.Msg):
|
|||||||
return self.cached
|
return self.cached
|
||||||
|
|
||||||
def short(self):
|
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):
|
def assemble(self):
|
||||||
"""
|
"""
|
||||||
|
4
mitmdump
4
mitmdump
@ -47,6 +47,9 @@ if __name__ == '__main__':
|
|||||||
parser.add_option("", "--respscript",
|
parser.add_option("", "--respscript",
|
||||||
action="store", dest="response_script", default=None,
|
action="store", dest="response_script", default=None,
|
||||||
help="Script to run when a response is recieved.")
|
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()
|
options, args = parser.parse_args()
|
||||||
|
|
||||||
@ -61,6 +64,7 @@ if __name__ == '__main__':
|
|||||||
wfile = options.wfile,
|
wfile = options.wfile,
|
||||||
request_script = options.request_script,
|
request_script = options.request_script,
|
||||||
response_script = options.response_script,
|
response_script = options.response_script,
|
||||||
|
replay = options.replay,
|
||||||
)
|
)
|
||||||
if args:
|
if args:
|
||||||
filt = " ".join(args)
|
filt = " ".join(args)
|
||||||
|
@ -23,6 +23,23 @@ class uDumpMaster(libpry.AutoTree):
|
|||||||
self._cycle(m, content)
|
self._cycle(m, content)
|
||||||
return cs.getvalue()
|
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):
|
def test_options(self):
|
||||||
o = dump.Options(verbosity = 2)
|
o = dump.Options(verbosity = 2)
|
||||||
assert o.verbosity == 2
|
assert o.verbosity == 2
|
||||||
|
@ -309,6 +309,24 @@ class uFlowMaster(libpry.AutoTree):
|
|||||||
err = proxy.Error(f.request, "msg")
|
err = proxy.Error(f.request, "msg")
|
||||||
fm.handle_error(err)
|
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 = [
|
tests = [
|
||||||
|
@ -1,3 +1,4 @@
|
|||||||
|
import os.path
|
||||||
from libmproxy import proxy, utils, filt, flow
|
from libmproxy import proxy, utils, filt, flow
|
||||||
|
|
||||||
def treq(conn=None):
|
def treq(conn=None):
|
||||||
@ -20,3 +21,4 @@ def tflow():
|
|||||||
r = treq()
|
r = treq()
|
||||||
return flow.Flow(r)
|
return flow.Flow(r)
|
||||||
|
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user