Add serialization hooks to flows and flow component objects.

7 files changed, 148 insertions, 1 deletions

diff --git a/libmproxy/controller.py b/libmproxy/controller.py
index 98d02c51..29e2f7a8 100644
--- a/libmproxy/controller.py
+++ b/libmproxy/controller.py
@@ -51,7 +51,7 @@ class Master:
 
     def tick(self, q):
         try:
-            # This endless loop is running until the 'Queue.Empty'
+            # This endless loop runs until the 'Queue.Empty'
             # exception is thrown. If more than one request is in
             # the queue, this speeds up every request by 0.1 seconds,
             # because get_input(..) function is not blocking.
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index bdca5ebd..41064829 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -31,6 +31,27 @@ class Flow:
         self.intercepting = False
         self._backup = None
 
+    def get_state(self):
+        return dict(
+            request = self.request.get_state() if self.request else None,
+            response = self.response.get_state() if self.response else None,
+            error = self.error.get_state() if self.error else None,
+        )
+
+    @classmethod
+    def from_state(klass, state):
+        f = Flow(ReplayConnection)
+        if state["request"]:
+            f.request = proxy.Request.from_state(state["request"])
+        if state["response"]:
+            f.response = proxy.Response.from_state(f.request, state["response"])
+        if state["error"]:
+            f.error = proxy.Error.from_state(state["error"])
+        return f
+
+    def __eq__(self, other):
+        return self.get_state() == other.get_state()
+
     def backup(self):
         if not self._backup:
             self._backup = [
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index 9f87f80e..f3b9fe53 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -88,6 +88,33 @@ class Request(controller.Msg):
         self.kill = False
         controller.Msg.__init__(self)
 
+    def get_state(self):
+        return dict(
+            host = self.host,
+            port = self.port,
+            scheme = self.scheme,
+            method = self.method,
+            path = self.path,
+            headers = self.headers.get_state(),
+            content = self.content
+        )
+
+    @classmethod
+    def from_state(klass, state):
+        return klass(
+            None,
+            state["host"],
+            state["port"],
+            state["scheme"],
+            state["method"],
+            state["path"],
+            utils.Headers.from_state(state["headers"]),
+            state["content"]
+        )
+
+    def __eq__(self, other):
+        return self.get_state() == other.get_state()
+
     def copy(self):
         c = copy.copy(self)
         c.headers = self.headers.copy()
@@ -137,6 +164,29 @@ class Response(controller.Msg):
         self.kill = False
         controller.Msg.__init__(self)
 
+    def get_state(self):
+        return dict(
+            code = self.code,
+            proto = self.proto,
+            msg = self.msg,
+            headers = self.headers.get_state(),
+            content = self.content
+        )
+
+    @classmethod
+    def from_state(klass, request, state):
+        return klass(
+            request,
+            state["code"],
+            state["proto"],
+            state["msg"],
+            utils.Headers.from_state(state["headers"]),
+            state["content"]
+        )
+
+    def __eq__(self, other):
+        return self.get_state() == other.get_state()
+
     def copy(self):
         c = copy.copy(self)
         c.headers = self.headers.copy()
@@ -181,6 +231,21 @@ class Error(controller.Msg):
     def copy(self):
         return copy.copy(self)
 
+    def get_state(self):
+        return dict(
+            msg = self.msg,
+        )
+
+    @classmethod
+    def from_state(klass, state):
+        return klass(
+            None,
+            state["msg"],
+        )
+
+    def __eq__(self, other):
+        return self.get_state() == other.get_state()
+
 
 class FileLike:
     def __init__(self, o):
diff --git a/libmproxy/utils.py b/libmproxy/utils.py
index 82ddf8fc..0614a464 100644
--- a/libmproxy/utils.py
+++ b/libmproxy/utils.py
@@ -156,6 +156,16 @@ class MultiDict:
             for j in self[i]:
                 yield (i, j)
 
+    def get_state(self):
+        return list(self.itemPairs())
+
+    @classmethod
+    def from_state(klass, state):
+        md = klass()
+        for i in state:
+            md.append(*i)
+        return md
+
 
 class Headers(MultiDict):
     """
diff --git a/test/test_console.py b/test/test_console.py
index cda61bd5..399cc485 100644
--- a/test/test_console.py
+++ b/test/test_console.py
@@ -196,6 +196,15 @@ class uFlow(libpry.AutoTree):
         f.backup()
         f.revert()
 
+    def test_getset_state(self):
+        f = tflow()
+        state = f.get_state() 
+        assert f == console.ConsoleFlow.from_state(state)
+        f.response = tresp()
+        f.request = f.response.request
+        state = f.get_state() 
+        assert f == console.ConsoleFlow.from_state(state)
+
     def test_simple(self):
         f = tflow()
         assert f.get_text()
diff --git a/test/test_proxy.py b/test/test_proxy.py
index 90cfbbfb..ea1c56aa 100644
--- a/test/test_proxy.py
+++ b/test/test_proxy.py
@@ -225,6 +225,14 @@ class uRequest(libpry.AutoTree):
         assert r.short()
         assert r.assemble()
 
+    def test_getset_state(self):
+        h = utils.Headers()
+        h["test"] = ["test"]
+        c = proxy.BrowserConnection("addr", 2222)
+        r = proxy.Request(c, "host", 22, "https", "GET", "/", h, "content")
+        state = r.get_state()
+        assert proxy.Request.from_state(state) == r
+
 
 class uResponse(libpry.AutoTree):
     def test_simple(self):
@@ -236,6 +244,24 @@ class uResponse(libpry.AutoTree):
         assert resp.short()
         assert resp.assemble()
 
+    def test_getset_state(self):
+        h = utils.Headers()
+        h["test"] = ["test"]
+        c = proxy.BrowserConnection("addr", 2222)
+        r = proxy.Request(c, "host", 22, "https", "GET", "/", h, "content")
+        req = proxy.Request(c, "host", 22, "https", "GET", "/", h, "content")
+        resp = proxy.Response(req, 200, "HTTP", "msg", h.copy(), "content")
+
+        state = resp.get_state()
+        assert proxy.Response.from_state(req, state) == resp
+
+
+class uError(libpry.AutoTree):
+    def test_getset_state(self):
+        e = proxy.Error(None, "Error")
+        state = e.get_state()
+        assert proxy.Error.from_state(state) == e
+
 
 class uProxyError(libpry.AutoTree):
     def test_simple(self):
@@ -252,6 +278,7 @@ tests = [
     uConfig(),
     u_parse_proxy_request(),
     u_parse_url(),
+    uError(),
     _TestServers(), [
         uSanity(),
         uProxy(),
diff --git a/test/test_utils.py b/test/test_utils.py
index 8a4da968..cd435005 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -102,6 +102,14 @@ class uMultiDict(libpry.AutoTree):
         assert ("foo", 2) in l
         assert ("bar", 3) in l
 
+    def test_getset_state(self):
+        self.md.append("foo", 1)
+        self.md.append("foo", 2)
+        self.md.append("bar", 3)
+        state = self.md.get_state()
+        nd = utils.MultiDict.from_state(state)
+        assert nd == self.md
+
 
 class uHeaders(libpry.AutoTree):
     def setUp(self):
@@ -176,6 +184,13 @@ class uHeaders(libpry.AutoTree):
         assert h.match_re("two: due")
         assert not h.match_re("nonono")
 
+    def test_getset_state(self):
+        self.hd.append("foo", 1)
+        self.hd.append("foo", 2)
+        self.hd.append("bar", 3)
+        state = self.hd.get_state()
+        nd = utils.Headers.from_state(state)
+        assert nd == self.hd
 
 
 class uisStringLike(libpry.AutoTree):