# Copyright DataStax, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import socket
import logging

try:
    import kerberos
    _have_kerberos = True
except ImportError:
    _have_kerberos = False

try:
    from puresasl.client import SASLClient
    _have_puresasl = True
except ImportError:
    _have_puresasl = False

try:
    from puresasl.client import SASLClient
except ImportError:
    SASLClient = None

log = logging.getLogger(__name__)

# Custom payload keys related to DSE Unified Auth
_proxy_execute_key = 'ProxyExecute'


class AuthProvider(object):
    """
    An abstract class that defines the interface that will be used for
    creating :class:`~.Authenticator` instances when opening new
    connections to Cassandra.

    .. versionadded:: 2.0.0
    """

    def new_authenticator(self, host):
        """
        Implementations of this class should return a new instance
        of :class:`~.Authenticator` or one of its subclasses.
        """
        raise NotImplementedError()


class Authenticator(object):
    """
    An abstract class that handles SASL authentication with Cassandra servers.

    Each time a new connection is created and the server requires authentication,
    a new instance of this class will be created by the corresponding
    :class:`~.AuthProvider` to handler that authentication. The lifecycle of the
    new :class:`~.Authenticator` will the be:

    1) The :meth:`~.initial_response()` method will be called. The return
    value will be sent to the server to initiate the handshake.

    2) The server will respond to each client response by either issuing a
    challenge or indicating that the authentication is complete (successful or not).
    If a new challenge is issued, :meth:`~.evaluate_challenge()`
    will be called to produce a response that will be sent to the
    server. This challenge/response negotiation will continue until the server
    responds that authentication is successful (or an :exc:`~.AuthenticationFailed`
    is raised).

    3) When the server indicates that authentication is successful,
    :meth:`~.on_authentication_success` will be called a token string that
    that the server may optionally have sent.

    The exact nature of the negotiation between the client and server is specific
    to the authentication mechanism configured server-side.

    .. versionadded:: 2.0.0
    """

    server_authenticator_class = None
    """ Set during the connection AUTHENTICATE phase """

    def initial_response(self):
        """
        Returns an message to send to the server to initiate the SASL handshake.
        :const:`None` may be returned to send an empty message.
        """
        return None

    def evaluate_challenge(self, challenge):
        """
        Called when the server sends a challenge message.  Generally, this method
        should return :const:`None` when authentication is complete from a
        client perspective.  Otherwise, a string should be returned.
        """
        raise NotImplementedError()

    def on_authentication_success(self, token):
        """
        Called when the server indicates that authentication was successful.
        Depending on the authentication mechanism, `token` may be :const:`None`
        or a string.
        """
        pass


class PlainTextAuthProvider(AuthProvider):
    """
    An :class:`~.AuthProvider` that works with Cassandra's PasswordAuthenticator.

    Example usage::

        from cassandra.cluster import Cluster
        from cassandra.auth import PlainTextAuthProvider

        auth_provider = PlainTextAuthProvider(
                username='cassandra', password='cassandra')
        cluster = Cluster(auth_provider=auth_provider)

    .. versionadded:: 2.0.0
    """

    def __init__(self, username, password):
        self.username = username
        self.password = password

    def new_authenticator(self, host):
        return PlainTextAuthenticator(self.username, self.password)


class TransitionalModePlainTextAuthProvider(object):
    """
    An :class:`~.AuthProvider` that works with DSE TransitionalModePlainTextAuthenticator.

    Example usage::

        from cassandra.cluster import Cluster
        from cassandra.auth import TransitionalModePlainTextAuthProvider

        auth_provider = TransitionalModePlainTextAuthProvider()
        cluster = Cluster(auth_provider=auth_provider)

    .. warning:: TransitionalModePlainTextAuthProvider will be removed in cassandra-driver
                 4.0. The transitional mode will be handled internally without the need
                 of any auth provider.
    """

    def __init__(self):
        # TODO remove next major
        log.warning("TransitionalModePlainTextAuthProvider will be removed in cassandra-driver "
                    "4.0. The transitional mode will be handled internally without the need "
                    "of any auth provider.")

    def new_authenticator(self, host):
        return TransitionalModePlainTextAuthenticator()


class SaslAuthProvider(AuthProvider):
    """
    An :class:`~.AuthProvider` supporting general SASL auth mechanisms

    Suitable for GSSAPI or other SASL mechanisms

    Example usage::

        from cassandra.cluster import Cluster
        from cassandra.auth import SaslAuthProvider

        sasl_kwargs = {'service': 'something',
                       'mechanism': 'GSSAPI',
                       'qops': 'auth'.split(',')}
        auth_provider = SaslAuthProvider(**sasl_kwargs)
        cluster = Cluster(auth_provider=auth_provider)

    .. versionadded:: 2.1.4
    """

    def __init__(self, **sasl_kwargs):
        if SASLClient is None:
            raise ImportError('The puresasl library has not been installed')
        if 'host' in sasl_kwargs:
            raise ValueError("kwargs should not contain 'host' since it is passed dynamically to new_authenticator")
        self.sasl_kwargs = sasl_kwargs

    def new_authenticator(self, host):
        return SaslAuthenticator(host, **self.sasl_kwargs)


class SaslAuthenticator(Authenticator):
    """
    A pass-through :class:`~.Authenticator` using the third party package
    'pure-sasl' for authentication

    .. versionadded:: 2.1.4
    """

    def __init__(self, host, service, mechanism='GSSAPI', **sasl_kwargs):
        if SASLClient is None:
            raise ImportError('The puresasl library has not been installed')
        self.sasl = SASLClient(host, service, mechanism, **sasl_kwargs)

    def initial_response(self):
        return self.sasl.process()

    def evaluate_challenge(self, challenge):
        return self.sasl.process(challenge)

# TODO remove me next major
DSEPlainTextAuthProvider = PlainTextAuthProvider


class DSEGSSAPIAuthProvider(AuthProvider):
    """
    Auth provider for GSS API authentication. Works with legacy `KerberosAuthenticator`
    or `DseAuthenticator` if `kerberos` scheme is enabled.
    """
    def __init__(self, service='dse', qops=('auth',), resolve_host_name=True, **properties):
        """
        :param service: name of the service
        :param qops: iterable of "Quality of Protection" allowed; see ``puresasl.QOP``
        :param resolve_host_name: boolean flag indicating whether the authenticator should reverse-lookup an FQDN when
            creating a new authenticator. Default is ``True``, which will resolve, or return the numeric address if there is no PTR
            record. Setting ``False`` creates the authenticator with the numeric address known by Cassandra
        :param properties: additional keyword properties to pass for the ``puresasl.mechanisms.GSSAPIMechanism`` class.
            Presently, 'principal' (user) is the only one referenced in the ``pure-sasl`` implementation
        """
        if not _have_puresasl:
            raise ImportError('The puresasl library has not been installed')
        if not _have_kerberos:
            raise ImportError('The kerberos library has not been installed')
        self.service = service
        self.qops = qops
        self.resolve_host_name = resolve_host_name
        self.properties = properties

    def new_authenticator(self, host):
        if self.resolve_host_name:
            host = socket.getnameinfo((host, 0), 0)[0]
        return GSSAPIAuthenticator(host, self.service, self.qops, self.properties)


class BaseDSEAuthenticator(Authenticator):
    def get_mechanism(self):
        raise NotImplementedError("get_mechanism not implemented")

    def get_initial_challenge(self):
        raise NotImplementedError("get_initial_challenge not implemented")

    def initial_response(self):
        if self.server_authenticator_class == "com.datastax.bdp.cassandra.auth.DseAuthenticator":
            return self.get_mechanism()
        else:
            return self.evaluate_challenge(self.get_initial_challenge())


class PlainTextAuthenticator(BaseDSEAuthenticator):

    def __init__(self, username, password):
        self.username = username
        self.password = password

    def get_mechanism(self):
        return b"PLAIN"

    def get_initial_challenge(self):
        return b"PLAIN-START"

    def evaluate_challenge(self, challenge):
        if challenge == b'PLAIN-START':
            data = "\x00%s\x00%s" % (self.username, self.password)
            return data.encode()
        raise Exception('Did not receive a valid challenge response from server')


class TransitionalModePlainTextAuthenticator(PlainTextAuthenticator):
    """
     Authenticator that accounts for DSE authentication is configured with transitional mode.
    """

    def __init__(self):
        super(TransitionalModePlainTextAuthenticator, self).__init__('', '')


class GSSAPIAuthenticator(BaseDSEAuthenticator):
    def __init__(self, host, service, qops, properties):
        properties = properties or {}
        self.sasl = SASLClient(host, service, 'GSSAPI', qops=qops, **properties)

    def get_mechanism(self):
        return b"GSSAPI"

    def get_initial_challenge(self):
        return b"GSSAPI-START"

    def evaluate_challenge(self, challenge):
        if challenge == b'GSSAPI-START':
            return self.sasl.process()
        else:
            return self.sasl.process(challenge)