Introduce short form object state, and connect the ends to send data to web app

7 files changed, 58 insertions, 20 deletions

diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index 1f3d6fdf..1912390a 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -106,6 +106,7 @@ class HTTPMessage(stateobject.StateObject):
         timestamp_start=float,
         timestamp_end=float
     )
+    _stateobject_long_attributes = set(["content"])
 
     def get_decoded_content(self):
         """
diff --git a/libmproxy/protocol/primitives.py b/libmproxy/protocol/primitives.py
index 77dc936d..3d87e888 100644
--- a/libmproxy/protocol/primitives.py
+++ b/libmproxy/protocol/primitives.py
@@ -78,8 +78,8 @@ class Flow(stateobject.StateObject):
         conntype=str
     )
 
-    def get_state(self):
-        d = super(Flow, self).get_state()
+    def get_state(self, short=False):
+        d = super(Flow, self).get_state(short)
         d.update(version=version.IVERSION)
         return d
 
diff --git a/libmproxy/proxy/connection.py b/libmproxy/proxy/connection.py
index e0e94a2b..ff36d39c 100644
--- a/libmproxy/proxy/connection.py
+++ b/libmproxy/proxy/connection.py
@@ -36,8 +36,8 @@ class ClientConnection(tcp.BaseHandler, stateobject.StateObject):
         timestamp_ssl_setup=float
     )
 
-    def get_state(self):
-        d = super(ClientConnection, self).get_state()
+    def get_state(self, short=False):
+        d = super(ClientConnection, self).get_state(short)
         d.update(
             address={"address": self.address(), "use_ipv6": self.address.use_ipv6},
             clientcert=self.cert.to_pem() if self.clientcert else None
@@ -107,8 +107,8 @@ class ServerConnection(tcp.TCPClient, stateobject.StateObject):
         sni=str
     )
 
-    def get_state(self):
-        d = super(ServerConnection, self).get_state()
+    def get_state(self, short=False):
+        d = super(ServerConnection, self).get_state(short)
         d.update(
             address={"address": self.address(),
                      "use_ipv6": self.address.use_ipv6},
diff --git a/libmproxy/stateobject.py b/libmproxy/stateobject.py
index 37b72c7e..39614c6e 100644
--- a/libmproxy/stateobject.py
+++ b/libmproxy/stateobject.py
@@ -9,27 +9,35 @@ class StateObject(object):
         or StateObject instances themselves.
     """
     # An attribute-name -> class-or-type dict containing all attributes that
-    # should be serialized. If the attribute is a class, it must be a subclass
-    # of StateObject.
+    # should be serialized. If the attribute is a class, it must implement the
+    # StateObject protocol.
     _stateobject_attributes = None
-
-    def _get_state_attr(self, attr, cls):
-        val = getattr(self, attr)
-        if hasattr(val, "get_state"):
-            return val.get_state()
-        else:
-            return val
+    # A set() of attributes that should be ignored for short state
+    _stateobject_long_attributes = frozenset([])
 
     def from_state(self):
         raise NotImplementedError
 
-    def get_state(self):
+    def get_state(self, short=False):
+        """
+            Retrieve object state. If short is true, return an abbreviated
+            format with long data elided.
+        """
         state = {}
         for attr, cls in self._stateobject_attributes.iteritems():
-            state[attr] = self._get_state_attr(attr, cls)
+            if short and attr in self._stateobject_long_attributes:
+                continue
+            val = getattr(self, attr)
+            if hasattr(val, "get_state"):
+                state[attr] = val.get_state(short)
+            else:
+                state[attr] = val
         return state
 
     def load_state(self, state):
+        """
+            Load object state from data returned by a get_state call.
+        """
         for attr, cls in self._stateobject_attributes.iteritems():
             if state.get(attr, None) is None:
                 setattr(self, attr, None)
diff --git a/libmproxy/web/__init__.py b/libmproxy/web/__init__.py
index c2597861..50b41b80 100644
--- a/libmproxy/web/__init__.py
+++ b/libmproxy/web/__init__.py
@@ -81,18 +81,31 @@ class WebMaster(flow.FlowMaster):
             self.shutdown()
 
     def handle_request(self, f):
-        pprint.pprint(f.get_state())
+        app.ClientConnection.broadcast("flow", f.get_state(True))
         flow.FlowMaster.handle_request(self, f)
         if f:
             f.reply()
         return f
 
     def handle_response(self, f):
+        app.ClientConnection.broadcast("flow", f.get_state(True))
         flow.FlowMaster.handle_response(self, f)
         if f:
             f.reply()
         return f
 
     def handle_error(self, f):
+        app.ClientConnection.broadcast("flow", f.get_state(True))
         flow.FlowMaster.handle_error(self, f)
         return f
+
+    def handle_log(self, l):
+        app.ClientConnection.broadcast(
+            "event", {
+                "message": l.msg,
+                "level": l.level
+            }
+        )
+        self.add_event(l.msg, l.level)
+        l.reply()
+
diff --git a/libmproxy/web/app.py b/libmproxy/web/app.py
index e9bcc526..e2765a6d 100644
--- a/libmproxy/web/app.py
+++ b/libmproxy/web/app.py
@@ -2,6 +2,7 @@ import os.path
 import tornado.web
 import tornado.websocket
 import logging
+import json
 
 
 class IndexHandler(tornado.web.RequestHandler):
@@ -22,7 +23,14 @@ class ClientConnection(tornado.websocket.WebSocketHandler):
     def broadcast(cls, type, data):
         for conn in cls.connections:
             try:
-                conn.write_message(type, data)
+                conn.write_message(
+                    json.dumps(
+                        {
+                            "type": type,
+                            "data": data
+                        }
+                    )
+                )
             except:
                 logging.error("Error sending message", exc_info=True)
 
diff --git a/libmproxy/web/static/js/app.js b/libmproxy/web/static/js/app.js
index 90a004ef..84eece87 100644
--- a/libmproxy/web/static/js/app.js
+++ b/libmproxy/web/static/js/app.js
@@ -221,7 +221,15 @@ _Connection.prototype.onopen = function (open) {
 };
 _Connection.prototype.onmessage = function (message) {
     //AppDispatcher.dispatchServerAction(...);
-    console.log("onmessage", this, arguments);
+    var m = JSON.parse(message.data);
+    switch (m.type){
+        case "flow":
+            console.log("flow", m.data);
+            break;
+        case "event":
+            console.log("event", m.data.message)
+            break;
+    }
 };
 _Connection.prototype.onerror = function (error) {
     console.log("onerror", this, arguments);