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

Add a "done" event for scripts.

Browse Source 
Called exactly once after all other events.
This commit is contained in:
Aldo Cortesi 2011-08-05 14:03:10 +12:00
parent cd0e2f18e6
commit f23818ceea
4 changed files with 41 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
examples/stub.py
View File

@ -1,30 +1,31 @@
"""
This is a script stub, with empty definitions for all events.
This is a script stub, with definitions for all events.
"""
def start(ctx):
"""
Called once on script startup, before any other events.
"""
pass
ctx.log("start")
def clientconnect(ctx, client_connect):
"""
Called when a client initiates a connection to the proxy. Note that a
connection can correspond to multiple HTTP requests
"""
pass
ctx.log("clientconnect")
def request(ctx, flow):
"""
Called when a client request has been received.
"""
ctx.log("request")
def response(ctx, flow):
"""
Called when a server response has been received.
"""
pass
ctx.log("response")
def error(ctx, flow):
"""
@ -32,16 +33,16 @@ def error(ctx, flow):
interrupted connections. This is distinct from a valid server HTTP error
response, which is simply a response with an HTTP error code.
"""
pass
ctx.log("error")
def clientdisconnect(ctx, client_disconnect):
"""
Called when a client disconnects from the proxy.
"""
pass
ctx.log("clientdisconnect")
def done(ctx):
"""
Called once on script shutdown, after any other events.
"""
pass
ctx.log("done")

16
libmproxy/console.py
View File

@ -964,12 +964,6 @@ class ConsoleMaster(flow.FlowMaster):
print >> sys.stderr, "Sticky auth error:", r
sys.exit(1)
if options.script:
err = self.load_script(options.script)
if err:
print >> sys.stderr, "Script load error:", r
sys.exit(1)
self.refresh_server_playback = options.refresh_server_playback
self.anticache = options.anticache
self.anticomp = options.anticomp
@ -987,6 +981,13 @@ class ConsoleMaster(flow.FlowMaster):
self.debug = options.debug
if options.script:
err = self.load_script(options.script)
if err:
print >> sys.stderr, "Script load error:", err
sys.exit(1)
def run_script_once(self, path, f):
ret = self.get_script(path)
if ret[0]:
@ -999,6 +1000,7 @@ class ConsoleMaster(flow.FlowMaster):
s.run("response", f)
if f.error:
s.run("error", f)
s.run("done")
self.refresh_connection(f)
self.state.last_script = path
@ -1660,7 +1662,7 @@ class ConsoleMaster(flow.FlowMaster):
k = None
elif k == "s":
if self.script:
self.script = None
self.load_script(None)
else:
self.path_prompt(
"Set script: ",

3
libmproxy/dump.py
View File

@ -191,10 +191,11 @@ class DumpMaster(flow.FlowMaster):
self._process_flow(f)
return f
# begin nocover
def run(self):
if self.o.rfile and not self.o.keepserving:
if self.script:
self.load_script(None)
return
try:
return flow.FlowMaster.run(self)

33
libmproxy/flow.py
View File

@ -1112,13 +1112,19 @@ class FlowMaster(controller.Master):
def load_script(self, path):
"""
Loads a script. Returns an error description if something went
wrong.
wrong. If path is None, the current script is terminated.
"""
r = self.get_script(path)
if r[0]:
return r[0]
if path is None:
self.run_script_hook("done")
self.script = None
else:
self.script = r[1]
r = self.get_script(path)
if r[0]:
return r[0]
else:
if self.script:
self.run_script_hook("done")
self.script = r[1]
def set_stickycookie(self, txt):
if txt:
@ -1250,7 +1256,7 @@ class FlowMaster(controller.Master):
rt.start()
#end nocover
def run_script(self, name, *args, **kwargs):
def run_script_hook(self, name, *args, **kwargs):
if self.script:
ret = self.script.run(name, *args, **kwargs)
if not ret[0] and ret[1]:
@ -1258,12 +1264,12 @@ class FlowMaster(controller.Master):
self.add_event(e, "error")
def handle_clientconnect(self, cc):
self.run_script("clientconnect", cc)
self.run_script_hook("clientconnect", cc)
self.add_event("Connect from: %s:%s"%cc.address)
cc._ack()
def handle_clientdisconnect(self, r):
self.run_script("clientdisconnect", r)
self.run_script_hook("clientdisconnect", r)
s = "Disconnect from: %s:%s"%r.client_conn.address
self.add_event(s)
if r.client_conn.requestcount:
@ -1278,7 +1284,7 @@ class FlowMaster(controller.Master):
def handle_error(self, r):
f = self.state.add_error(r)
if f:
self.run_script("error", f)
self.run_script_hook("error", f)
if self.client_playback:
self.client_playback.clear(f)
r._ack()
@ -1286,14 +1292,14 @@ class FlowMaster(controller.Master):
def handle_request(self, r):
f = self.state.add_request(r)
self.run_script("request", f)
self.run_script_hook("request", f)
self.process_new_request(f)
return f
def handle_response(self, r):
f = self.state.add_response(r)
if f:
self.run_script("response", f)
self.run_script_hook("response", f)
if self.client_playback:
self.client_playback.clear(f)
if not f:
@ -1301,6 +1307,11 @@ class FlowMaster(controller.Master):
self.process_new_response(f)
return f
def shutdown(self):
if self.script:
self.load_script(None)
controller.Master.shutdown(self)
class FlowWriter:
def __init__(self, fo):