Simple TCP layer that just relays messages right now.
| 59 | |
| 60 | |
| 61 | class TCPLayer(layer.Layer): |
| 62 | """ |
| 63 | Simple TCP layer that just relays messages right now. |
| 64 | """ |
| 65 | |
| 66 | flow: tcp.TCPFlow | None |
| 67 | |
| 68 | def __init__(self, context: Context, ignore: bool = False): |
| 69 | super().__init__(context) |
| 70 | if ignore: |
| 71 | self.flow = None |
| 72 | else: |
| 73 | self.flow = tcp.TCPFlow(self.context.client, self.context.server, True) |
| 74 | |
| 75 | @expect(events.Start) |
| 76 | def start(self, _) -> layer.CommandGenerator[None]: |
| 77 | if self.flow: |
| 78 | yield TcpStartHook(self.flow) |
| 79 | |
| 80 | if self.context.server.timestamp_start is None: |
| 81 | err = yield commands.OpenConnection(self.context.server) |
| 82 | if err: |
| 83 | if self.flow: |
| 84 | self.flow.error = flow.Error(str(err)) |
| 85 | yield TcpErrorHook(self.flow) |
| 86 | yield commands.CloseConnection(self.context.client) |
| 87 | self._handle_event = self.done |
| 88 | return |
| 89 | self._handle_event = self.relay_messages |
| 90 | |
| 91 | _handle_event = start |
| 92 | |
| 93 | @expect(events.DataReceived, events.ConnectionClosed, TcpMessageInjected) |
| 94 | def relay_messages(self, event: events.Event) -> layer.CommandGenerator[None]: |
| 95 | if isinstance(event, TcpMessageInjected): |
| 96 | # we just spoof that we received data here and then process that regularly. |
| 97 | event = events.DataReceived( |
| 98 | self.context.client |
| 99 | if event.message.from_client |
| 100 | else self.context.server, |
| 101 | event.message.content, |
| 102 | ) |
| 103 | |
| 104 | assert isinstance(event, events.ConnectionEvent) |
| 105 | |
| 106 | from_client = event.connection == self.context.client |
| 107 | send_to: Connection |
| 108 | if from_client: |
| 109 | send_to = self.context.server |
| 110 | else: |
| 111 | send_to = self.context.client |
| 112 | |
| 113 | if isinstance(event, events.DataReceived): |
| 114 | if self.flow: |
| 115 | tcp_message = tcp.TCPMessage(from_client, event.data) |
| 116 | self.flow.messages.append(tcp_message) |
| 117 | yield TcpMessageHook(self.flow) |
| 118 | yield commands.SendData(send_to, tcp_message.content) |
no outgoing calls
searching dependent graphs…