mitmproxy/libmproxy/protocol/base.py

from __future__ import (absolute_import, print_function, division)
import sys

import six

from netlib import tcp
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: nocover
        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.
    """