# Copyright: (c) 2020, Jordan Borean (@jborean93) <jborean93@gmail.com>
# MIT License (see LICENSE or https://opensource.org/licenses/MIT)

from __future__ import (absolute_import, division, print_function)

__metaclass__ = type  # noqa (fixes E402 for the imports below)

import base64
import collections
import logging

from spnego._compat import (
    List,
    Optional,
    Tuple,
)

from spnego._context import (
    ContextProxy,
    ContextReq,
    GSSMech,
)

from spnego._spnego import (
    NegState,
    NegTokenInit,
    NegTokenResp,
    pack_mech_type_list,
    unpack_token,
)

from spnego._text import (
    to_text,
)

from spnego.exceptions import (
    BadMechanismError,
    InvalidTokenError,
)

from spnego.gss import (
    GSSAPIProxy,
)

from spnego.ntlm import (
    NTLMProxy,
)

log = logging.getLogger(__name__)


class NegotiateProxy(ContextProxy):
    """A context wrapper for a Python managed SPNEGO context.

    This is a context that can be used on Linux to generate SPNEGO tokens based on the raw Kerberos or NTLM tokens
    generated by gssapi or our Python NTLM provider This is used as a fallback if gssapi is not available or cannot
    generate SPNEGO tokens.
    """

    def __init__(self, username=None, password=None, hostname=None, service=None, channel_bindings=None,
                 context_req=ContextReq.default, usage='initiate', protocol='negotiate', options=0):
        super(NegotiateProxy, self).__init__(username, password, hostname, service, channel_bindings, context_req,
                                             usage, protocol, options, False)

        self._hostname = hostname  # type: str
        self._service = service  # type: str
        self._complete = False  # type: bool
        self._context_list = None  # type: Optional[collections.OrderedDict[GSSMech, Tuple[ContextProxy, Optional[bytes]]]]  # noqa
        self.__chosen_mech = None  # type: Optional[GSSMech]
        self._mech_list = []  # type: (List[str])

        self._init_sent = False  # type: bool
        self._mech_sent = False  # type: bool
        self._mic_sent = False  # type: bool
        self._mic_recv = False  # type: bool
        self._mic_required = False  # type: bool

    @classmethod
    def available_protocols(cls, options=None):
        # We always support Negotiate and NTLM as we have our builtin NTLM backend and only support kerberos if gssapi
        # is present.
        protocols = ['ntlm', 'negotiate']

        # Make sure we add Kerberos first as the order is important.
        if 'kerberos' in GSSAPIProxy.available_protocols(options=options):
            protocols.insert(0, 'kerberos')

        return protocols

    @classmethod
    def iov_available(cls):
        # NTLM does not support IOV so we can only say that IOV is available if GSSAPI IOV is available.
        return GSSAPIProxy.iov_available()

    @property
    def client_principal(self):
        return self._context.client_principal

    @property
    def complete(self):
        return self._complete

    @property
    def context_attr(self):
        return self._context.context_attr

    @property
    def negotiated_protocol(self):
        return self._context.negotiated_protocol

    @property
    def session_key(self):
        return self._context.session_key if self._context else False

    def step(self, in_token=None):
        log.debug("SPNEGO step input: %s", to_text(base64.b64encode(in_token or b"")))

        # Step 1. Process SPNEGO mechs.
        mech_token_in, mech_list_mic, is_spnego = self._step_spnego_input(in_token=in_token)

        mech_token_out = None
        if mech_token_in or self.usage == 'initiate':
            # Step 2. Process the inner context tokens.
            mech_token_out = self._step_spnego_token(in_token=mech_token_in)

        if is_spnego:
            # Step 3. Process / generate the mechListMIC.
            out_mic = self._step_spnego_mic(in_mic=mech_list_mic)

            # Step 4. Generate the output SPNEGO token.
            out_token = self._step_spnego_output(out_token=mech_token_out, out_mic=out_mic)

        else:
            out_token = mech_token_out
            self._complete = self._context.complete

        if self.complete:
            # Remove the leftover contexts if there are still others remaining.
            self._context_list = collections.OrderedDict([(self._chosen_mech, (self._context, None))])

        log.debug("SPNEGO step output: %s" % to_text(base64.b64encode(out_token or b"")))

        return out_token

    def _step_spnego_input(self, in_token=None):
        mech_list_mic = None
        token = None
        is_spnego = True

        if in_token:
            in_token = unpack_token(in_token)

            if isinstance(in_token, NegTokenInit):
                mech_list_mic = in_token.mech_list_mic
                token = in_token.mech_token

                # This is the first token of the exchange, we should build our context list based on the mechs the
                # opposite end supports.
                mech_list = self._rebuild_context_list(mech_types=in_token.mech_types)

                if self.usage == 'initiate':
                    # If initiate processes a NegTokenInit2 token that's just used as a hint, use the actually
                    # supported mechs as the true mech list.
                    self._mech_list = mech_list

                else:
                    # If accept processes a NegTokenInit token we treat that as an actual init is sent so it does not
                    # send it's own and uses the initiate mech list as the true mech list.
                    self._init_sent = True
                    self._mech_list = in_token.mech_types

            elif isinstance(in_token, NegTokenResp):
                mech_list_mic = in_token.mech_list_mic
                token = in_token.response_token

                # If we have received the supported_mech then we don't need to send our own.
                if in_token.supported_mech:
                    if in_token.supported_mech != self._chosen_mech.value:
                        self._mic_required = True

                    self.__chosen_mech = GSSMech.from_oid(in_token.supported_mech)
                    self._mech_sent = True

                # Raise exception if we are rejected and have no error info (mechToken) that will give us more info.
                if in_token.neg_state == NegState.reject and not token:
                    raise InvalidTokenError(context_msg="Received SPNEGO rejection with no token error message")

                if in_token.neg_state == NegState.request_mic:
                    self._mic_required = True
                elif in_token.neg_state == NegState.accept_complete:
                    self._complete = True

            else:
                # This usually indicates the token is a raw NTLM or Kerberos token, return as is.
                is_spnego = False
                token = in_token

                self.__chosen_mech = GSSMech.ntlm if token.startswith(b"NTLMSSP\x00") else GSSMech.kerberos

                if not self._context_list:
                    self._rebuild_context_list(mech_types=[self.__chosen_mech.value])

        else:
            self._mech_list = self._rebuild_context_list()

        return token, mech_list_mic, is_spnego

    def _step_spnego_token(self, in_token=None):
        chosen_mech = self._chosen_mech
        context, generated_token = self._context_list[chosen_mech]

        if not context.complete:
            # When usage == 'initiate', the context_list may contain a pre-cached token which we use instead.

            if generated_token:
                out_token = generated_token
                self._context_list[chosen_mech] = (context, None)  # Clear the value as it's no longer required.

            else:
                out_token = context.step(in_token=in_token)

            # NTLM has a special case where we need to tell it it's ok to generate the MIC and also determine if
            # it actually did set the MIC as that controls the mechListMIC for the SPNEGO token.
            if self._requires_mech_list_mic:
                self._mic_required = True

            return out_token

    def _step_spnego_mic(self, in_mic=None):
        if in_mic:
            self.verify(pack_mech_type_list(self._mech_list), in_mic)
            self._reset_ntlm_crypto_state(outgoing=False)

            self._mic_required = True  # If we received a mechListMIC we need to send one back.
            self._mic_recv = True

            if self._mic_sent:
                self._complete = True

        if self._mic_required and not self._mic_sent:
            out_mic = self.sign(pack_mech_type_list(self._mech_list))
            self._reset_ntlm_crypto_state()

            self._mic_sent = True

            return out_mic

    def _step_spnego_output(self, out_token=None, out_mic=None):
        if not self._init_sent:
            self._init_sent = True

            init_kwargs = {
                'mech_token': out_token,
                'mech_list_mic': out_mic,
            }
            if self.usage == 'accept':
                init_kwargs['hint_name'] = b'not_defined_in_RFC4178@please_ignore'

            return NegTokenInit(self._mech_list, **init_kwargs).pack()

        elif not self.complete:
            # As per RFC 4178 - 4.2.2: supportedMech should only be present in the first reply from the target.
            # https://tools.ietf.org/html/rfc4178#section-4.2.2
            supported_mech = None
            if not self._mech_sent:
                supported_mech = self._chosen_mech.value
                self._mech_sent = True

            state = NegState.accept_incomplete

            if self._context.complete:
                if self._mic_sent and not self._mic_recv:
                    state = NegState.request_mic
                else:
                    state = NegState.accept_complete
                    self._complete = True

            return NegTokenResp(neg_state=state, supported_mech=supported_mech, response_token=out_token,
                                mech_list_mic=out_mic).pack()

    def wrap(self, data, encrypt=True, qop=None):
        return self._context.wrap(data, encrypt=encrypt, qop=qop)

    def wrap_iov(self, iov, encrypt=True, qop=None):
        return self._context.wrap_iov(iov, encrypt=encrypt, qop=qop)

    def wrap_winrm(self, data):
        return self._context.wrap_winrm(data)

    def unwrap(self, data):
        return self._context.unwrap(data)

    def unwrap_iov(self, iov):
        return self._context.unwrap_iov(iov)

    def unwrap_winrm(self, header, data):
        return self._context.unwrap_winrm(header, data)

    def sign(self, data, qop=None):
        return self._context.sign(data, qop=qop)

    def verify(self, data, mic):
        return self._context.verify(data, mic)

    @property
    def _context(self):
        if self._context_list:
            return self._context_list[self._chosen_mech][0]

    @property
    def _chosen_mech(self):
        if self.__chosen_mech:
            return self.__chosen_mech

        return next(iter(self._context_list))

    @property
    def _context_attr_map(self):
        return []  # SPNEGO layer uses the generic commands, the underlying context has it's own specific map.

    @property
    def _requires_mech_list_mic(self):
        return self._context._requires_mech_list_mic

    def _convert_iov_buffer(self, iov):
        pass  # Handled in the underlying context.  # pragma: no cover

    def _rebuild_context_list(self, mech_types=None):  # type: (Optional[List[str]]) -> List[str]
        """ Builds a new context list that are available to the client. """
        context_kwargs = {
            'username': self.username,
            'password': self.password,
            'hostname': self._hostname,
            'service': self._service,
            'channel_bindings': self.channel_bindings,
            'context_req': self.context_req,
            'usage': self.usage,
            'options': self.options,
            '_is_wrapped': True,
        }
        gssapi_protocols = [p for p in GSSAPIProxy.available_protocols(options=self.options) if p != 'negotiate']
        all_protocols = gssapi_protocols[:]
        if 'ntlm' not in all_protocols:
            all_protocols.append('ntlm')

        self._context_list = collections.OrderedDict()
        mech_list = []
        last_err = None
        for protocol in all_protocols:
            mech = getattr(GSSMech, protocol)
            if mech_types and mech.value not in mech_types:
                continue

            try:
                proxy_obj = GSSAPIProxy if protocol in gssapi_protocols else NTLMProxy
                context = proxy_obj(protocol=protocol, **context_kwargs)
                first_token = context.step() if self.usage == 'initiate' else None
            except Exception as e:
                last_err = e
                log.debug("Failed to create gssapi context for SPNEGO protocol %s: %s", protocol, to_text(e))
                continue

            self._context_list[mech] = (context, first_token)
            mech_list.append(mech.value)

        if not mech_list:
            raise BadMechanismError(context_msg="Unable to negotiate common mechanism", base_error=last_err)

        return mech_list

    def _reset_ntlm_crypto_state(self, outgoing=True):
        return self._context._reset_ntlm_crypto_state(outgoing=outgoing)
