Clean up and clarify StateObject

- Flatten the class hierarchy
- get_state, load_state, from_state are public
- Simplify code
- Remove __eq__ and __neq__. This fundamentally changes the semantics of
inherited objects in a way that's not part of the core function of the
class

3 files changed, 43 insertions, 24 deletions

diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index 46da7a2b..1f3d6fdf 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -85,7 +85,7 @@ class decoded(object):
             self.o.encode(self.ce)
 
 
-class HTTPMessage(stateobject.SimpleStateObject):
+class HTTPMessage(stateobject.StateObject):
     """
     Base class for HTTPRequest and HTTPResponse
     """
@@ -275,9 +275,9 @@ class HTTPRequest(HTTPMessage):
     )
 
     @classmethod
-    def _from_state(cls, state):
+    def from_state(cls, state):
         f = cls(None, None, None, None, None, None, None, None, None, None, None)
-        f._load_state(state)
+        f.load_state(state)
         return f
 
     def __repr__(self):
@@ -626,9 +626,9 @@ class HTTPResponse(HTTPMessage):
     )
 
     @classmethod
-    def _from_state(cls, state):
+    def from_state(cls, state):
         f = cls(None, None, None, None, None)
-        f._load_state(state)
+        f.load_state(state)
         return f
 
     def __repr__(self):
@@ -814,9 +814,9 @@ class HTTPFlow(Flow):
     )
 
     @classmethod
-    def _from_state(cls, state):
+    def from_state(cls, state):
         f = cls(None, None)
-        f._load_state(state)
+        f.load_state(state)
         return f
 
     def __repr__(self):
@@ -1311,4 +1311,4 @@ class RequestReplayThread(threading.Thread):
             self.flow.error = Error(repr(v))
             self.channel.ask("error", self.flow)
         finally:
-            r.form_out = form_out_backup
\ No newline at end of file
+            r.form_out = form_out_backup
diff --git a/libmproxy/protocol/primitives.py b/libmproxy/protocol/primitives.py
index a8c5856c..77dc936d 100644
--- a/libmproxy/protocol/primitives.py
+++ b/libmproxy/protocol/primitives.py
@@ -8,7 +8,7 @@ from ..proxy.connection import ClientConnection, ServerConnection
 KILL = 0  # const for killed requests
 
 
-class Error(stateobject.SimpleStateObject):
+class Error(stateobject.StateObject):
     """
         An Error.
 
@@ -41,11 +41,11 @@ class Error(stateobject.SimpleStateObject):
         return self.msg
 
     @classmethod
-    def _from_state(cls, state):
+    def from_state(cls, state):
         # the default implementation assumes an empty constructor. Override
         # accordingly.
         f = cls(None)
-        f._load_state(state)
+        f.load_state(state)
         return f
 
     def copy(self):
@@ -53,7 +53,7 @@ class Error(stateobject.SimpleStateObject):
         return c
 
 
-class Flow(stateobject.SimpleStateObject):
+class Flow(stateobject.StateObject):
     """
     A Flow is a collection of objects representing a single transaction.
     This class is usually subclassed for each protocol, e.g. HTTPFlow.
@@ -78,8 +78,8 @@ class Flow(stateobject.SimpleStateObject):
         conntype=str
     )
 
-    def _get_state(self):
-        d = super(Flow, self)._get_state()
+    def get_state(self):
+        d = super(Flow, self).get_state()
         d.update(version=version.IVERSION)
         return d
 
@@ -101,7 +101,7 @@ class Flow(stateobject.SimpleStateObject):
             Has this Flow been modified?
         """
         if self._backup:
-            return self._backup != self._get_state()
+            return self._backup != self.get_state()
         else:
             return False
 
@@ -111,14 +111,14 @@ class Flow(stateobject.SimpleStateObject):
             call to .revert().
         """
         if not self._backup:
-            self._backup = self._get_state()
+            self._backup = self.get_state()
 
     def revert(self):
         """
             Revert to the last backed up state.
         """
         if self._backup:
-            self._load_state(self._backup)
+            self.load_state(self._backup)
             self._backup = None
 
 
diff --git a/libmproxy/protocol/tcp.py b/libmproxy/protocol/tcp.py
index 990c502a..a56bf07b 100644
--- a/libmproxy/protocol/tcp.py
+++ b/libmproxy/protocol/tcp.py
@@ -1,8 +1,10 @@
 from __future__ import absolute_import
-import select, socket
+import select
+import socket
 from .primitives import ProtocolHandler
 from netlib.utils import cleanBin
 
+
 class TCPHandler(ProtocolHandler):
     """
     TCPHandler acts as a generic TCP forwarder.
@@ -34,7 +36,9 @@ class TCPHandler(ProtocolHandler):
                     closed = False
                     if src.ssl_established:
                         # Unfortunately, pyOpenSSL lacks a recv_into function.
-                        contents = src.rfile.read(1)  # We need to read a single byte before .pending() becomes usable
+                        # We need to read a single byte before .pending()
+                        # becomes usable
+                        contents = src.rfile.read(1)
                         contents += src.rfile.read(src.connection.pending())
                         if not contents:
                             closed = True
@@ -56,15 +60,30 @@ class TCPHandler(ProtocolHandler):
                         continue
 
                     if src.ssl_established or dst.ssl_established:
-                        # if one of the peers is over SSL, we need to send bytes/strings
-                        if not src.ssl_established:  # only ssl to dst, i.e. we revc'd into buf but need bytes/string now.
+                        # if one of the peers is over SSL, we need to send
+                        # bytes/strings
+                        if not src.ssl_established:
+                            # only ssl to dst, i.e. we revc'd into buf but need
+                            # bytes/string now.
                             contents = buf[:size].tobytes()
-                        self.c.log("%s %s\r\n%s" % (direction, dst_str, cleanBin(contents)), "debug")
+                        self.c.log(
+                            "%s %s\r\n%s" % (
+                                direction, dst_str, cleanBin(contents)
+                            ),
+                            "debug"
+                        )
                         dst.connection.send(contents)
                     else:
                         # socket.socket.send supports raw bytearrays/memoryviews
-                        self.c.log("%s %s\r\n%s" % (direction, dst_str, cleanBin(buf.tobytes())), "debug")
+                        self.c.log(
+                            "%s %s\r\n%s" % (
+                                direction,
+                                dst_str,
+                                cleanBin(buf.tobytes())
+                            ),
+                            "debug"
+                        )
                         dst.connection.send(buf[:size])
         except socket.error as e:
             self.c.log("TCP connection closed unexpectedly.", "debug")
-            return
\ No newline at end of file
+            return