From 85ddc5056b982f16451fc97c655e54843b76b377 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@corte.si>
Date: Thu, 16 Mar 2017 11:21:44 +1300
Subject: [PATCH] Move mitmproxy console handlers to addon

---
 mitmproxy/tools/console/master.py | 62 +++++++++++++++----------------
 1 file changed, 30 insertions(+), 32 deletions(-)

diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py
index b6339817c..d39e5e5d3 100644
--- a/mitmproxy/tools/console/master.py
+++ b/mitmproxy/tools/console/master.py
@@ -40,6 +40,35 @@ class Logger:
         signals.add_log(evt.msg, evt.level)
 
 
+class UnsupportedLog:
+    """
+        A small addon to dump info on flow types we don't support yet.
+    """
+    def websocket_message(self, f):
+        message = f.messages[-1]
+        signals.add_log(f.message_info(message), "info")
+        signals.add_log(strutils.bytes_to_escaped_str(message.content), "debug")
+
+    def websocket_end(self, f):
+        signals.add_log("WebSocket connection closed by {}: {} {}, {}".format(
+            f.close_sender,
+            f.close_code,
+            f.close_message,
+            f.close_reason), "info")
+
+    def tcp_message(self, f):
+        message = f.messages[-1]
+        direction = "->" if message.from_client else "<-"
+        signals.add_log("{client_host}:{client_port} {direction} tcp {direction} {server_host}:{server_port}".format(
+            client_host=f.client_conn.address[0],
+            client_port=f.client_conn.address[1],
+            server_host=f.server_conn.address[0],
+            server_port=f.server_conn.address[1],
+            direction=direction,
+        ), "info")
+        signals.add_log(strutils.bytes_to_escaped_str(message.content), "debug")
+
+
 class ConsoleMaster(master.Master):
     palette = []
 
@@ -63,7 +92,7 @@ class ConsoleMaster(master.Master):
         signals.sig_add_log.connect(self.sig_add_log)
         self.addons.add(Logger())
         self.addons.add(*addons.default_addons())
-        self.addons.add(intercept.Intercept(), self.view)
+        self.addons.add(intercept.Intercept(), self.view, UnsupportedLog())
 
         def sigint_handler(*args, **kwargs):
             self.prompt_for_exit()
@@ -396,34 +425,3 @@ class ConsoleMaster(master.Master):
 
     def clear_events(self):
         self.logbuffer[:] = []
-
-    # Handlers
-    @controller.handler
-    def websocket_message(self, f):
-        super().websocket_message(f)
-        message = f.messages[-1]
-        signals.add_log(f.message_info(message), "info")
-        signals.add_log(strutils.bytes_to_escaped_str(message.content), "debug")
-
-    @controller.handler
-    def websocket_end(self, f):
-        super().websocket_end(f)
-        signals.add_log("WebSocket connection closed by {}: {} {}, {}".format(
-            f.close_sender,
-            f.close_code,
-            f.close_message,
-            f.close_reason), "info")
-
-    @controller.handler
-    def tcp_message(self, f):
-        super().tcp_message(f)
-        message = f.messages[-1]
-        direction = "->" if message.from_client else "<-"
-        signals.add_log("{client_host}:{client_port} {direction} tcp {direction} {server_host}:{server_port}".format(
-            client_host=f.client_conn.address[0],
-            client_port=f.client_conn.address[1],
-            server_host=f.server_conn.address[0],
-            server_port=f.server_conn.address[1],
-            direction=direction,
-        ), "info")
-        signals.add_log(strutils.bytes_to_escaped_str(message.content), "debug")