Merge pull request #1127 from mitmproxy/tcp-flows

mitmdump: Add Basic Support for TCP Flows

1 files changed, 53 insertions, 10 deletions

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()