Move script hooks into the flow primitives.

This lets handle scripts in corner cases like playback situations more easily.

libmproxy/dump.py

@@ -32,6 +32,11 @@ class DumpMaster(flow.FlowMaster):
         else:
             self.filt = None
 
+        if self.o.response_script:
+            self.set_response_script(self.o.response_script)  
+        if self.o.request_script:
+            self.set_request_script(self.o.request_script)  
+
         if options.wfile:
             path = os.path.expanduser(options.wfile)
             try:
@@ -47,7 +52,7 @@ class DumpMaster(flow.FlowMaster):
                 flows = list(flow.FlowReader(f).stream())
             except IOError, v:
                 raise DumpError(v.strerror)
-            self.start_playback(flows)
+            self.start_playback(flows, options.kill)
 
     def _runscript(self, f, script):
         try:
@@ -65,17 +70,7 @@ class DumpMaster(flow.FlowMaster):
 
     def handle_request(self, r):
         f = flow.FlowMaster.handle_request(self, r)
-        if self.o.request_script:
-            self._runscript(f, self.o.request_script)
-
-        if self.o.replay:
-            pb = self.playback(f)
-            if not pb:
-                if self.o.kill:
-                    self.state.kill_flow(f)
-                else:
-                    r.ack()
-        else:
+        if f:
             r.ack()
 
     def indent(self, n, t):
@@ -85,8 +80,6 @@ class DumpMaster(flow.FlowMaster):
     def handle_response(self, msg):
         f = flow.FlowMaster.handle_response(self, msg)
         if f:
-            if self.o.response_script:
-                self._runscript(f, self.o.response_script)
             msg.ack()
             if self.filt and not f.match(self.filt):
                     return

libmproxy/flow.py

@@ -35,7 +35,7 @@ class ServerPlaybackState:
     def __init__(self):
         self.fmap = {}
 
-    def __len__(self):
+    def count(self):
         return sum([len(i) for i in self.fmap.values()])
     
     def load(self, flows):
@@ -329,19 +329,35 @@ class FlowMaster(controller.Master):
     def __init__(self, server, state):
         controller.Master.__init__(self, server)
         self.state = state
-        self._playback_state = None
+        self.playback = None
+        self.scripts = {}
+        self.kill_nonreplay = False
 
-    def start_playback(self, flows):
-        self._playback_state = ServerPlaybackState()
-        self._playback_state.load(flows)
+    def _runscript(self, f, script):
+        return f.run_script(script)
 
-    def playback(self, flow):
+    def set_response_script(self, s):
+        self.scripts["response"] = s
+
+    def set_request_script(self, s):
+        self.scripts["request"] = s
+
+    def start_playback(self, flows, kill):
+        """
+            flows: A list of flows.
+            kill: Boolean, should we kill requests not part of the replay?
+        """
+        self.playback = ServerPlaybackState()
+        self.playback.load(flows)
+        self.kill_nonreplay = kill
+
+    def do_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 self.playback:
+            rflow = self.playback.next_flow(flow)
             if not rflow:
                 return None
             response = proxy.Response.from_state(flow.request, rflow.response.get_state())
@@ -365,12 +381,24 @@ class FlowMaster(controller.Master):
         return f
 
     def handle_request(self, r):
-        return self.state.add_request(r)
+        f = self.state.add_request(r)
+        if "request" in self.scripts:
+            self._runscript(f, self.scripts["request"])
+        if self.playback:
+            pb = self.do_playback(f)
+            if not pb:
+                if self.kill_nonreplay:
+                    self.state.kill_flow(f)
+                else:
+                    r.ack()
+        return f
 
     def handle_response(self, r):
         f = self.state.add_response(r)
         if not f:
             r.ack()
+        if "response" in self.scripts:
+            self._runscript(f, self.scripts["response"])
         return f
 
 

test/test_flow.py

@@ -26,16 +26,16 @@ class uServerPlaybackState(libpry.AutoTree):
         r2.request.headers["key"] = ["two"]
 
         s.load([r, r2])
-        assert len(s) == 2
+        assert s.count() == 2
         assert len(s.fmap.keys()) == 1
 
         n = s.next_flow(r)
         assert n.request.headers["key"] == ["one"]
-        assert len(s) == 1
+        assert s.count() == 1
 
         n = s.next_flow(r)
         assert n.request.headers["key"] == ["two"]
-        assert len(s) == 0
+        assert s.count() == 0
 
         assert not s.next_flow(r)
 
@@ -317,15 +317,15 @@ class uFlowMaster(libpry.AutoTree):
         pb = [f]
 
         fm = flow.FlowMaster(None, s)
-        assert not fm.playback(utils.tflow())
+        assert not fm.do_playback(utils.tflow())
 
-        fm.start_playback(pb)
-        assert fm.playback(utils.tflow())
+        fm.start_playback(pb, False)
+        assert fm.do_playback(utils.tflow())
 
-        fm.start_playback(pb)
+        fm.start_playback(pb, False)
         r = utils.tflow()
         r.request.content = "gibble"
-        assert not fm.playback(r)
+        assert not fm.do_playback(r)