Extract flow-specific Master operations into FlowMaster.

3 files changed, 49 insertions, 20 deletions

diff --git a/libmproxy/console.py b/libmproxy/console.py
index dbeef04a..794d5483 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -779,7 +779,7 @@ VIEW_CONNLIST = 0
 VIEW_FLOW = 1
 VIEW_HELP = 2
 
-class ConsoleMaster(controller.Master):
+class ConsoleMaster(flow.FlowMaster):
     palette = []
     footer_text_default = [
         ('key', "?"), ":help ",
@@ -794,10 +794,10 @@ class ConsoleMaster(controller.Master):
         ('key', "q"), ":back ",
     ]
     def __init__(self, server, options):
+        flow.FlowMaster.__init__(self, server, ConsoleState())
+
         self.conn_list_view = None
         self.set_palette()
-        controller.Master.__init__(self, server)
-        self.state = ConsoleState()
 
         r = self.set_limit(options.limit)
         if r:
@@ -1313,23 +1313,18 @@ class ConsoleMaster(controller.Master):
 
     # Handlers
     def handle_clientconnection(self, r):
-        f = flow.Flow(r)
-        self.state.add_browserconnect(f)
-        r.ack()
-        self.sync_list_view()
+        f = flow.FlowMaster.handle_clientconnection(self, r)
+        if f:
+            self.sync_list_view()
 
     def handle_error(self, r):
-        f = self.state.add_error(r)
-        if not f:
-            r.ack()
-        else:
+        f = flow.FlowMaster.handle_error(self, r)
+        if f:
             self.process_flow(f, r)
 
     def handle_request(self, r):
-        f = self.state.add_request(r)
-        if not f:
-            r.ack()
-        else:
+        f = flow.FlowMaster.handle_request(self, r)
+        if f:
             if f.match(self.stickycookie):
                 hid = (f.request.host, f.request.port)
                 if f.request.headers.has_key("cookie"):
@@ -1339,10 +1334,8 @@ class ConsoleMaster(controller.Master):
             self.process_flow(f, r)
 
     def handle_response(self, r):
-        f = self.state.add_response(r)
-        if not f:
-            r.ack()
-        else:
+        f = flow.FlowMaster.handle_request(self, r)
+        if f:
             if f.match(self.stickycookie):
                 hid = (f.request.host, f.request.port)
                 if f.response.headers.has_key("set-cookie"):
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 921abe1d..e791055f 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -4,6 +4,7 @@
 """
 import subprocess, base64, sys, json
 import proxy, threading, netstring
+import controller
 
 class RunException(Exception):
     def __init__(self, msg, returncode, errout):
@@ -286,6 +287,41 @@ class State:
         #end nocover
 
 
+class FlowMaster(controller.Master):
+    def __init__(self, server, state):
+        controller.Master.__init__(self, server)
+        self.state = state
+
+    # Handlers
+    def handle_clientconnection(self, r):
+        f = Flow(r)
+        self.state.add_browserconnect(f)
+        r.ack()
+        return f
+
+    def handle_error(self, r):
+        f = self.state.add_error(r)
+        if not f:
+            r.ack()
+        return f
+
+    def handle_request(self, r):
+        f = self.state.add_request(r)
+        if not f:
+            r.ack()
+        return f
+
+    def handle_response(self, r):
+        f = self.state.add_response(r)
+        if not f:
+            r.ack()
+        else:
+            if f.match(self.stickycookie):
+                hid = (f.request.host, f.request.port)
+                if f.response.headers.has_key("set-cookie"):
+                    self.stickyhosts[hid] = f.response.headers["set-cookie"]
+            self.process_flow(f, r)
+
 
 class FlowWriter:
     def __init__(self, fo):
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index 4a6d3614..6093c936 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -307,7 +307,7 @@ class ClientConnection(controller.Msg):
         controller.Msg.__init__(self)
 
     def get_state(self):
-        return list(self.address)
+        return list(self.address) if self.address else None
 
     @classmethod
     def from_state(klass, state):