Merge pull request #1127 from mitmproxy/tcp-flows

mitmdump: Add Basic Support for TCP Flows

13 files changed, 180 insertions, 68 deletions

diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py
index aae397cd..f1eabdb8 100644
--- a/mitmproxy/dump.py
+++ b/mitmproxy/dump.py
@@ -320,7 +320,6 @@ class DumpMaster(flow.FlowMaster):
             self.outfile.flush()
 
     def _process_flow(self, f):
-        self.state.delete_flow(f)
         if self.filt and not f.match(self.filt):
             return
 
@@ -328,6 +327,7 @@ class DumpMaster(flow.FlowMaster):
 
     def handle_request(self, f):
         flow.FlowMaster.handle_request(self, f)
+        self.state.delete_flow(f)
         if f:
             f.reply()
         return f
diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py
index a9018e16..a09a81a7 100644
--- a/mitmproxy/flow.py
+++ b/mitmproxy/flow.py
@@ -3,7 +3,6 @@
 """
 from __future__ import absolute_import
 
-import traceback
 from abc import abstractmethod, ABCMeta
 import hashlib
 import sys
@@ -18,12 +17,13 @@ from typing import List, Optional, Set
 from netlib import wsgi, odict
 from netlib.exceptions import HttpException
 from netlib.http import Headers, http1, cookies
+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
+from .models import ClientConnection, ServerConnection, HTTPFlow, HTTPRequest, FLOW_TYPES, TCPFlow
 from collections import defaultdict
 
 
@@ -651,8 +651,9 @@ class FlowMaster(controller.ServerMaster):
         if server:
             self.add_server(server)
         self.state = state
-        self.server_playback = None
-        self.client_playback = None
+        self.active_flows = set()  # type: Set[Flow]
+        self.server_playback = None  # type: Optional[ServerPlaybackState]
+        self.client_playback = None  # type: Optional[ClientPlaybackState]
         self.kill_nonreplay = False
         self.scripts = []  # type: List[script.Script]
         self.pause_scripts = False
@@ -898,6 +899,17 @@ 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()
 
@@ -1020,6 +1032,7 @@ class FlowMaster(controller.ServerMaster):
                 return
         if f not in self.state.flows:  # don't add again on replay
             self.state.add_flow(f)
+        self.active_flows.add(f)
         self.replacehooks.run(f)
         self.setheaders.run(f)
         self.process_new_request(f)
@@ -1040,6 +1053,7 @@ class FlowMaster(controller.ServerMaster):
         return f
 
     def handle_response(self, f):
+        self.active_flows.discard(f)
         self.state.update_flow(f)
         self.replacehooks.run(f)
         self.setheaders.run(f)
@@ -1085,18 +1099,47 @@ class FlowMaster(controller.ServerMaster):
             self.add_event('"{}" reloaded.'.format(s.filename), 'info')
         return ok
 
-    def handle_tcp_message(self, m):
-        self.run_script_hook("tcp_message", m)
-        m.reply()
+    def handle_tcp_open(self, flow):
+        # TODO: This would break mitmproxy currently.
+        # self.state.add_flow(flow)
+        self.active_flows.add(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):
+        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.active_flows.discard(flow)
+        if self.stream:
+            self.stream.add(flow)
+        self.run_script_hook("tcp_close", flow)
+        flow.reply()
 
     def shutdown(self):
         super(FlowMaster, self).shutdown()
 
         # Add all flows that are still active
         if self.stream:
-            for i in self.state.flows:
-                if not i.response:
-                    self.stream.add(i)
+            for flow in self.active_flows:
+                self.stream.add(flow)
             self.stop_stream()
 
         self.unload_scripts()
diff --git a/mitmproxy/models/__init__.py b/mitmproxy/models/__init__.py
index df86eff4..3d9d9dae 100644
--- a/mitmproxy/models/__init__.py
+++ b/mitmproxy/models/__init__.py
@@ -7,9 +7,11 @@ 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
+    http=HTTPFlow,
+    tcp=TCPFlow,
 )
 
 __all__ = [
@@ -18,5 +20,6 @@ __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 594147ec..1019c9fb 100644
--- a/mitmproxy/models/flow.py
+++ b/mitmproxy/models/flow.py
@@ -40,6 +40,9 @@ 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
@@ -99,6 +102,12 @@ 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 77a809cf..75ffbfd0 100644
--- a/mitmproxy/models/http.py
+++ b/mitmproxy/models/http.py
@@ -191,12 +191,6 @@ 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/models/tcp.py b/mitmproxy/models/tcp.py
new file mode 100644
index 00000000..7e966b95
--- /dev/null
+++ b/mitmproxy/models/tcp.py
@@ -0,0 +1,50 @@
+import time
+from typing import List
+
+from netlib.utils import Serializable
+from .flow import Flow
+
+
+class TCPMessage(Serializable):
+    def __init__(self, from_client, content, timestamp=None):
+        self.content = content
+        self.from_client = from_client
+        if timestamp is None:
+            timestamp = time.time()
+        self.timestamp = timestamp
+
+    @classmethod
+    def from_state(cls, state):
+        return cls(*state)
+
+    def get_state(self):
+        return self.from_client, self.content, self.timestamp
+
+    def set_state(self, state):
+        self.from_client = state.pop("from_client")
+        self.content = state.pop("content")
+        self.timestamp = state.pop("timestamp")
+
+    def __repr__(self):
+        return "{direction} {content}".format(
+            direction="->" if self.from_client else "<-",
+            content=repr(self.content)
+        )
+
+
+class TCPFlow(Flow):
+    """
+    A TCPFlow is a simplified representation of a TCP session.
+    """
+
+    def __init__(self, client_conn, server_conn, live=None):
+        super(TCPFlow, self).__init__("tcp", client_conn, server_conn, live)
+        self.messages = []  # type: List[TCPMessage]
+
+    _stateobject_attributes = Flow._stateobject_attributes.copy()
+    _stateobject_attributes.update(
+        messages=List[TCPMessage]
+    )
+
+    def __repr__(self):
+        return "<TCPFlow ({} messages)>".format(len(self.messages))
diff --git a/mitmproxy/protocol/rawtcp.py b/mitmproxy/protocol/rawtcp.py
index 7d18025e..1b546c40 100644
--- a/mitmproxy/protocol/rawtcp.py
+++ b/mitmproxy/protocol/rawtcp.py
@@ -9,29 +9,26 @@ from netlib.exceptions import TcpException
 from netlib.tcp import ssl_read_select
 from netlib.utils import clean_bin
 from ..exceptions import ProtocolException
-from .base import Layer
-
+from ..models import Error
+from ..models.tcp import TCPFlow, TCPMessage
 
-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
+from .base import Layer
 
 
 class RawTCPLayer(Layer):
     chunk_size = 4096
 
-    def __init__(self, ctx, logging=True):
-        self.logging = logging
+    def __init__(self, ctx, ignore=False):
+        self.ignore = ignore
         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
@@ -59,30 +56,16 @@ class RawTCPLayer(Layer):
                             return
                         continue
 
-                    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"
-                        )
+                    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)
 
         except (socket.error, TcpException, SSL.Error) as e:
-            six.reraise(
-                ProtocolException,
-                ProtocolException("TCP connection closed unexpectedly: {}".format(repr(e))),
-                sys.exc_info()[2]
-            )
+            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)
diff --git a/mitmproxy/proxy/root_context.py b/mitmproxy/proxy/root_context.py
index c55105ec..96e7aab6 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.sni, 443))
             if ignore:
-                return RawTCPLayer(top_layer, logging=False)
+                return RawTCPLayer(top_layer, ignore=True)
 
         # 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/mitmproxy/scripts/tcp_stream_modify.py b/test/mitmproxy/scripts/tcp_stream_modify.py
index 93b0d5c8..d7953ef9 100644
--- a/test/mitmproxy/scripts/tcp_stream_modify.py
+++ b/test/mitmproxy/scripts/tcp_stream_modify.py
@@ -1,3 +1,4 @@
-def tcp_message(ctx, tm):
-    if tm.sender == tm.server_conn:
-        tm.message = tm.message.replace("foo", "bar")
+def tcp_message(ctx, flow):
+    message = flow.messages[-1]
+    if not message.from_client:
+        message.content = message.content.replace("foo", "bar")
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index bf417423..3e78a5c4 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -680,6 +680,10 @@ class TestSerialize:
         for i in range(3):
             f = tutils.tflow(err=True)
             w.add(f)
+        f = tutils.ttcpflow()
+        w.add(f)
+        f = tutils.ttcpflow(err=True)
+        w.add(f)
 
         sio.seek(0)
         return flow.FlowReader(sio)
@@ -1151,6 +1155,10 @@ class TestError:
         e3 = e.copy()
         assert e3.get_state() == e.get_state()
 
+    def test_repr(self):
+        e = Error("yay")
+        assert repr(e)
+
 
 class TestClientConnection:
 
diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py
index 454736d4..0701d52b 100644
--- a/test/mitmproxy/test_server.py
+++ b/test/mitmproxy/test_server.py
@@ -14,7 +14,7 @@ from pathod import pathoc, pathod
 
 from mitmproxy.proxy.config import HostMatcher
 from mitmproxy.exceptions import Kill
-from mitmproxy.models import Error, HTTPResponse
+from mitmproxy.models import Error, HTTPResponse, HTTPFlow
 
 from . import tutils, tservers
 
@@ -177,9 +177,9 @@ class TcpMixin:
         assert n.status_code == 304
         assert i.status_code == 305
         assert i2.status_code == 306
-        assert any(f.response.status_code == 304 for f in self.master.state.flows)
-        assert not any(f.response.status_code == 305 for f in self.master.state.flows)
-        assert not any(f.response.status_code == 306 for f in self.master.state.flows)
+        assert any(f.response.status_code == 304 for f in self.master.state.flows if isinstance(f, HTTPFlow))
+        assert not any(f.response.status_code == 305 for f in self.master.state.flows if isinstance(f, HTTPFlow))
+        assert not any(f.response.status_code == 306 for f in self.master.state.flows if isinstance(f, HTTPFlow))
 
         # Test that we get the original SSL cert
         if self.ssl:
diff --git a/test/mitmproxy/tservers.py b/test/mitmproxy/tservers.py
index 4fa519cc..c9d68cfd 100644
--- a/test/mitmproxy/tservers.py
+++ b/test/mitmproxy/tservers.py
@@ -50,9 +50,8 @@ class TestMaster(flow.FlowMaster):
     def clear_log(self):
         self.log = []
 
-    def handle_log(self, l):
-        self.log.append(l.msg)
-        l.reply()
+    def add_event(self, message, level=None):
+        self.log.append(message)
 
 
 class ProxyThread(threading.Thread):
diff --git a/test/mitmproxy/tutils.py b/test/mitmproxy/tutils.py
index 2dfd710e..118f849c 100644
--- a/test/mitmproxy/tutils.py
+++ b/test/mitmproxy/tutils.py
@@ -3,6 +3,8 @@ import shutil
 import tempfile
 import argparse
 import sys
+
+from mitmproxy.models.tcp import TCPMessage
 from six.moves import cStringIO as StringIO
 from contextlib import contextmanager
 
@@ -12,7 +14,7 @@ import netlib.utils
 import netlib.tutils
 from mitmproxy import utils, controller
 from mitmproxy.models import (
-    ClientConnection, ServerConnection, Error, HTTPRequest, HTTPResponse, HTTPFlow
+    ClientConnection, ServerConnection, Error, HTTPRequest, HTTPResponse, HTTPFlow, TCPFlow
 )
 
 
@@ -45,6 +47,26 @@ def skip_appveyor(fn):
         return fn
 
 
+def ttcpflow(client_conn=True, server_conn=True, messages=True, err=None):
+    if client_conn is True:
+        client_conn = tclient_conn()
+    if server_conn is True:
+        server_conn = tserver_conn()
+    if messages is True:
+        messages = [
+            TCPMessage(True, b"hello"),
+            TCPMessage(False, b"it's me"),
+        ]
+    if err is True:
+        err = terr()
+
+    f = TCPFlow(client_conn, server_conn)
+    f.messages = messages
+    f.error = err
+    f.reply = controller.DummyReply()
+    return f
+
+
 def tflow(client_conn=True, server_conn=True, req=True, resp=None, err=None):
     """
     @type client_conn: bool | None | mitmproxy.proxy.connection.ClientConnection
@@ -52,7 +74,7 @@ def tflow(client_conn=True, server_conn=True, req=True, resp=None, err=None):
     @type req:         bool | None | mitmproxy.protocol.http.HTTPRequest
     @type resp:        bool | None | mitmproxy.protocol.http.HTTPResponse
     @type err:         bool | None | mitmproxy.protocol.primitives.Error
-    @return:           bool | None | mitmproxy.protocol.http.HTTPFlow
+    @return:           mitmproxy.protocol.http.HTTPFlow
     """
     if client_conn is True:
         client_conn = tclient_conn()