codekingpro/portable-devtools

1from dataclasses import dataclass2 3from mitmproxy import flow4from mitmproxy import tcp5from mitmproxy.connection import Connection6from mitmproxy.connection import ConnectionState7from mitmproxy.proxy import commands8from mitmproxy.proxy import events9from mitmproxy.proxy import layer10from mitmproxy.proxy.commands import StartHook11from mitmproxy.proxy.context import Context12from mitmproxy.proxy.events import MessageInjected13from mitmproxy.proxy.utils import expect14 15 16@dataclass17class TcpStartHook(StartHook):18    """19    A TCP connection has started.20    """21 22    flow: tcp.TCPFlow23 24 25@dataclass26class TcpMessageHook(StartHook):27    """28    A TCP connection has received a message. The most recent message29    will be flow.messages[-1]. The message is user-modifiable.30    """31 32    flow: tcp.TCPFlow33 34 35@dataclass36class TcpEndHook(StartHook):37    """38    A TCP connection has ended.39    """40 41    flow: tcp.TCPFlow42 43 44@dataclass45class TcpErrorHook(StartHook):46    """47    A TCP error has occurred.48 49    Every TCP flow will receive either a tcp_error or a tcp_end event, but not both.50    """51 52    flow: tcp.TCPFlow53 54 55class TcpMessageInjected(MessageInjected[tcp.TCPMessage]):56    """57    The user has injected a custom TCP message.58    """59 60 61class TCPLayer(layer.Layer):62    """63    Simple TCP layer that just relays messages right now.64    """65 66    flow: tcp.TCPFlow | None67 68    def __init__(self, context: Context, ignore: bool = False):69        super().__init__(context)70        if ignore:71            self.flow = None72        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.done88                return89        self._handle_event = self.relay_messages90 91    _handle_event = start92 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.client99                if event.message.from_client100                else self.context.server,101                event.message.content,102            )103 104        assert isinstance(event, events.ConnectionEvent)105 106        from_client = event.connection == self.context.client107        send_to: Connection108        if from_client:109            send_to = self.context.server110        else:111            send_to = self.context.client112 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)119            else:120                yield commands.SendData(send_to, event.data)121 122        elif isinstance(event, events.ConnectionClosed):123            all_done = not (124                (self.context.client.state & ConnectionState.CAN_READ)125                or (self.context.server.state & ConnectionState.CAN_READ)126            )127            if all_done:128                self._handle_event = self.done129                if self.context.server.state is not ConnectionState.CLOSED:130                    yield commands.CloseConnection(self.context.server)131                if self.context.client.state is not ConnectionState.CLOSED:132                    yield commands.CloseConnection(self.context.client)133                if self.flow:134                    yield TcpEndHook(self.flow)135                    self.flow.live = False136            else:137                yield commands.CloseTcpConnection(send_to, half_close=True)138        else:139            raise AssertionError(f"Unexpected event: {event}")140 141    @expect(events.DataReceived, events.ConnectionClosed, TcpMessageInjected)142    def done(self, _) -> layer.CommandGenerator[None]:143        yield from ()144