From bb4f9611f53771103b9b790d7fa3baddb706cfe9 Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Tue, 10 May 2016 12:18:05 -0600
Subject: Revert "fix tests"

This reverts commit f315dc1eb9f4701548fd57bbb38de1e9caa34e40.
---
 mitmproxy/flow.py                | 50 +++-----------------------------
 mitmproxy/models/__init__.py     |  5 +---
 mitmproxy/models/flow.py         |  9 ------
 mitmproxy/models/http.py         |  6 ++++
 mitmproxy/protocol/rawtcp.py     | 61 +++++++++++++++++++++++++---------------
 mitmproxy/proxy/root_context.py  |  2 +-
 test/netlib/http/test_request.py |  4 +--
 7 files changed, 53 insertions(+), 84 deletions(-)

diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py
index 9e6d3f89..9292e76a 100644
--- a/mitmproxy/flow.py
+++ b/mitmproxy/flow.py
@@ -16,13 +16,12 @@ import re
 from netlib import wsgi
 from netlib.exceptions import HttpException
 from netlib.http import Headers, http1
-from netlib.utils import clean_bin
 from . import controller, tnetstring, filt, script, version, flow_format_compat
 from .onboarding import app
 from .proxy.config import HostMatcher
 from .protocol.http_replay import RequestReplayThread
 from .exceptions import Kill, FlowReadException
-from .models import ClientConnection, ServerConnection, HTTPFlow, HTTPRequest, FLOW_TYPES, TCPFlow
+from .models import ClientConnection, ServerConnection, HTTPFlow, HTTPRequest, FLOW_TYPES
 from collections import defaultdict
 
 
@@ -893,17 +892,6 @@ class FlowMaster(controller.ServerMaster):
                 self.handle_response(f)
             if f.error:
                 self.handle_error(f)
-        elif isinstance(f, TCPFlow):
-            messages = f.messages
-            f.messages = []
-            f.reply = controller.DummyReply()
-            self.handle_tcp_open(f)
-            while messages:
-                f.messages.append(messages.pop(0))
-                self.handle_tcp_message(f)
-            if f.error:
-                self.handle_tcp_error(f)
-            self.handle_tcp_close(f)
         else:
             raise NotImplementedError()
 
@@ -1091,39 +1079,9 @@ class FlowMaster(controller.ServerMaster):
             self.add_event('"{}" reloaded.'.format(s.filename), 'info')
         return ok
 
-    def handle_tcp_open(self, flow):
-        self.state.add_flow(flow)
-        self.run_script_hook("tcp_open", flow)
-        flow.reply()
-
-    def handle_tcp_message(self, flow):
-        self.run_script_hook("tcp_message", flow)
-        message = flow.messages[-1]
-        direction = "->" if message.from_client else "<-"
-        self.add_event("{client} {direction} tcp {direction} {server}".format(
-            client=repr(flow.client_conn.address),
-            server=repr(flow.server_conn.address),
-            direction=direction,
-        ), "info")
-        self.add_event(clean_bin(message.content), "debug")
-        flow.reply()
-
-    def handle_tcp_error(self, flow):
-        if self.stream:
-            self.stream.add(flow)
-        self.add_event("Error in TCP connection to {}: {}".format(
-            repr(flow.server_conn.address),
-            flow.error
-        ), "info")
-        self.run_script_hook("tcp_error", flow)
-        flow.reply()
-
-    def handle_tcp_close(self, flow):
-        self.state.delete_flow(flow)
-        if self.stream:
-            self.stream.add(flow)
-        self.run_script_hook("tcp_close", flow)
-        flow.reply()
+    def handle_tcp_message(self, m):
+        self.run_script_hook("tcp_message", m)
+        m.reply()
 
     def shutdown(self):
         super(FlowMaster, self).shutdown()
diff --git a/mitmproxy/models/__init__.py b/mitmproxy/models/__init__.py
index 3d9d9dae..df86eff4 100644
--- a/mitmproxy/models/__init__.py
+++ b/mitmproxy/models/__init__.py
@@ -7,11 +7,9 @@ from .http import (
 from netlib.http import decoded
 from .connections import ClientConnection, ServerConnection
 from .flow import Flow, Error
-from .tcp import TCPFlow
 
 FLOW_TYPES = dict(
-    http=HTTPFlow,
-    tcp=TCPFlow,
+    http=HTTPFlow
 )
 
 __all__ = [
@@ -20,6 +18,5 @@ __all__ = [
     "make_connect_response", "expect_continue_response",
     "ClientConnection", "ServerConnection",
     "Flow", "Error",
-    "TCPFlow"
     "FLOW_TYPES"
 ]
diff --git a/mitmproxy/models/flow.py b/mitmproxy/models/flow.py
index 1019c9fb..594147ec 100644
--- a/mitmproxy/models/flow.py
+++ b/mitmproxy/models/flow.py
@@ -40,9 +40,6 @@ class Error(stateobject.StateObject):
     def __str__(self):
         return self.msg
 
-    def __repr__(self):
-        return self.msg
-
     @classmethod
     def from_state(cls, state):
         # the default implementation assumes an empty constructor. Override
@@ -102,12 +99,6 @@ class Flow(stateobject.StateObject):
             self._backup = state.pop("backup")
         super(Flow, self).set_state(state)
 
-    @classmethod
-    def from_state(cls, state):
-        f = cls(None, None)
-        f.set_state(state)
-        return f
-
     def copy(self):
         f = copy.copy(self)
 
diff --git a/mitmproxy/models/http.py b/mitmproxy/models/http.py
index 75ffbfd0..77a809cf 100644
--- a/mitmproxy/models/http.py
+++ b/mitmproxy/models/http.py
@@ -191,6 +191,12 @@ class HTTPFlow(Flow):
         response=HTTPResponse
     )
 
+    @classmethod
+    def from_state(cls, state):
+        f = cls(None, None)
+        f.set_state(state)
+        return f
+
     def __repr__(self):
         s = "<HTTPFlow"
         for a in ("request", "response", "error", "client_conn", "server_conn"):
diff --git a/mitmproxy/protocol/rawtcp.py b/mitmproxy/protocol/rawtcp.py
index 5f6fca75..7d18025e 100644
--- a/mitmproxy/protocol/rawtcp.py
+++ b/mitmproxy/protocol/rawtcp.py
@@ -9,26 +9,29 @@ from netlib.exceptions import TcpException
 from netlib.tcp import ssl_read_select
 from netlib.utils import clean_bin
 from ..exceptions import ProtocolException
-from ..models import Error
-from ..models.tcp import TCPFlow, TCPMessage
-
 from .base import Layer
 
 
+class TcpMessage(object):
+
+    def __init__(self, client_conn, server_conn, sender, receiver, message):
+        self.client_conn = client_conn
+        self.server_conn = server_conn
+        self.sender = sender
+        self.receiver = receiver
+        self.message = message
+
+
 class RawTCPLayer(Layer):
     chunk_size = 4096
 
-    def __init__(self, ctx, ignore=False):
-        self.ignore = ignore
+    def __init__(self, ctx, logging=True):
+        self.logging = logging
         super(RawTCPLayer, self).__init__(ctx)
 
     def __call__(self):
         self.connect()
 
-        if not self.ignore:
-            flow = TCPFlow(self.client_conn, self.server_conn, self)
-            self.channel.ask("tcp_open", flow)
-
         buf = memoryview(bytearray(self.chunk_size))
 
         client = self.client_conn.connection
@@ -48,24 +51,38 @@ class RawTCPLayer(Layer):
                         if isinstance(conn, SSL.Connection):
                             # We can't half-close a connection, so we just close everything here.
                             # Sockets will be cleaned up on a higher level.
-                            break
+                            return
                         else:
                             dst.shutdown(socket.SHUT_WR)
 
                         if len(conns) == 0:
-                            break
+                            return
                         continue
 
-                    tcp_message = TCPMessage(dst == server, buf[:size].tobytes())
-                    if not self.ignore:
-                        flow.messages.append(tcp_message)
-                        self.channel.ask("tcp_message", flow)
-                    dst.sendall(tcp_message.content)
+                    tcp_message = TcpMessage(
+                        self.client_conn, self.server_conn,
+                        self.client_conn if dst == server else self.server_conn,
+                        self.server_conn if dst == server else self.client_conn,
+                        buf[:size].tobytes())
+                    self.channel.ask("tcp_message", tcp_message)
+                    dst.sendall(tcp_message.message)
+
+                    if self.logging:
+                        # log messages are prepended with the client address,
+                        # hence the "weird" direction string.
+                        if dst == server:
+                            direction = "-> tcp -> {}".format(repr(self.server_conn.address))
+                        else:
+                            direction = "<- tcp <- {}".format(repr(self.server_conn.address))
+                        data = clean_bin(tcp_message.message)
+                        self.log(
+                            "{}\r\n{}".format(direction, data),
+                            "info"
+                        )
 
         except (socket.error, TcpException, SSL.Error) as e:
-            if not self.ignore:
-                flow.error = Error("TCP connection closed unexpectedly: {}".format(repr(e)))
-                self.channel.tell("tcp_error", flow)
-        finally:
-            if not self.ignore:
-                self.channel.tell("tcp_close", flow)
+            six.reraise(
+                ProtocolException,
+                ProtocolException("TCP connection closed unexpectedly: {}".format(repr(e))),
+                sys.exc_info()[2]
+            )
diff --git a/mitmproxy/proxy/root_context.py b/mitmproxy/proxy/root_context.py
index 21478dfd..9caae02a 100644
--- a/mitmproxy/proxy/root_context.py
+++ b/mitmproxy/proxy/root_context.py
@@ -65,7 +65,7 @@ class RootContext(object):
                 else:
                     ignore = self.config.check_ignore((client_hello.client_sni, 443))
             if ignore:
-                return RawTCPLayer(top_layer, ignore=True)
+                return RawTCPLayer(top_layer, logging=False)
 
         # 2. Always insert a TLS layer, even if there's neither client nor server tls.
         # An inline script may upgrade from http to https,
diff --git a/test/netlib/http/test_request.py b/test/netlib/http/test_request.py
index d57dca13..91fd8ce3 100644
--- a/test/netlib/http/test_request.py
+++ b/test/netlib/http/test_request.py
@@ -108,7 +108,7 @@ class TestRequestUtils(object):
             request.url = "not-a-url"
 
     def test_url_options(self):
-        request = treq(method=b"OPTIONS", path=b"*")
+        request = treq(method="OPTIONS", path="*")
         assert request.url == "http://address:22"
 
     def test_url_authority(self):
@@ -149,7 +149,7 @@ class TestRequestUtils(object):
         assert request.pretty_url == "http://address:22/path"
 
     def test_pretty_url_options(self):
-        request = treq(method=b"OPTIONS", path=b"*")
+        request = treq(method="OPTIONS", path="*")
         assert request.pretty_url == "http://address:22"
 
     def test_pretty_url_authority(self):
-- 
cgit v1.2.3