# -*- coding: utf-8 -*-
"""Main classes for performing all the SSL connections within the plugins.
"""
from __future__ import absolute_import
from __future__ import unicode_literals

import random
import socket
from typing import Text
import struct
import time
from base64 import b64encode
try:
    # Python 3
    # noinspection PyCompatibility
    from urllib.parse import quote
except ImportError:
    # Python 2
    # noinspection PyCompatibility
    from urllib import quote

from nassl import _nassl
from nassl.debug_ssl_client import DebugSslClient
from nassl.ssl_client import ClientCertificateRequested, OpenSslVerifyEnum
from sslyze.utils.http_request_generator import HttpRequestGenerator

from sslyze.utils.http_response_parser import HttpResponseParser


class SSLHandshakeRejected(IOError):
    """The server explicitly rejected the SSL handshake.
    """
    pass


class StartTLSError(IOError):
    """The server rejected the StartTLS negotiation.
    """
    pass


class ProxyError(IOError):
    """The proxy was offline or did not return HTTP 200 to our CONNECT request.
    """
    pass


class SSLConnection(DebugSslClient):
    """Base SSL connection class.
    """

    # The following errors mean that the server explicitly rejected the handshake. The goal to differentiate rejected
    # handshakes from random network errors such as the server going offline, etc.
    HANDSHAKE_REJECTED_SOCKET_ERRORS = {'was forcibly closed': 'Received FIN',
                                        'reset by peer': 'Received RST'}

    HANDSHAKE_REJECTED_SSL_ERRORS = {'sslv3 alert handshake failure': 'Alert handshake failure',
                                     'no ciphers available': 'No ciphers available',
                                     'excessive message size': 'Excessive message size',
                                     'bad mac decode': 'Bad mac decode',
                                     'wrong version number': 'Wrong version number',
                                     'no cipher match': 'No cipher match',
                                     'bad decompression': 'Bad decompression',
                                     'peer error no cipher': 'Peer error no cipher',
                                     'no cipher list': 'No ciphers list',
                                     'insufficient security': 'Insufficient security',
                                     'block type is not 01': 'block type is not 01',  # Actually an RSA error
                                     'tlsv1 alert protocol version': 'Alert: protocol version '}

    # Constants for tunneling the traffic through a proxy
    HTTP_CONNECT_REQ = 'CONNECT {0}:{1} HTTP/1.1\r\n\r\n'
    HTTP_CONNECT_REQ_PROXY_AUTH_BASIC = 'CONNECT {0}:{1} HTTP/1.1\r\nProxy-Authorization: Basic {2}\r\n\r\n'

    # Errors caused by the proxy
    ERR_CONNECT_REJECTED = 'The proxy rejected the CONNECT request for this host'
    ERR_PROXY_OFFLINE = 'Could not connect to the proxy: "{0}"'

    # Restrict cipher list to make the client hello smaller so we don't run into
    # https://bugs.debian.org/cgi-bin/bugreport.cgi?bug=665452
    DEFAULT_SSL_CIPHER_LIST = 'HIGH:MEDIUM:-aNULL:-eNULL:-3DES:-SRP:-PSK:-CAMELLIA'

    # Default socket settings global to all SSLyze connections; can be overridden
    NETWORK_MAX_RETRIES = 3
    NETWORK_TIMEOUT = 5

    @classmethod
    def set_global_network_settings(cls, network_max_retries, network_timeout):
        cls.NETWORK_MAX_RETRIES = network_max_retries
        cls.NETWORK_TIMEOUT = network_timeout

    def __init__(self, host, ip, port, ssl_version, ssl_verify_locations=None, client_auth_creds=None,
                 should_ignore_client_auth=False):
        if client_auth_creds:
            # A client certificate and private key were provided
            super(SSLConnection, self).__init__(ssl_version=ssl_version,
                                                ssl_verify=OpenSslVerifyEnum.NONE,
                                                ssl_verify_locations=ssl_verify_locations,
                                                client_certchain_file=client_auth_creds.client_certificate_chain_path,
                                                client_key_file=client_auth_creds.client_key_path,
                                                client_key_type=client_auth_creds.client_key_type,
                                                client_key_password=client_auth_creds.client_key_password,
                                                ignore_client_authentication_requests=False)
        else:
            # No client cert and key
            super(SSLConnection, self).__init__(ssl_version=ssl_version,
                                                ssl_verify=OpenSslVerifyEnum.NONE,
                                                ssl_verify_locations=ssl_verify_locations,
                                                ignore_client_authentication_requests=should_ignore_client_auth)

        self._ssl_version = ssl_version
        self._sock = None
        self._host = host
        self._ip = ip
        self._port = port
        self._tunnel_host = None
        self._tunnel_port = None
        self._tunnel_basic_auth_token = None
        self.set_cipher_list(self.DEFAULT_SSL_CIPHER_LIST)

    def enable_http_connect_tunneling(self, tunnel_host, tunnel_port, tunnel_user=None, tunnel_password=None):
        # type: (Text, int, Text, Text) -> None
        """Proxy the traffic through an HTTP Connect proxy.
        """
        self._tunnel_host = tunnel_host
        self._tunnel_port = tunnel_port
        self._tunnel_basic_auth_token = None
        if tunnel_user is not None:
            self._tunnel_basic_auth_token = b64encode(
                '{0}:{1}'.format(quote(tunnel_user), quote(tunnel_password)).encode('utf-8')
            )

    def do_pre_handshake(self, network_timeout):
        # type: (int) -> None
        """Open a socket to the server; setup HTTP tunneling if a proxy was configured.
        """
        if self._tunnel_host:
            # Proxy configured; setup HTTP tunneling
            try:
                self._sock = socket.create_connection((self._tunnel_host, self._tunnel_port), network_timeout)
            except socket.timeout as e:
                raise ProxyError(self.ERR_PROXY_OFFLINE.format(e[0]))
            except socket.error as e:
                raise ProxyError(self.ERR_PROXY_OFFLINE.format(e[1]))

            # Send a CONNECT request with the host we want to tunnel to
            if self._tunnel_basic_auth_token is None:
                self._sock.send(self.HTTP_CONNECT_REQ.format(self._host, self._port).encode('utf-8'))
            else:
                self._sock.send(self.HTTP_CONNECT_REQ_PROXY_AUTH_BASIC.format(
                    self._host, self._port, self._tunnel_basic_auth_token
                ).encode('utf-8'))
            http_response = HttpResponseParser.parse(self._sock)

            # Check if the proxy was able to connect to the host
            if http_response.status != 200:
                raise ProxyError(self.ERR_CONNECT_REJECTED)
        else:
            # No proxy; connect directly to the server
            self._sock = socket.create_connection(address=(self._ip, self._port), timeout=network_timeout)

    def connect(self, network_timeout=None, network_max_retries=None):
        # type: (int, int) -> None
        final_timeout = self.NETWORK_TIMEOUT if network_timeout is None else network_timeout
        final_max_retries = self.NETWORK_MAX_RETRIES if network_max_retries is None else network_max_retries
        retry_attempts = 0
        delay = 0
        while True:
            try:
                # Sleep if it's a retry attempt
                time.sleep(delay)

                # StartTLS negotiation or proxy setup if needed
                self.do_pre_handshake(final_timeout)

                try:
                    # SSL handshake
                    self.do_handshake()

                except ClientCertificateRequested:
                    # Server expected a client certificate and we didn't provide one
                    raise
                except socket.timeout:
                    # Network timeout, propagate the error to trigger a retry
                    raise
                except (socket.error, IOError) as e:
                    # On Python 3.3+ socket.error == IOError but on Python 2.7 they are different
                    # We use the same except block so it works on all versions of Python
                    # This section is meant to handle IOErrors
                    if 'Nassl SSL handshake failed' in str(e.args):
                        raise SSLHandshakeRejected('TLS / Unexpected EOF')

                    # This section is meant to handle socket.errors
                    for error_msg in self.HANDSHAKE_REJECTED_SOCKET_ERRORS.keys():
                        if error_msg in str(e.args):
                            raise SSLHandshakeRejected('TCP / ' + self.HANDSHAKE_REJECTED_SOCKET_ERRORS[error_msg])

                    # Unknown socket error
                    raise
                except _nassl.OpenSSLError as e:
                    for error_msg in self.HANDSHAKE_REJECTED_SSL_ERRORS.keys():
                        if error_msg in str(e.args):
                            raise SSLHandshakeRejected('TLS / ' + self.HANDSHAKE_REJECTED_SSL_ERRORS[error_msg])
                    raise  # Unknown SSL error if we get there

            # Pass on exceptions for rejected handshakes
            except SSLHandshakeRejected:
                raise
            except ClientCertificateRequested:
                raise
            except _nassl.OpenSSLError:
                # Raise unknown OpenSSL errors
                raise
            except socket.timeout:
                # Attempt to retry connection if a network error occurred during connection or the handshake
                retry_attempts += 1
                if retry_attempts >= final_max_retries:
                    # Exhausted the number of retry attempts, give up
                    raise
                elif retry_attempts == 1:
                    delay = random.random()
                else:
                    # Exponential back off
                    delay = min(6, 2 * delay)  # Cap max delay at 6 seconds

            else:
                # No network error occurred
                break

    def close(self):
        # type: () -> None
        self.shutdown()
        if self._sock:
            self._sock.close()

    def post_handshake_check(self):
        # type: () -> Text
        return ''


class HTTPSConnection(SSLConnection):
    """SSL connection class that sends an HTTP GET request after the SSL handshake.
    """

    GET_RESULT_FORMAT = 'HTTP {0} {1}{2}'

    ERR_HTTP_TIMEOUT = 'Timeout on HTTP GET'
    ERR_NOT_HTTP = 'Server response was not HTTP'
    ERR_GENERIC = 'Error sending HTTP GET'


    def post_handshake_check(self):

        try:
            # TODO: This is code only used by OpenSSLCipherSuitesPlugin anf should be moved there
            # Send an HTTP GET to the server and store the HTTP Status Code
            self.write(HttpRequestGenerator.get_request(self._host))

            # Parse the response and print the Location header
            http_response = HttpResponseParser.parse(self)
            if http_response.version == 9:
                # HTTP 0.9 => Probably not an HTTP response
                result = self.ERR_NOT_HTTP
            else:
                redirect = ''
                if 300 <= http_response.status < 400:
                    if http_response.getheader('Location', None):
                        # Add redirection URL to the result
                        redirect = ' - ' + http_response.getheader('Location', None)

                result = self.GET_RESULT_FORMAT.format(http_response.status, http_response.reason, redirect)
        except socket.timeout:
            result = self.ERR_HTTP_TIMEOUT
        except IOError:
            result = self.ERR_GENERIC

        return result


class SMTPConnection(SSLConnection):
    """SSL connection class that performs an SMTP StartTLS negotiation before the SSL handshake and sends a NOOP after
    the handshake.
    """

    ERR_SMTP_REJECTED = 'SMTP EHLO was rejected'
    ERR_NO_SMTP_STARTTLS = 'SMTP STARTTLS not supported'


    def do_pre_handshake(self, network_timeout):
        super(SMTPConnection, self).do_pre_handshake(network_timeout)

        # Get the SMTP banner
        self._sock.recv(2048)

        # Send a EHLO and wait for the 250 status
        self._sock.send(b'EHLO sslyze.scan\r\n')
        if b'250 ' not in self._sock.recv(2048):
            raise StartTLSError(self.ERR_SMTP_REJECTED)

        # Send a STARTTLS
        self._sock.send(b'STARTTLS\r\n')
        if b'220' not in self._sock.recv(2048):
            raise StartTLSError(self.ERR_NO_SMTP_STARTTLS)


    def post_handshake_check(self):
        try:
            self.write(b'NOOP\r\n')
            result = self.read(2048).strip()
        except socket.timeout:
            result = 'Timeout on SMTP NOOP'
        return result


class XMPPConnection(SSLConnection):
    """SSL connection class that performs an XMPP StartTLS negotiation before the SSL handshake.
    """

    ERR_XMPP_REJECTED = 'Error opening XMPP stream, try --xmpp_to'
    ERR_XMPP_HOST_UNKNOWN = 'Error opening XMPP stream: server returned host-unknown error, try --xmpp_to'
    ERR_XMPP_NO_STARTTLS = 'XMPP STARTTLS not supported'

    XMPP_OPEN_STREAM = "<stream:stream xmlns='jabber:client' xmlns:stream='http://etherx.jabber.org/streams' " \
                       "xmlns:tls='http://www.ietf.org/rfc/rfc2595.txt' to='{xmpp_to}' xml:lang='en' version='1.0'>"
    XMPP_STARTTLS = b"<starttls xmlns='urn:ietf:params:xml:ns:xmpp-tls'/>"


    def __init__(self, host, ip, port, ssl_version, ssl_verify_locations=None, client_auth_creds=None,
                 should_ignore_client_auth=False):
        super(XMPPConnection, self).__init__(host, ip, port, ssl_version, ssl_verify_locations, client_auth_creds,
                                             should_ignore_client_auth)
        self._xmpp_to = host


    def set_xmpp_to(self, xmpp_to):
        """XMPP host specified with the XMPP handshake."""
        self._xmpp_to = xmpp_to


    def do_pre_handshake(self, network_timeout):
        """Connect to a host on a given (SSL) port, send a STARTTLS command, and perform the SSL handshake.
        """
        super(XMPPConnection, self).do_pre_handshake(network_timeout)

        # Open an XMPP stream
        self._sock.send(self.XMPP_OPEN_STREAM.format(xmpp_to=self._xmpp_to).encode('utf-8'))

        # Get the server's features and check for an error
        server_resp = self._sock.recv(4096)
        if b'<stream:error>' in server_resp:
            raise StartTLSError(self.ERR_XMPP_REJECTED)
        elif b'</stream:features>' not in server_resp:
            # Get all the server features before initiating startTLS
            self._sock.recv(4096)

        # Send a STARTTLS message
        self._sock.send(self.XMPP_STARTTLS)
        xmpp_resp = self._sock.recv(2048)

        if b'host-unknown' in xmpp_resp:
            raise StartTLSError(self.ERR_XMPP_HOST_UNKNOWN)

        if b'proceed' not in xmpp_resp:
            raise StartTLSError(self.ERR_XMPP_NO_STARTTLS)


class XMPPServerConnection(XMPPConnection):
    XMPP_OPEN_STREAM = "<stream:stream xmlns='jabber:server' xmlns:stream='http://etherx.jabber.org/streams' " \
                       "xmlns:tls='http://www.ietf.org/rfc/rfc2595.txt' to='{xmpp_to}' xml:lang='en' version='1.0'>"


class LDAPConnection(SSLConnection):
    """SSL connection class that performs an LDAP StartTLS negotiation before the SSL handshake.
    """

    ERR_NO_STARTTLS = 'LDAP AUTH TLS was rejected'

    START_TLS_CMD = b'0\x1d\x02\x01\x01w\x18\x80\x161.3.6.1.4.1.1466.20037'
    START_TLS_OK = b'\x30\x0c\x02\x01\x01\x78\x07\x0a\x01\x00\x04\x00\x04'
    START_TLS_OK2 = b'Start TLS request accepted'
    START_TLS_OK_APACHEDS = b'\x30\x26\x02\x01\x01\x78\x21\x0a\x01\x00\x04\x00\x04\x00\x8a\x16\x31\x2e\x33\x2e\x36' \
                            b'\x2e\x31\x2e\x34\x2e\x31\x2e\x31\x34\x36\x36\x2e\x32\x30\x30\x33\x37\x8b\x00'


    def do_pre_handshake(self, network_timeout):
        """Connect to a host on a given (SSL) port, send a STARTTLS command, and perform the SSL handshake.
        """
        super(LDAPConnection, self).do_pre_handshake(network_timeout)

        # Send Start TLS
        self._sock.send(self.START_TLS_CMD)
        data = self._sock.recv(2048)
        if self.START_TLS_OK not in data and self.START_TLS_OK_APACHEDS not in data and self.START_TLS_OK2 not in data:
            raise StartTLSError(self.ERR_NO_STARTTLS + ', returned: "' + data + '" (hex: "' + data.encode('hex') + '")')


class RDPConnection(SSLConnection):
    """SSL connection class that performs an RDP StartTLS negotiation before the SSL handshake.
    """

    ERR_NO_STARTTLS = 'RDP AUTH TLS was rejected'

    START_TLS_CMD = b'\x03\x00\x00\x13\x0E\xE0\x00\x00\x00\x00\x00\x01\x00\x08\x00\x03\x00\x00\x00'
    START_TLS_OK = b'Start TLS request accepted.'

    def do_pre_handshake(self, network_timeout):
        """
        Connect to a host on a given (SSL) port, send a STARTTLS command,
        and perform the SSL handshake.
        """
        super(RDPConnection, self).do_pre_handshake(network_timeout)

        self._sock.send(self.START_TLS_CMD)
        data = self._sock.recv(4)
        if not data or len(data) != 4 or data[:2] != b'\x03\x00':
            raise StartTLSError(self.ERR_NO_STARTTLS)
        packet_len = struct.unpack(">H", data[2:])[0] - 4
        data = self._sock.recv(packet_len)

        if not data or len(data) != packet_len:
            raise StartTLSError(self.ERR_NO_STARTTLS)


class GenericStartTLSConnection(SSLConnection):
    """SSL connection class that performs a StartTLS negotiation before the SSL handshake.
    """

    # To be defined in subclasses
    ERR_NO_STARTTLS = b''
    START_TLS_CMD = b''
    START_TLS_OK = b''
    SHOULD_WAIT_FOR_SERVER_BANNER = True

    def do_pre_handshake(self, network_timeout):
        """Connect to a host on a given (SSL) port, send a STARTTLS command, and perform the SSL handshake.
        """
        super(GenericStartTLSConnection, self).do_pre_handshake(network_timeout)

        # Grab the banner
        if self.SHOULD_WAIT_FOR_SERVER_BANNER:
            self._sock.recv(2048)

        # Send Start TLS
        self._sock.send(self.START_TLS_CMD)
        if self.START_TLS_OK not in self._sock.recv(2048):
            raise StartTLSError(self.ERR_NO_STARTTLS)


class IMAPConnection(GenericStartTLSConnection):
    """SSL connection class that performs an IMAP StartTLS negotiation before the SSL handshake.
    """

    ERR_NO_STARTTLS = b'IMAP START TLS was rejected'

    START_TLS_CMD = b'. STARTTLS\r\n'
    START_TLS_OK = b'. OK'


class POP3Connection(GenericStartTLSConnection):
    """SSL connection class that performs a POP3 StartTLS negotiation before the SSL handshake.
    """

    ERR_NO_STARTTLS = b'POP START TLS was rejected'

    START_TLS_CMD = b'STLS\r\n'
    START_TLS_OK = b'+OK'


class FTPConnection(GenericStartTLSConnection):
    """SSL connection class that performs an FTP StartTLS negotiation before the SSL handshake.
    """

    ERR_NO_STARTTLS = b'FTP AUTH TLS was rejected'

    START_TLS_CMD = b'AUTH TLS\r\n'
    START_TLS_OK = b'234'


class PostgresConnection(GenericStartTLSConnection):
    """PostgreSQL SSL Connection.
    """

    ERR_NO_STARTTLS = b'Postgres AUTH TLS was rejected'

    START_TLS_CMD = b'\x00\x00\x00\x08\x04\xD2\x16\x2F'
    START_TLS_OK = b'S'
    SHOULD_WAIT_FOR_SERVER_BANNER = False
