Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-26 10:16:27 +00:00
Code Issues Projects Releases Wiki Activity

make eventsequence.iterate extendable

Browse Source
This commit is contained in:
Maximilian Hils 2017-07-15 12:22:46 +02:00
parent af777f9b4f
commit a81cd18d92
1 changed files with 55 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

94
mitmproxy/eventsequence.py
View File

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