Add a "done" event for scripts.

Called exactly once after all other events.

examples/stub.py

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

libmproxy/console.py

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

libmproxy/dump.py

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

libmproxy/flow.py

@@ -1112,12 +1112,18 @@ 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.
         """
+        if path is None:
+            self.run_script_hook("done")
+            self.script = None
+        else:
             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):
@@ -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):