Implement state loading that doesn't change object identity.

We need this to let us load state from copied Flows returned from scripts.

2 files changed, 57 insertions, 5 deletions

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 03b8b309..3520cc93 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -99,11 +99,26 @@ class Flow:
 
     def load_state(self, state):
         self._backup = state["backup"]
-        self.request = proxy.Request.from_state(state["request"])
+        if self.request:
+            self.request.load_state(state["request"])
+        else:
+            self.request = proxy.Request.from_state(state["request"])
+
         if state["response"]:
-            self.response = proxy.Response.from_state(self.request, state["response"])
+            if self.response:
+                self.response.load_state(state["response"])
+            else:
+                self.response = proxy.Response.from_state(self.request, state["response"])
+        else:
+            self.response = None
+
         if state["error"]:
-            self.error = proxy.Error.from_state(state["error"])
+            if self.error:
+                self.error.load_state(state["error"])
+            else:
+                self.error = proxy.Error.from_state(state["error"])
+        else:
+            self.error = None
 
     @classmethod
     def from_state(klass, state):
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index 88c62b25..4ab19694 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -148,6 +148,23 @@ class Request(controller.Msg):
     def is_cached(self):
         return False
 
+    def load_state(self, state):
+        if state["client_conn"]:
+            if self.client_conn:
+                self.client_conn.load_state(state["client_conn"])
+            else:
+                self.client_conn = ClientConnect.from_state(state["client_conn"])
+        else:
+            self.client_conn = None
+        self.host = state["host"]
+        self.port = state["port"]
+        self.scheme = state["scheme"]
+        self.method = state["method"]
+        self.path = state["path"]
+        self.headers = utils.Headers.from_state(state["headers"])
+        self.content = base64.decodestring(state["content"])
+        self.timestamp = state["timestamp"]
+
     def get_state(self):
         return dict(
             client_conn = self.client_conn.get_state() if self.client_conn else None,
@@ -164,7 +181,7 @@ class Request(controller.Msg):
     @classmethod
     def from_state(klass, state):
         return klass(
-            ClientConnect.from_state(state["client_conn"]) if state["client_conn"] else None,
+            ClientConnect.from_state(state["client_conn"]),
             state["host"],
             state["port"],
             state["scheme"],
@@ -249,6 +266,13 @@ class Response(controller.Msg):
         self.cached = False
         controller.Msg.__init__(self)
 
+    def load_state(self, state):
+        self.code = state["code"]
+        self.msg = state["msg"]
+        self.headers = utils.Headers.from_state(state["headers"])
+        self.content = base64.decodestring(state["content"])
+        self.timestamp = state["timestamp"]
+
     def get_state(self):
         return dict(
             code = self.code,
@@ -325,12 +349,21 @@ class ClientConnect(controller.Msg):
         self.close = False
         controller.Msg.__init__(self)
 
+    def __eq__(self, other):
+        return self.get_state() == other.get_state()
+
+    def load_state(self, state):
+        self.address = state
+
     def get_state(self):
         return list(self.address) if self.address else None
 
     @classmethod
     def from_state(klass, state):
-        return klass(state)
+        if state:
+            return klass(state)
+        else:
+            return None
 
     def copy(self):
         return copy.copy(self)
@@ -342,6 +375,10 @@ class Error(controller.Msg):
         self.timestamp = timestamp or time.time()
         controller.Msg.__init__(self)
 
+    def load_state(self, state):
+        self.msg = state["msg"]
+        self.timestamp = state["timestamp"]
+
     def copy(self):
         return copy.copy(self)