from base64 import b64decode, b64encode
import functools
import socket
from typing import Union

from flask import request, Response, make_response, g
import gssapi
from gssapi.raw import RequirementFlag

_server_name = None


def init_spnego(service_name: str, fqdn: Union[str, None] = None):
    """Set the server principal which will be used for SPNEGO."""
    global _server_name
    if fqdn is None:
        fqdn = socket.getfqdn()
    _server_name = gssapi.Name('ceod/' + fqdn)

    # make sure that we're actually capable of acquiring credentials
    gssapi.Credentials(usage='accept', name=_server_name)


def requires_authentication(f):
    """
    Requires that all requests to f have a GSSAPI initiator token.
    The initiator principal will be passed to the first argument of f
    in the form user@REALM.
    """
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        if 'authorization' not in request.headers:
            return Response('Unauthorized', 401, {'WWW-Authenticate': 'Negotiate'})
        header = request.headers['authorization']
        client_token = b64decode(header.split()[1])
        creds = gssapi.Credentials(usage='accept', name=_server_name)
        ctx = gssapi.SecurityContext(creds=creds, usage='accept')
        server_token = ctx.step(client_token)

        # OK so we're going to cheat a bit here by assuming that Kerberos is the
        # mechanism being used (which we know will be true). We know that Kerberos
        # only requires one round-trip for the service handshake, so we don't need
        # to store state between requests. Just to be sure, we assert that this is
        # indeed the case.
        # (This isn't compliant with the GSSAPI spec, but why write more code than
        # necessary?)
        assert ctx.complete, 'only one round trip expected'

        # Store the username in flask.g
        client_princ = str(ctx.initiator_name)
        g.auth_user = client_princ[:client_princ.index('@')]

        # Store the delegated credentials, if they were given
        if ctx.actual_flags & RequirementFlag.delegate_to_peer:
            # For some reason, shit gets screwed up when you try to use a
            # gssapi.Credentials object which was created in another function.
            # So we're going to export the token instead (which is a bytes
            # object) and pass it to Credentials() whenever we need it.
            g.client_token = ctx.delegated_creds.export()

        # TODO: don't pass client_princ to f anymore since it's stored in flask.g
        resp = make_response(f(client_princ, *args, **kwargs))
        # RFC 2744, section 5.1:
        # "If no token need be sent, gss_accept_sec_context will indicate this
        # by setting the length field of the output_token argument to zero."
        if server_token is not None:
            resp.headers['WWW-Authenticate'] = 'Negotiate ' + b64encode(server_token).decode()
        return resp
    return wrapper