From f7e4e89b124e0bc20cf8192227add04bd92846ae Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Wed, 3 Aug 2011 13:33:18 +1200
Subject: Move the event notification mechanism into flow.py

---
 libmproxy/console.py | 25 ++++++-------------------
 libmproxy/dump.py    | 17 ++++-------------
 libmproxy/flow.py    | 16 ++++++++++++++++
 3 files changed, 26 insertions(+), 32 deletions(-)

(limited to 'libmproxy')

diff --git a/libmproxy/console.py b/libmproxy/console.py
index ae3210e7..5509f275 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -1766,31 +1766,18 @@ class ConsoleMaster(flow.FlowMaster):
         self.sync_list_view()
         self.refresh_connection(f)
 
-    def add_event(self, e):
+    def add_event(self, e, level="info"):
+        if level == "info":
+            e = urwid.Text(e)
+        elif level == "error":
+            e = urwid.Text(("error", e))
+
         self.eventlist.append(e)
         if len(self.eventlist) > EVENTLOG_SIZE:
             self.eventlist.pop(0)
         self.eventlist.set_focus(len(self.eventlist))
 
     # Handlers
-    def handle_clientconnect(self, r):
-        self.add_event(urwid.Text("Connect from: %s:%s"%r.address))
-        return flow.FlowMaster.handle_clientconnect(self, r)
-
-    def handle_clientdisconnect(self, r):
-        s = "Disconnect from: %s:%s"%r.client_conn.address
-        self.add_event(urwid.Text(s))
-        if r.client_conn.requestcount:
-            s = "    -> handled %s requests"%r.client_conn.requestcount
-            self.add_event(urwid.Text(s))
-        if r.client_conn.connection_error:
-            self.add_event(
-                urwid.Text(
-                        ("error", "   -> error: %s"%r.client_conn.connection_error)
-                )
-            )
-        return flow.FlowMaster.handle_clientdisconnect(self, r)
-
     def handle_error(self, r):
         f = flow.FlowMaster.handle_error(self, r)
         if f:
diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index 4d437b38..a47066d9 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -20,6 +20,7 @@ class Options(object):
         "rfile",
         "rheaders",
         "server_replay",
+        "script",
         "stickycookie",
         "stickyauth",
         "verbosity",
@@ -132,23 +133,13 @@ class DumpMaster(flow.FlowMaster):
                     "%s: %s\n%s"%(script, e.args[0], eout)
                 )
 
-    def handle_clientconnect(self, c):
+    def add_event(self, e, level="info"):
         if self.eventlog:
-            print >> self.outfile, "Connection from: %s:%s"%c.address
-        return flow.FlowMaster.handle_clientconnect(self, c)
-
-    def handle_clientdisconnect(self, c):
-        if self.eventlog:
-            print >> self.outfile, "Disconnect from: %s:%s"%tuple(c.client_conn.address),
-            print >> self.outfile, "(handled %s requests)"%c.client_conn.requestcount
-            if c.client_conn.connection_error:
-                print >> self.outfile, "\terror: %s"%c.client_conn.connection_error
-        return flow.FlowMaster.handle_clientconnect(self, c)
+            print >> self.outfile, e
 
     def handle_request(self, r):
         f = flow.FlowMaster.handle_request(self, r)
-        if self.eventlog:
-            print >> self.outfile, "Request: %s"%str_request(r)
+        self.add_event("Request: %s"%str_request(r))
         if f:
             r.ack()
         return f
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 92a4e812..49cf0796 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -481,6 +481,12 @@ class FlowMaster(controller.Master):
         raise NotImplementedError
         #end nocover
 
+    def add_event(self, e, level="info"):
+        """
+            level: info, error
+        """
+        pass
+
     def set_plugin(self, p):
         self.plugin = p
 
@@ -615,9 +621,19 @@ class FlowMaster(controller.Master):
         #end nocover
 
     def handle_clientconnect(self, r):
+        self.add_event("Connect from: %s:%s"%r.address)
         r.ack()
 
     def handle_clientdisconnect(self, r):
+        s = "Disconnect from: %s:%s"%r.client_conn.address
+        self.add_event(s)
+        if r.client_conn.requestcount:
+            s = "    -> handled %s requests"%r.client_conn.requestcount
+            self.add_event(s)
+        if r.client_conn.connection_error:
+            self.add_event(
+                "   -> error: %s"%r.client_conn.connection_error, "error"
+            )
         r.ack()
 
     def handle_error(self, r):
-- 
cgit v1.2.3