aboutsummaryrefslogtreecommitdiffstats
path: root/mitmproxy/proxy/protocol/base.py
blob: f1b8da6cb978d15df19f9825eb0cc149136f5fd0 (plain)
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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
from mitmproxy import exceptions
from mitmproxy import connections
from mitmproxy import controller  # noqa
from mitmproxy.proxy import config  # noqa


class _LayerCodeCompletion:

    """
    Dummy class that provides type hinting in PyCharm, which simplifies development a lot.
    """

    def __init__(self, **mixin_args):  # pragma: no cover
        super().__init__(**mixin_args)
        if True:
            return
        self.config = None  # type: config.ProxyConfig
        self.client_conn = None  # type: connections.ClientConnection
        self.server_conn = None  # type: connections.ServerConnection
        self.channel = None  # type: controller.Channel
        self.ctx = None
        """@type: mitmproxy.proxy.protocol.Layer"""


class Layer(_LayerCodeCompletion):

    """
    Base class for all layers. All other protocol layers should inherit from this class.
    """

    def __init__(self, ctx, **mixin_args):
        """
        Each layer usually passes itself to its child layers as a context. Properties of the
        context are transparently mapped to the layer, so that the following works:

        .. code-block:: python

            root_layer = Layer(None)
            root_layer.client_conn = 42
            sub_layer = Layer(root_layer)
            print(sub_layer.client_conn) # 42

        The root layer is passed a :py:class:`mitmproxy.proxy.RootContext` object,
        which provides access to :py:attr:`.client_conn <mitmproxy.proxy.RootContext.client_conn>`,
        :py:attr:`.next_layer <mitmproxy.proxy.RootContext.next_layer>` and other basic attributes.

        Args:
            ctx: The (read-only) parent layer / context.
        """
        self.ctx = ctx
        """
        The parent layer.

        :type: :py:class:`Layer`
        """
        super().__init__(**mixin_args)

    def __call__(self):
        """Logic of the layer.

        Returns:
            Once the protocol has finished without exceptions.

        Raises:
            ~mitmproxy.exceptions.ProtocolException: if an exception occurs. No other exceptions must be raised.
        """
        raise NotImplementedError()

    def __getattr__(self, name):
        """
        Attributes not present on the current layer are looked up on the context.
        """
        return getattr(self.ctx, name)


class ServerConnectionMixin:

    """
    Mixin that provides a layer with the capabilities to manage a server connection.
    The server address can be passed in the constructor or set by calling :py:meth:`set_server`.
    Subclasses are responsible for calling :py:meth:`disconnect` before returning.

    Recommended Usage:

    .. code-block:: python

        class MyLayer(Layer, ServerConnectionMixin):
            def __call__(self):
                try:
                    # Do something.
                finally:
                    if self.server_conn.connected():
                        self.disconnect()
    """

    def __init__(self, server_address=None):
        super().__init__()

        self.server_conn = self.__make_server_conn(server_address)

        self.__check_self_connect()

    def __check_self_connect(self):
        """
        We try to protect the proxy from _accidentally_ connecting to itself,
        e.g. because of a failed transparent lookup or an invalid configuration.
        """
        address = self.server_conn.address
        if address:
            self_connect = (
                address[1] == self.config.options.listen_port and
                address[0] in ("localhost", "127.0.0.1", "::1")
            )
            if self_connect:
                raise exceptions.ProtocolException(
                    "Invalid server address: {}\r\n"
                    "The proxy shall not connect to itself.".format(repr(address))
                )

    def __make_server_conn(self, server_address):
        if self.config.options.spoof_source_address and self.config.options.upstream_bind_address == '':
            return connections.ServerConnection(
                server_address, (self.ctx.client_conn.address[0], 0), True)
        else:
            return connections.ServerConnection(
                server_address, (self.config.options.upstream_bind_address, 0),
                self.config.options.spoof_source_address
            )

    def set_server(self, address):
        """
        Sets a new server address. If there is an existing connection, it will be closed.
        """
        if self.server_conn.connected():
            self.disconnect()
        self.log("Set new server address: {}:{}".format(address[0], address[1]), "debug")
        self.server_conn.address = address
        self.__check_self_connect()

    def disconnect(self):
        """
        Deletes (and closes) an existing server connection.
        Must not be called if there is no existing connection.
        """
        self.log("serverdisconnect", "debug", [repr(self.server_conn.address)])
        address = self.server_conn.address
        self.server_conn.finish()
        self.server_conn.close()
        self.channel.tell("serverdisconnect", self.server_conn)

        self.server_conn = self.__make_server_conn(address)

    def connect(self):
        """
        Establishes a server connection.
        Must not be called if there is an existing connection.

        Raises:
            ~mitmproxy.exceptions.ProtocolException: if the connection could not be established.
        """
        if not self.server_conn.address:
            raise exceptions.ProtocolException("Cannot connect to server, no server address given.")
        self.log("serverconnect", "debug", [repr(self.server_conn.address)])
        self.channel.ask("serverconnect", self.server_conn)
        try:
            self.server_conn.connect()
        except exceptions.TcpException as e:
            raise exceptions.ProtocolException(
                "Server connection to {} failed: {}".format(
                    repr(self.server_conn.address), str(e)
                )
            )