aboutsummaryrefslogtreecommitdiffstats
path: root/libmproxy/protocol/base.py
blob: 40fcaf65fcc087d02208399c1291022fc80d9545 (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
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
from __future__ import (absolute_import, print_function, division)
import sys

import six

from ..models import ServerConnection
from ..exceptions import ProtocolException
from netlib.exceptions import TcpException


class _LayerCodeCompletion(object):

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

    def __init__(self, **mixin_args):  # pragma: no cover
        super(_LayerCodeCompletion, self).__init__(**mixin_args)
        if True:
            return
        self.config = None
        """@type: libmproxy.proxy.ProxyConfig"""
        self.client_conn = None
        """@type: libmproxy.models.ClientConnection"""
        self.server_conn = None
        """@type: libmproxy.models.ServerConnection"""
        self.channel = None
        """@type: libmproxy.controller.Channel"""
        self.ctx = None
        """@type: libmproxy.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:`libmproxy.proxy.RootContext` object,
        which provides access to :py:attr:`.client_conn <libmproxy.proxy.RootContext.client_conn>`,
        :py:attr:`.next_layer <libmproxy.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(Layer, self).__init__(**mixin_args)

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

        Returns:
            Once the protocol has finished without exceptions.

        Raises:
            ~libmproxy.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)

    @property
    def layers(self):
        """
        List of all layers, including the current layer (``[self, self.ctx, self.ctx.ctx, ...]``)
        """
        return [self] + self.ctx.layers

    def __repr__(self):
        return type(self).__name__


class ServerConnectionMixin(object):

    """
    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:
                        self.disconnect()
    """

    def __init__(self, server_address=None):
        super(ServerConnectionMixin, self).__init__()
        self.server_conn = ServerConnection(server_address, (self.config.host, 0))
        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.port == self.config.port and
                address.host in ("localhost", "127.0.0.1", "::1")
            )
            if self_connect:
                raise ProtocolException(
                    "Invalid server address: {}\r\n"
                    "The proxy shall not connect to itself.".format(repr(address))
                )

    def set_server(self, address, server_tls=None, sni=None):
        """
        Sets a new server address. If there is an existing connection, it will be closed.

        Raises:
            ~libmproxy.exceptions.ProtocolException:
                if ``server_tls`` is ``True``, but there was no TLS layer on the
                protocol stack which could have processed this.
        """
        if self.server_conn:
            self.disconnect()
        self.log("Set new server address: " + repr(address), "debug")
        self.server_conn.address = address
        self.__check_self_connect()
        if server_tls:
            raise ProtocolException(
                "Cannot upgrade to TLS, no TLS layer on the protocol stack."
            )

    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
        source_address = self.server_conn.source_address
        self.server_conn.finish()
        self.server_conn.close()
        self.channel.tell("serverdisconnect", self.server_conn)
        self.server_conn = ServerConnection(address, source_address)

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

        Raises:
            ~libmproxy.exceptions.ProtocolException: if the connection could not be established.
        """
        if not self.server_conn.address:
            raise 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 TcpException as e:
            six.reraise(
                ProtocolException,
                ProtocolException(
                    "Server connection to {} failed: {}".format(
                        repr(self.server_conn.address), str(e)
                    )
                ),
                sys.exc_info()[2]
            )


class Kill(Exception):

    """
    Signal that both client and server connection(s) should be killed immediately.
    """