make eventsequence.iterate extendable

1 files changed, 55 insertions, 39 deletions

diff --git a/mitmproxy/eventsequence.py b/mitmproxy/eventsequence.py
index 4e199972..c2455f24 100644
--- a/mitmproxy/eventsequence.py
+++ b/mitmproxy/eventsequence.py
@@ -1,4 +1,6 @@
-from mitmproxy import controller
+import typing
+
+from mitmproxy import controller, flow
 from mitmproxy import http
 from mitmproxy import tcp
 from mitmproxy import websocket
@@ -8,27 +10,22 @@ Events = frozenset([
     "clientdisconnect",
     "serverconnect",
     "serverdisconnect",
-
     "tcp_start",
     "tcp_message",
     "tcp_error",
     "tcp_end",
-
     "http_connect",
     "request",
     "requestheaders",
     "response",
     "responseheaders",
     "error",
-
     "websocket_handshake",
     "websocket_start",
     "websocket_message",
     "websocket_error",
     "websocket_end",
-
     "next_layer",
-
     "configure",
     "done",
     "log",
@@ -39,37 +36,56 @@ Events = frozenset([
 ])
 
 
-def iterate(f):
-    if isinstance(f, http.HTTPFlow):
-        if f.request:
-            yield "requestheaders", f
-            yield "request", f
-        if f.response:
-            yield "responseheaders", f
-            yield "response", f
-        if f.error:
-            yield "error", f
-    elif isinstance(f, websocket.WebSocketFlow):
-        messages = f.messages
-        f.messages = []
-        f.reply = controller.DummyReply()
-        yield "websocket_start", f
-        while messages:
-            f.messages.append(messages.pop(0))
-            yield "websocket_message", f
-        if f.error:
-            yield "websocket_error", f
-        yield "websocket_end", f
-    elif isinstance(f, tcp.TCPFlow):
-        messages = f.messages
-        f.messages = []
-        f.reply = controller.DummyReply()
-        yield "tcp_start", f
-        while messages:
-            f.messages.append(messages.pop(0))
-            yield "tcp_message", f
-        if f.error:
-            yield "tcp_error", f
-        yield "tcp_end", f
+def _iterate_http(f: http.HTTPFlow):
+    if f.request:
+        yield "requestheaders", f
+        yield "request", f
+    if f.response:
+        yield "responseheaders", f
+        yield "response", f
+    if f.error:
+        yield "error", f
+
+
+def _iterate_websocket(f: websocket.WebSocketFlow):
+    messages = f.messages
+    f.messages = []
+    f.reply = controller.DummyReply()
+    yield "websocket_start", f
+    while messages:
+        f.messages.append(messages.pop(0))
+        yield "websocket_message", f
+    if f.error:
+        yield "websocket_error", f
+    yield "websocket_end", f
+
+
+def _iterate_tcp(f: tcp.TCPFlow):
+    messages = f.messages
+    f.messages = []
+    f.reply = controller.DummyReply()
+    yield "tcp_start", f
+    while messages:
+        f.messages.append(messages.pop(0))
+        yield "tcp_message", f
+    if f.error:
+        yield "tcp_error", f
+    yield "tcp_end", f
+
+
+TEventGenerator = typing.Iterator[typing.Tuple[str, typing.Any]]
+
+_iterate_map: typing.Dict[flow.Flow, typing.Callable[[flow.Flow], TEventGenerator]] = {
+    http.HTTPFlow: _iterate_http,
+    websocket.WebSocketFlow: _iterate_websocket,
+    tcp.TCPFlow: _iterate_tcp
+}
+
+
+def iterate(f: flow.Flow) -> TEventGenerator:
+    try:
+        e = _iterate_map[type(f)]
+    except KeyError as e:
+        raise TypeError("Unknown flow type: {}".format(f)) from e
     else:
-        raise TypeError()
+        yield from e(f)