push failing tests down to 43

4 files changed, 165 insertions, 155 deletions

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index b4b939c7..b1971469 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -5,137 +5,20 @@
 import hashlib, Cookie, cookielib, copy, re, urlparse, threading
 import time, urllib
 import types
-import tnetstring, filt, script, utils, encoding, proxy
+import tnetstring, filt, script, utils, encoding
 from email.utils import parsedate_tz, formatdate, mktime_tz
 from netlib import odict, http, certutils, wsgi
 from .proxy import ClientConnection, ServerConnection
 import controller, version, protocol, stateobject
 import app
-
-
-HDR_FORM_URLENCODED = "application/x-www-form-urlencoded"
-CONTENT_MISSING = 0
+from .protocol import KILL
+from .protocol.http import HTTPResponse, CONTENT_MISSING
+from .proxy import RequestReplayThread
 
 ODict = odict.ODict
 ODictCaseless = odict.ODictCaseless
 
 
-class BackreferenceMixin(object):
-    """
-    If an attribute from the _backrefattr tuple is set,
-    this mixin sets a reference back on the attribute object.
-    Example:
-        e = Error()
-        f = Flow()
-        f.error = e
-        assert f is e.flow
-    """
-    _backrefattr = tuple()
-
-    def __setattr__(self, key, value):
-        super(BackreferenceMixin, self).__setattr__(key, value)
-        if key in self._backrefattr and value is not None:
-            setattr(value, self._backrefname, self)
-
-
-class Error(stateobject.SimpleStateObject):
-    """
-        An Error.
-
-        This is distinct from an HTTP error response (say, a code 500), which
-        is represented by a normal Response object. This class is responsible
-        for indicating errors that fall outside of normal HTTP communications,
-        like interrupted connections, timeouts, protocol errors.
-
-        Exposes the following attributes:
-
-            flow: Flow object
-            msg: Message describing the error
-            timestamp: Seconds since the epoch
-    """
-    def __init__(self, msg, timestamp=None):
-        """
-        @type msg: str
-        @type timestamp: float
-        """
-        self.msg = msg
-        self.timestamp = timestamp or utils.timestamp()
-
-    _stateobject_attributes = dict(
-        msg=str,
-        timestamp=float
-    )
-
-    @classmethod
-    def _from_state(cls, state):
-        f = cls(None)  # the default implementation assumes an empty constructor. Override accordingly.
-        f._load_state(state)
-        return f
-
-    def copy(self):
-        c = copy.copy(self)
-        return c
-
-
-class Flow(stateobject.SimpleStateObject, BackreferenceMixin):
-    def __init__(self, conntype, client_conn, server_conn):
-        self.conntype = conntype
-        self.client_conn = client_conn
-        self.server_conn = server_conn
-        self.error = None
-
-    _backrefattr = ("error",)
-    _backrefname = "flow"
-
-    _stateobject_attributes = dict(
-        error=Error,
-        client_conn=ClientConnection,
-        server_conn=ServerConnection,
-        conntype=str
-    )
-
-    def _get_state(self):
-        d = super(Flow, self)._get_state()
-        d.update(version=version.IVERSION)
-        return d
-
-    @classmethod
-    def _from_state(cls, state):
-        f = cls(None, None, None)
-        f._load_state(state)
-        return f
-
-    def copy(self):
-        f = copy.copy(self)
-        if self.error:
-            f.error = self.error.copy()
-        return f
-
-    def modified(self):
-        """
-            Has this Flow been modified?
-        """
-        if self._backup:
-            return self._backup != self._get_state()
-        else:
-            return False
-
-    def backup(self, force=False):
-        """
-            Save a backup of this Flow, which can be reverted to using a
-            call to .revert().
-        """
-        if not self._backup:
-            self._backup = self._get_state()
-
-    def revert(self):
-        """
-            Revert to the last backed up state.
-        """
-        if self._backup:
-            self._load_state(self._backup)
-            self._backup = None
-
 
 class AppRegistry:
     def __init__(self):
@@ -660,10 +543,8 @@ class FlowMaster(controller.Master):
             rflow = self.server_playback.next_flow(flow)
             if not rflow:
                 return None
-            # FIXME
-            response = Response._from_state(flow.request, rflow.response._get_state())
-            response._set_replay()
-            flow.response = response
+            response = HTTPResponse._from_state(rflow.response._get_state())
+            response.is_replay = True
             if self.refresh_server_playback:
                 response.refresh()
             flow.request.reply(response)
@@ -742,13 +623,13 @@ class FlowMaster(controller.Master):
         if f.request.content == CONTENT_MISSING:
             return "Can't replay request with missing content..."
         if f.request:
-            f.request._set_replay()
+            f.request.is_replay = True
             if f.request.content:
                 f.request.headers["Content-Length"] = [str(len(f.request.content))]
             f.response = None
             f.error = None
             self.process_new_request(f)
-            rt = proxy.RequestReplayThread(
+            rt = RequestReplayThread(
                     self.server.config,
                     f,
                     self.masterq,
@@ -791,7 +672,7 @@ class FlowMaster(controller.Master):
                 err = app.serve(r, r.wfile, **{"mitmproxy.master": self})
                 if err:
                     self.add_event("Error in wsgi app. %s"%err, "error")
-                r.reply(proxy.KILL)
+                r.reply(KILL)
                 return
         f = self.state.add_request(r)
         self.replacehooks.run(f)
diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index 8c44461e..be60f374 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -6,7 +6,7 @@ from netlib.odict import ODict, ODictCaseless
 from . import ProtocolHandler, ConnectionTypeChange, KILL
 from .. import encoding, utils, version, filt, controller, stateobject
 from ..proxy import ProxyError, AddressPriority
-from ..flow import Flow, Error
+from .primitives import Flow, Error
 
 
 HDR_FORM_URLENCODED = "application/x-www-form-urlencoded"
@@ -340,7 +340,7 @@ class HTTPRequest(HTTPMessage):
             Raises an Exception if the request cannot be assembled.
         """
         if self.content == CONTENT_MISSING:
-            raise RuntimeError("Cannot assemble flow with CONTENT_MISSING")
+            raise ProxyError(502, "Cannot assemble flow with CONTENT_MISSING")
         head = self._assemble_head(form)
         if self.content:
             return head + self.content
@@ -444,6 +444,8 @@ class HTTPRequest(HTTPMessage):
             If hostheader is True, we use the value specified in the request
             Host header to construct the URL.
         """
+        raise NotImplementedError
+        # FIXME: Take server_conn into account.
         host = None
         if hostheader:
             host = self.headers.get_first("host")
@@ -462,6 +464,8 @@ class HTTPRequest(HTTPMessage):
 
             Returns False if the URL was invalid, True if the request succeeded.
         """
+        raise NotImplementedError
+        # FIXME: Needs to update server_conn as well.
         parts = http.parse_url(url)
         if not parts:
             return False
@@ -595,7 +599,7 @@ class HTTPResponse(HTTPMessage):
             Raises an Exception if the request cannot be assembled.
         """
         if self.content == CONTENT_MISSING:
-            raise RuntimeError("Cannot assemble flow with CONTENT_MISSING")
+            raise ProxyError(502, "Cannot assemble flow with CONTENT_MISSING")
         head = self._assemble_head()
         if self.content:
             return head + self.content
@@ -711,7 +715,7 @@ class HTTPFlow(Flow):
         if self.request:
             f.request = self.request.copy()
         if self.response:
-            f.response = self.request.copy()
+            f.response = self.response.copy()
         return f
 
     def match(self, f):
@@ -795,8 +799,7 @@ class HTTPHandler(ProtocolHandler):
 
         for i in range(2):
             try:
-                self.c.server_conn.wfile.write(request_raw)
-                self.c.server_conn.wfile.flush()
+                self.c.server_conn.send(request_raw)
                 return HTTPResponse.from_stream(self.c.server_conn.rfile, request.method,
                                                 body_size_limit=self.c.config.body_size_limit)
             except (tcp.NetLibDisconnect, http.HttpErrorConnClosed), v:
@@ -821,6 +824,7 @@ class HTTPHandler(ProtocolHandler):
             flow.request = HTTPRequest.from_stream(self.c.client_conn.rfile,
                                                    body_size_limit=self.c.config.body_size_limit)
             self.c.log("request", [flow.request._assemble_first_line(flow.request.form_in)])
+            self.process_request(flow.request)
 
             request_reply = self.c.channel.ask("request" if LEGACY else "httprequest",
                                                flow.request if LEGACY else flow)
@@ -830,7 +834,6 @@ class HTTPHandler(ProtocolHandler):
             if isinstance(request_reply, HTTPResponse):
                 flow.response = request_reply
             else:
-                self.process_request(flow.request)
                 self.c.establish_server_connection()
                 flow.response = self.get_response_from_server(flow.request)
 
diff --git a/libmproxy/protocol/primitives.py b/libmproxy/protocol/primitives.py
new file mode 100644
index 00000000..f77e097b
--- /dev/null
+++ b/libmproxy/protocol/primitives.py
@@ -0,0 +1,124 @@
+from .. import stateobject, utils, version
+from ..proxy import ServerConnection, ClientConnection
+import copy
+
+
+class _BackreferenceMixin(object):
+    """
+    If an attribute from the _backrefattr tuple is set,
+    this mixin sets a reference back on the attribute object.
+    Example:
+        e = Error()
+        f = Flow()
+        f.error = e
+        assert f is e.flow
+    """
+    _backrefattr = tuple()
+
+    def __setattr__(self, key, value):
+        super(_BackreferenceMixin, self).__setattr__(key, value)
+        if key in self._backrefattr and value is not None:
+            setattr(value, self._backrefname, self)
+
+
+class Error(stateobject.SimpleStateObject):
+    """
+        An Error.
+
+        This is distinct from an HTTP error response (say, a code 500), which
+        is represented by a normal Response object. This class is responsible
+        for indicating errors that fall outside of normal HTTP communications,
+        like interrupted connections, timeouts, protocol errors.
+
+        Exposes the following attributes:
+
+            flow: Flow object
+            msg: Message describing the error
+            timestamp: Seconds since the epoch
+    """
+    def __init__(self, msg, timestamp=None):
+        """
+        @type msg: str
+        @type timestamp: float
+        """
+        self.msg = msg
+        self.timestamp = timestamp or utils.timestamp()
+
+    _stateobject_attributes = dict(
+        msg=str,
+        timestamp=float
+    )
+
+    @classmethod
+    def _from_state(cls, state):
+        f = cls(None)  # the default implementation assumes an empty constructor. Override accordingly.
+        f._load_state(state)
+        return f
+
+    def copy(self):
+        c = copy.copy(self)
+        return c
+
+
+class Flow(stateobject.SimpleStateObject, _BackreferenceMixin):
+    def __init__(self, conntype, client_conn, server_conn):
+        self.conntype = conntype
+        self.client_conn = client_conn
+        self.server_conn = server_conn
+        self.error = None
+
+    _backrefattr = ("error",)
+    _backrefname = "flow"
+
+    _stateobject_attributes = dict(
+        error=Error,
+        client_conn=ClientConnection,
+        server_conn=ServerConnection,
+        conntype=str
+    )
+
+    def _get_state(self):
+        d = super(Flow, self)._get_state()
+        d.update(version=version.IVERSION)
+        return d
+
+    @classmethod
+    def _from_state(cls, state):
+        f = cls(None, None, None)
+        f._load_state(state)
+        return f
+
+    def copy(self):
+        f = copy.copy(self)
+
+        f.client_conn = self.client_conn.copy()
+        f.server_conn = self.server_conn.copy()
+
+        if self.error:
+            f.error = self.error.copy()
+        return f
+
+    def modified(self):
+        """
+            Has this Flow been modified?
+        """
+        if self._backup:
+            return self._backup != self._get_state()
+        else:
+            return False
+
+    def backup(self, force=False):
+        """
+            Save a backup of this Flow, which can be reverted to using a
+            call to .revert().
+        """
+        if not self._backup:
+            self._backup = self._get_state()
+
+    def revert(self):
+        """
+            Revert to the last backed up state.
+        """
+        if self._backup:
+            self._load_state(self._backup)
+            self._backup = None
\ No newline at end of file
diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py
index 4842a81f..feff2259 100644
--- a/libmproxy/proxy.py
+++ b/libmproxy/proxy.py
@@ -1,9 +1,8 @@
-import os, socket, time, threading
+import os, socket, time, threading, copy
 from OpenSSL import SSL
 from netlib import tcp, http, certutils, http_auth
 import utils, version, platform, controller, stateobject
 
-
 TRANSPARENT_SSL_PORTS = [443, 8443]
 
 
@@ -82,6 +81,9 @@ class ClientConnection(tcp.BaseHandler, stateobject.SimpleStateObject):
         self.address = tcp.Address(**state["address"]) if state["address"] else None
         self.clientcert = certutils.SSLCert.from_pem(state["clientcert"]) if state["clientcert"] else None
 
+    def copy(self):
+        return copy.copy(self)
+
     @classmethod
     def _from_state(cls, state):
         f = cls(None, None, None)
@@ -115,7 +117,9 @@ class ServerConnection(tcp.TCPClient, stateobject.SimpleStateObject):
         timestamp_ssl_setup=float,
         address=tcp.Address,
         source_address=tcp.Address,
-        cert=certutils.SSLCert
+        cert=certutils.SSLCert,
+        ssl_established=bool,
+        sni=str
     )
 
     def _get_state(self):
@@ -141,6 +145,9 @@ class ServerConnection(tcp.TCPClient, stateobject.SimpleStateObject):
         f._load_state(state)
         return f
 
+    def copy(self):
+        return copy.copy(self)
+
     def connect(self):
         self.timestamp_start = utils.timestamp()
         tcp.TCPClient.connect(self)
@@ -167,8 +174,10 @@ class ServerConnection(tcp.TCPClient, stateobject.SimpleStateObject):
         tcp.TCPClient.finish(self)
         self.timestamp_end = utils.timestamp()
 
+from . import protocol
+from .protocol.http import HTTPResponse
+
 
-"""
 class RequestReplayThread(threading.Thread):
     def __init__(self, config, flow, masterq):
         self.config, self.flow, self.channel = config, flow, controller.Channel(masterq)
@@ -177,24 +186,17 @@ class RequestReplayThread(threading.Thread):
     def run(self):
         try:
             r = self.flow.request
-            server = ServerConnection(self.config, r.scheme, r.host, r.port, r.host)
+            server = ServerConnection(self.flow.server_conn.address())
             server.connect()
-            server.send(r)
-            httpversion, code, msg, headers, content = http.read_response(
-                server.rfile, r.method, self.config.body_size_limit
-            )
-            response = flow.Response(
-                self.flow.request, httpversion, code, msg, headers, content, server.cert, 
-                server.rfile.first_byte_timestamp
-            )
-            self.channel.ask("response", response)
+            if self.flow.server_conn.ssl_established:
+                server.establish_ssl(self.config.clientcerts,
+                                     self.flow.server_conn.sni)
+            server.send(r._assemble())
+            self.flow.response = HTTPResponse.from_stream(server.rfile, r.method, body_size_limit=self.config.body_size_limit)
+            self.channel.ask("response", self.flow.response)
         except (ProxyError, http.HttpError, tcp.NetLibError), v:
-            err = flow.Error(str(v))
-            self.channel.ask("error", err)
-"""
-
-
-import protocol
+            self.flow.error = protocol.primitives.Error(str(v))
+            self.channel.ask("error", self.flow.error)
 
 class ConnectionHandler:
     def __init__(self, config, client_connection, client_address, server, channel, server_version):