Restructure to make subclassing Flow unnecessary.

3 files changed, 99 insertions, 91 deletions

diff --git a/libmproxy/console.py b/libmproxy/console.py
index 4c51e786..12f6dec9 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -44,6 +44,55 @@ def format_keyvals(lst, key="key", val="text", space=5, indent=0):
     return ret
 
 
+def format_flow(flow, focus, padding=3):
+    if not flow.request and not flow.response:
+        txt = [
+            ("title", " Connection from %s..."%(flow.connection.address)),
+        ]
+    else:
+        txt = [
+            ("ack", "!") if flow.intercepting and not flow.request.acked else " ",
+            ("method", flow.request.method),
+            " ",
+            (
+                "text" if (flow.response or flow.error) else "title",
+                flow.request.url(),
+            ),
+        ]
+        if flow.response or flow.error or flow.is_replay():
+            txt.append("\n" + " "*(padding+2))
+            if flow.is_replay():
+                txt.append(("method", "[replay] "))
+            if not (flow.response or flow.error):
+                txt.append(("text", "waiting for response..."))
+
+        if flow.response:
+            txt.append(
+               ("ack", "!") if flow.intercepting and not flow.response.acked else " "
+            )
+            txt.append("<- ")
+            if flow.response.code in [200, 304]:
+                txt.append(("goodcode", str(flow.response.code)))
+            else:
+                txt.append(("error", str(flow.response.code)))
+            t = flow.response.headers.get("content-type")
+            if t:
+                t = t[0].split(";")[0]
+                txt.append(("text", " %s"%t))
+            if flow.response.content:
+                txt.append(", %s"%utils.pretty_size(len(flow.response.content)))
+        elif flow.error:
+            txt.append(
+               ("error", flow.error.msg)
+            )
+    if focus:
+        txt.insert(0, ("focus", ">>" + " "*(padding-2)))
+    else:
+        txt.insert(0, " "*padding)
+    return txt
+
+
+
 #begin nocover
 
 def int_version(v):
@@ -67,8 +116,9 @@ class WWrap(urwid.WidgetWrap):
 
 
 class ConnectionItem(WWrap):
-    def __init__(self, master, state, flow):
+    def __init__(self, master, state, flow, focus):
         self.master, self.state, self.flow = master, state, flow
+        self.focus = focus
         w = self.get_text()
         WWrap.__init__(self, w)
 
@@ -77,7 +127,7 @@ class ConnectionItem(WWrap):
         self.w = self.get_text()
 
     def get_text(self):
-        return urwid.Text(self.flow.get_text())
+        return urwid.Text(format_flow(self.flow, self.focus))
 
     def selectable(self):
         return True
@@ -123,7 +173,7 @@ class ConnectionListView(urwid.ListWalker):
 
     def get_focus(self):
         f, i = self.state.get_focus()
-        f = ConnectionItem(self.master, self.state, f) if f else None
+        f = ConnectionItem(self.master, self.state, f, True) if f else None
         return f, i
 
     def set_focus(self, focus):
@@ -133,23 +183,24 @@ class ConnectionListView(urwid.ListWalker):
 
     def get_next(self, pos):
         f, i = self.state.get_next(pos)
-        f = ConnectionItem(self.master, self.state, f) if f else None
+        f = ConnectionItem(self.master, self.state, f, False) if f else None
         return f, i
 
     def get_prev(self, pos):
         f, i = self.state.get_prev(pos)
-        f = ConnectionItem(self.master, self.state, f) if f else None
+        f = ConnectionItem(self.master, self.state, f, False) if f else None
         return f, i
 
 
 class ConnectionViewHeader(WWrap):
-    def __init__(self, flow):
-        self.flow = flow
-        self.w = urwid.Text(flow.get_text(nofocus=True, padding=0))
+    def __init__(self, master, flow):
+        self.master, self.flow = master, flow
+        self.w = urwid.Text(format_flow(flow, False, padding=0))
 
-    def refresh_connection(self, f):
+    def refresh_connection(self, flow):
         if f == self.flow:
-            self.w = urwid.Text(f.get_text(nofocus=True, padding=0))
+            self.w = urwid.Text(format_flow(flow, False, padding=0))
+
 
 VIEW_BODY_RAW = 0
 VIEW_BODY_BINARY = 1
@@ -602,59 +653,6 @@ class StatusBar(WWrap):
 
 #end nocover
 
-class ConsoleFlow(flow.Flow):
-    def __init__(self, connection):
-        flow.Flow.__init__(self, connection)
-        self.focus = False
-
-    def get_text(self, nofocus=False, padding=3):
-        if not self.request and not self.response:
-            txt = [
-                ("title", " Connection from %s..."%(self.connection.address)),
-            ]
-        else:
-            txt = [
-                ("ack", "!") if self.intercepting and not self.request.acked else " ",
-                ("method", self.request.method),
-                " ",
-                (
-                    "text" if (self.response or self.error) else "title",
-                    self.request.url(),
-                ),
-            ]
-            if self.response or self.error or self.is_replay():
-                txt.append("\n" + " "*(padding+2))
-                if self.is_replay():
-                    txt.append(("method", "[replay] "))
-                if not (self.response or self.error):
-                    txt.append(("text", "waiting for response..."))
-
-            if self.response:
-                txt.append(
-                   ("ack", "!") if self.intercepting and not self.response.acked else " "
-                )
-                txt.append("<- ")
-                if self.response.code in [200, 304]:
-                    txt.append(("goodcode", str(self.response.code)))
-                else:
-                    txt.append(("error", str(self.response.code)))
-                t = self.response.headers.get("content-type")
-                if t:
-                    t = t[0].split(";")[0]
-                    txt.append(("text", " %s"%t))
-                if self.response.content:
-                    txt.append(", %s"%utils.pretty_size(len(self.response.content)))
-            elif self.error:
-                txt.append(
-                   ("error", self.error.msg)
-                )
-        if self.focus and not nofocus:
-            txt.insert(0, ("focus", ">>" + " "*(padding-2)))
-        else:
-            txt.insert(0, " "*padding)
-        return txt
-
-
 class ConsoleState(flow.State):
     def __init__(self):
         flow.State.__init__(self)
@@ -694,13 +692,10 @@ class ConsoleState(flow.State):
 
     def set_focus(self, idx):
         if self.view:
-            for i in self.view:
-                i.focus = False
             if idx >= len(self.view):
                 idx = len(self.view) - 1
             elif idx < 0:
                 idx = 0
-            self.view[idx].focus = True
             self.focus = idx
 
     def get_from_pos(self, pos):
@@ -817,7 +812,6 @@ class ConsoleMaster(controller.Master):
         sys.stderr.flush()
         self.shutdown()
 
-
     def make_view(self):
         self.view = urwid.Frame(
                         self.body,
@@ -850,7 +844,7 @@ class ConsoleMaster(controller.Master):
     def view_flow(self, flow):
         self.statusbar = StatusBar(self, self.footer_text_connview)
         self.body = ConnectionView(self, self.state, flow)
-        self.header = ConnectionViewHeader(flow)
+        self.header = ConnectionViewHeader(self, flow)
         self.viewstate = VIEW_FLOW
         self.currentflow = flow
         self.make_view()
@@ -903,7 +897,7 @@ class ConsoleMaster(controller.Master):
             f.close()
         except IOError, v:
             return v.strerror
-        self.state.load_flows(data, ConsoleFlow)
+        self.state.load_flows(data)
         if self.conn_list_view:
             self.conn_list_view.set_focus(0)
             self.sync_list_view()
@@ -1209,7 +1203,7 @@ class ConsoleMaster(controller.Master):
 
     # Handlers
     def handle_browserconnection(self, r):
-        f = ConsoleFlow(r)
+        f = flow.Flow(r)
         self.state.add_browserconnect(f)
         r.ack()
         self.sync_list_view()
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 7f4394b1..a014f8cb 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -33,6 +33,12 @@ class Flow:
         self.intercepting = False
         self._backup = None
 
+    def run_script(self):
+        """
+            Run a script on a flow, returning the modified flow.
+        """
+        pass
+
     def dump(self):
         data = dict(
                 flows = [self.get_state()]
@@ -160,9 +166,9 @@ class State:
                )
         return bson.dumps(data)
 
-    def load_flows(self, js, klass):
+    def load_flows(self, js):
         data = bson.loads(js)
-        data = [klass.from_state(i) for i in data["flows"]]
+        data = [Flow.from_state(i) for i in data["flows"]]
         self.flow_list.extend(data)
 
     def set_limit(self, limit):
diff --git a/test/test_console.py b/test/test_console.py
index cfafed96..c2791bb7 100644
--- a/test/test_console.py
+++ b/test/test_console.py
@@ -19,14 +19,14 @@ def tresp(req=None):
 
 def tflow():
     bc = proxy.BrowserConnection("address", 22)
-    return console.ConsoleFlow(bc)
+    return flow.Flow(bc)
 
 
 class uState(libpry.AutoTree):
     def test_backup(self):
         bc = proxy.BrowserConnection("address", 22)
         c = console.ConsoleState()
-        f = console.ConsoleFlow(bc)
+        f = flow.Flow(bc)
         c.add_browserconnect(f)
 
         f.backup()
@@ -40,7 +40,7 @@ class uState(libpry.AutoTree):
         """
         bc = proxy.BrowserConnection("address", 22)
         c = console.ConsoleState()
-        f = console.ConsoleFlow(bc)
+        f = flow.Flow(bc)
         c.add_browserconnect(f)
         assert c.lookup(bc)
         assert c.get_focus() == (f, 0)
@@ -66,7 +66,7 @@ class uState(libpry.AutoTree):
     def test_err(self):
         bc = proxy.BrowserConnection("address", 22)
         c = console.ConsoleState()
-        f = console.ConsoleFlow(bc)
+        f = flow.Flow(bc)
         c.add_browserconnect(f)
         e = proxy.Error(bc, "message")
         assert c.add_error(e)
@@ -104,7 +104,7 @@ class uState(libpry.AutoTree):
         c = console.ConsoleState()
 
         bc = proxy.BrowserConnection("address", 22)
-        f = console.ConsoleFlow(bc)
+        f = flow.Flow(bc)
         c.add_browserconnect(f)
         assert c.get_focus() == (f, 0)
         assert c.get_from_pos(0) == (f, 0)
@@ -112,7 +112,7 @@ class uState(libpry.AutoTree):
         assert c.get_next(0) == (None, None)
 
         bc2 = proxy.BrowserConnection("address", 22)
-        f2 = console.ConsoleFlow(bc2)
+        f2 = flow.Flow(bc2)
         c.add_browserconnect(f2)
         assert c.get_focus() == (f, 1)
         assert c.get_next(0) == (f, 1)
@@ -193,8 +193,8 @@ class uState(libpry.AutoTree):
 
         dump = c.dump_flows()
         c.clear()
-        c.load_flows(dump, console.ConsoleFlow)
-        assert isinstance(c.flow_list[0], console.ConsoleFlow)
+        c.load_flows(dump)
+        assert isinstance(c.flow_list[0], flow.Flow)
 
 
 class uFlow(libpry.AutoTree):
@@ -212,36 +212,44 @@ class uFlow(libpry.AutoTree):
     def test_getset_state(self):
         f = tflow()
         state = f.get_state() 
-        assert f == console.ConsoleFlow.from_state(state)
+        assert f == flow.Flow.from_state(state)
         f.response = tresp()
         f.request = f.response.request
         state = f.get_state() 
-        assert f == console.ConsoleFlow.from_state(state)
+        assert f == flow.Flow.from_state(state)
 
     def test_simple(self):
         f = tflow()
-        assert f.get_text()
+        assert console.format_flow(f, True)
+        assert console.format_flow(f, False)
 
         f.request = treq()
-        assert f.get_text()
+        assert console.format_flow(f, True)
+        assert console.format_flow(f, False)
 
         f.response = tresp()
         f.response.headers["content-type"] = ["text/html"]
-        assert f.get_text()
+        assert console.format_flow(f, True)
+        assert console.format_flow(f, False)
         f.response.code = 404
-        assert f.get_text()
+        assert console.format_flow(f, True)
+        assert console.format_flow(f, False)
 
         f.focus = True
-        assert f.get_text()
+        assert console.format_flow(f, True)
+        assert console.format_flow(f, False)
 
         f.connection = flow.ReplayConnection()
-        assert f.get_text()
+        assert console.format_flow(f, True)
+        assert console.format_flow(f, False)
 
         f.response = None
-        assert f.get_text()
+        assert console.format_flow(f, True)
+        assert console.format_flow(f, False)
 
         f.error = proxy.Error(200, "test")
-        assert f.get_text()
+        assert console.format_flow(f, True)
+        assert console.format_flow(f, False)
 
     def test_kill(self):
         f = tflow()
@@ -274,7 +282,7 @@ class uFlow(libpry.AutoTree):
         assert f.response.acked
 
     def test_serialization(self):
-        f = console.ConsoleFlow(None)
+        f = flow.Flow(None)
         f.request = treq()