/
tcp.py
86 lines (67 loc) · 2.64 KB
/
tcp.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
from typing import Optional
from mitmproxy import flow, tcp
from mitmproxy.proxy2 import commands, events
from mitmproxy.proxy2.commands import Hook
from mitmproxy.proxy2.context import Context
from mitmproxy.proxy2.layer import Layer
from mitmproxy.proxy2.utils import expect
class TcpStartHook(Hook):
flow: tcp.TCPFlow
class TcpMessageHook(Hook):
flow: tcp.TCPFlow
class TcpEndHook(Hook):
flow: tcp.TCPFlow
class TcpErrorHook(Hook):
flow: tcp.TCPFlow
class TCPLayer(Layer):
"""
Simple TCP layer that just relays messages right now.
"""
context: Context
flow: Optional[tcp.TCPFlow]
def __init__(self, context: Context, ignore: bool = False):
super().__init__(context)
if ignore:
self.flow = None
else:
self.flow = tcp.TCPFlow(self.context.client, self.context.server, True)
@expect(events.Start)
def start(self, _) -> commands.TCommandGenerator:
if self.flow:
yield TcpStartHook(self.flow)
if not self.context.server.connected:
err = yield commands.OpenConnection(self.context.server)
if err:
if self.flow:
self.flow.error = flow.Error(str(err))
yield TcpErrorHook(self.flow)
yield commands.CloseConnection(self.context.client)
self._handle_event = self.done
return
self._handle_event = self.relay_messages
_handle_event = start
@expect(events.DataReceived, events.ConnectionClosed)
def relay_messages(self, event: events.ConnectionEvent) -> commands.TCommandGenerator:
from_client = event.connection == self.context.client
if from_client:
send_to = self.context.server
else:
send_to = self.context.client
if isinstance(event, events.DataReceived):
if self.flow:
tcp_message = tcp.TCPMessage(from_client, event.data)
self.flow.messages.append(tcp_message)
yield TcpMessageHook(self.flow)
yield commands.SendData(send_to, tcp_message.content)
else:
yield commands.SendData(send_to, event.data)
elif isinstance(event, events.ConnectionClosed):
yield commands.CloseConnection(send_to)
all_done = (not self.context.client.connected and not self.context.server.connected)
if all_done:
self._handle_event = self.done
if self.flow:
yield TcpEndHook(self.flow)
@expect(events.DataReceived, events.ConnectionClosed)
def done(self, _):
yield from ()