From 317b4beccaa7e30f11095cfaf0364e325a875fd2 Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Fri, 29 Apr 2016 20:18:32 -0700
Subject: mitmdump: don't fail for other flowtypes

---
 mitmproxy/flow.py            | 39 +++++++++++++++++++++------------------
 mitmproxy/models/__init__.py |  5 +++++
 test/mitmproxy/test_flow.py  |  2 +-
 3 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py
index 2bed4e1a..555c6895 100644
--- a/mitmproxy/flow.py
+++ b/mitmproxy/flow.py
@@ -21,7 +21,7 @@ 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
+from .models import ClientConnection, ServerConnection, HTTPFlow, HTTPRequest, FLOW_TYPES
 from collections import defaultdict
 
 
@@ -873,23 +873,24 @@ class FlowMaster(controller.ServerMaster):
 
     def load_flow(self, f):
         """
-            Loads a flow, and returns a new flow object.
+        Loads a flow
         """
-
-        if self.server and self.server.config.mode == "reverse":
-            f.request.host = self.server.config.upstream_server.address.host
-            f.request.port = self.server.config.upstream_server.address.port
-            f.request.scheme = self.server.config.upstream_server.scheme
-
-        f.reply = controller.DummyReply()
-        if f.request:
-            self.handle_request(f)
-        if f.response:
-            self.handle_responseheaders(f)
-            self.handle_response(f)
-        if f.error:
-            self.handle_error(f)
-        return f
+        if isinstance(f, HTTPFlow):
+            if self.server and self.server.config.mode == "reverse":
+                f.request.host = self.server.config.upstream_server.address.host
+                f.request.port = self.server.config.upstream_server.address.port
+                f.request.scheme = self.server.config.upstream_server.scheme
+
+            f.reply = controller.DummyReply()
+            if f.request:
+                self.handle_request(f)
+            if f.response:
+                self.handle_responseheaders(f)
+                self.handle_response(f)
+            if f.error:
+                self.handle_error(f)
+        else:
+            raise NotImplementedError()
 
     def load_flows(self, fr):
         """
@@ -1166,7 +1167,9 @@ class FlowReader:
                     raise FlowReadException(str(e))
                 if can_tell:
                     off = self.fo.tell()
-                yield HTTPFlow.from_state(data)
+                if data["type"] not in FLOW_TYPES:
+                    raise FlowReadException("Unknown flow type: {}".format(data["type"]))
+                yield FLOW_TYPES[data["type"]].from_state(data)
         except ValueError:
             # Error is due to EOF
             if can_tell and self.fo.tell() == off and self.fo.read() == '':
diff --git a/mitmproxy/models/__init__.py b/mitmproxy/models/__init__.py
index f5f0213a..df86eff4 100644
--- a/mitmproxy/models/__init__.py
+++ b/mitmproxy/models/__init__.py
@@ -8,10 +8,15 @@ from netlib.http import decoded
 from .connections import ClientConnection, ServerConnection
 from .flow import Flow, Error
 
+FLOW_TYPES = dict(
+    http=HTTPFlow
+)
+
 __all__ = [
     "HTTPFlow", "HTTPRequest", "HTTPResponse", "Headers", "decoded",
     "make_error_response", "make_connect_request",
     "make_connect_response", "expect_continue_response",
     "ClientConnection", "ServerConnection",
     "Flow", "Error",
+    "FLOW_TYPES"
 ]
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index 1d69a3f8..145e91cf 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -818,7 +818,7 @@ class TestFlowMaster:
         s = flow.State()
         fm = flow.FlowMaster(None, s)
         f = tutils.tflow(resp=True)
-        f = fm.load_flow(f)
+        fm.load_flow(f)
         assert s.flow_count() == 1
         f2 = fm.duplicate_flow(f)
         assert f2.response
-- 
cgit v1.2.3