GSSAPI="BASE"  # This ensures that a full module is generated by Cython

from libc.stdlib cimport malloc, calloc, free
from libc.string cimport memcpy

from gssapi.raw.cython_types cimport *
from gssapi.raw.sec_contexts cimport SecurityContext

from gssapi.raw.misc import GSSError
from gssapi.raw import types as gssapi_types
from gssapi.raw.named_tuples import IOVUnwrapResult
from collections import namedtuple
from collections.abc import Sequence

from enum import IntEnum
from gssapi.raw._enum_extensions import ExtendableEnum

# Kept for backwards compatibility - functions used to be declared here
try:
    from gssapi.raw.ext_dce_aead import wrap_aead, unwrap_aead
except ImportError:
    pass


cdef extern from "python_gssapi_ext.h":
    """
    #ifdef OSX_HAS_GSS_FRAMEWORK
    #define gss_wrap_iov __ApplePrivate_gss_wrap_iov
    #define gss_unwrap_iov __ApplePrivate_gss_unwrap_iov
    #define gss_wrap_iov_length __ApplePrivate_gss_wrap_iov_length
    #define gss_release_iov_buffer __ApplePrivate_gss_release_iov_buffer
    #endif
    """

    # NB(directxman12): this wiki page has a different argument order
    #                   than the header file, and uses size_t instead of int
    #                   (this file matches the header file)
    OM_uint32 gss_wrap_iov(OM_uint32 *min_stat, gss_ctx_id_t ctx_handle,
                           int conf_req_flag, gss_qop_t qop_req, int *conf_ret,
                           gss_iov_buffer_desc *iov, int iov_count) nogil

    OM_uint32 gss_unwrap_iov(OM_uint32 *min_stat, gss_ctx_id_t ctx_handle,
                             int* conf_ret, gss_qop_t *qop_ret,
                             gss_iov_buffer_desc *iov, int iov_count) nogil

    OM_uint32 gss_wrap_iov_length(OM_uint32 *min_stat, gss_ctx_id_t ctx_handle,
                                  int conf_req, gss_qop_t qop_req,
                                  int *conf_ret, gss_iov_buffer_desc *iov,
                                  int iov_count) nogil

    OM_uint32 gss_release_iov_buffer(OM_uint32 *min_stat,
                                     gss_iov_buffer_desc *iov,
                                     int iov_count) nogil

    gss_iov_buffer_t GSS_C_NO_IOV_BUFFER

    OM_uint32 GSS_IOV_BUFFER_TYPE_EMPTY
    OM_uint32 GSS_IOV_BUFFER_TYPE_DATA
    OM_uint32 GSS_IOV_BUFFER_TYPE_HEADER
    OM_uint32 GSS_IOV_BUFFER_TYPE_MECH_PARAMS
    OM_uint32 GSS_IOV_BUFFER_TYPE_TRAILER
    OM_uint32 GSS_IOV_BUFFER_TYPE_PADDING
    OM_uint32 GSS_IOV_BUFFER_TYPE_STREAM
    OM_uint32 GSS_IOV_BUFFER_TYPE_SIGN_ONLY

    OM_uint32 GSS_IOV_BUFFER_FLAG_MASK
    OM_uint32 GSS_IOV_BUFFER_FLAG_ALLOCATE
    OM_uint32 GSS_IOV_BUFFER_FLAG_ALLOCATED

    # a few more are in the enum extension file


class IOVBufferType(IntEnum, metaclass=ExtendableEnum):
    empty = GSS_IOV_BUFFER_TYPE_EMPTY
    data = GSS_IOV_BUFFER_TYPE_DATA
    header = GSS_IOV_BUFFER_TYPE_HEADER
    mech_params = GSS_IOV_BUFFER_TYPE_MECH_PARAMS
    trailer = GSS_IOV_BUFFER_TYPE_TRAILER
    padding = GSS_IOV_BUFFER_TYPE_PADDING
    stream = GSS_IOV_BUFFER_TYPE_STREAM
    sign_only = GSS_IOV_BUFFER_TYPE_SIGN_ONLY


IOVBuffer = namedtuple('IOVBuffer', ['type', 'allocate', 'value'])


cdef class IOV:
    # defined in ext_dce.pxd

    # cdef int iov_len
    # cdef bint c_changed

    # cdef gss_iov_buffer_desc *_iov
    # cdef bint _unprocessed
    # cdef list _buffs

    AUTO_ALLOC_BUFFERS = set([IOVBufferType.header, IOVBufferType.padding,
                              IOVBufferType.trailer])

    def __init__(IOV self, *args, std_layout=True, auto_alloc=True):
        self._unprocessed = True
        self.c_changed = False

        self._buffs = []

        if std_layout:
            self._buffs.append(IOVBuffer(IOVBufferType.header,
                                         auto_alloc, None))

        cdef char *val_copy
        for buff_desc in args:
            if isinstance(buff_desc, tuple):
                if len(buff_desc) > 3 or len(buff_desc) < 2:
                    raise ValueError("Buffer description tuples must be "
                                     "length 2 or 3")

                buff_type = buff_desc[0]

                if len(buff_desc) == 2:
                    if buff_type in self.AUTO_ALLOC_BUFFERS:
                        alloc = buff_desc[1]
                        data = None
                    else:
                        data = buff_desc[1]
                        alloc = False
                else:
                    (buff_type, alloc, data) = buff_desc

                self._buffs.append(IOVBuffer(buff_type, alloc, data))
            elif isinstance(buff_desc, bytes):  # assume type data
                val = buff_desc
                self._buffs.append(IOVBuffer(IOVBufferType.data, False, val))
            else:
                alloc = False
                if buff_desc in self.AUTO_ALLOC_BUFFERS:
                    alloc = auto_alloc

                self._buffs.append(IOVBuffer(buff_desc, alloc, None))

        if std_layout:
            self._buffs.append(IOVBuffer(IOVBufferType.padding, auto_alloc,
                                         None))
            self._buffs.append(IOVBuffer(IOVBufferType.trailer, auto_alloc,
                                         None))

    cdef gss_iov_buffer_desc* __cvalue__(IOV self) except NULL:
        cdef OM_uint32 tmp_min_stat
        cdef int i
        if self._unprocessed:
            if self._iov is not NULL:
                gss_release_iov_buffer(&tmp_min_stat, self._iov, self.iov_len)
                free(self._iov)

            self.iov_len = len(self._buffs)
            self._iov = <gss_iov_buffer_desc *>calloc(
                self.iov_len, sizeof(gss_iov_buffer_desc))
            if self._iov is NULL:
                raise MemoryError("Cannot calloc for IOV buffer array")

            for i in range(self.iov_len):
                buff = self._buffs[i]
                self._iov[i].type = buff.type

                if buff.allocate:
                    self._iov[i].type |= GSS_IOV_BUFFER_FLAG_ALLOCATE
                elif buff.allocate is None:
                    self._iov[i].type |= GSS_IOV_BUFFER_FLAG_ALLOCATED

                if buff.value is None:
                    self._iov[i].buffer.length = 0
                    self._iov[i].buffer.value = NULL
                else:
                    self._iov[i].buffer.length = len(buff.value)
                    self._iov[i].buffer.value = <char *>malloc(
                        self._iov[i].buffer.length)
                    if self._iov[i].buffer.value is NULL:
                        raise MemoryError("Cannot malloc for buffer value")

                    memcpy(self._iov[i].buffer.value, <char *>buff.value,
                           self._iov[i].buffer.length)

        return self._iov

    cdef _recreate_python_values(IOV self):
        cdef i
        cdef bint val_change = False
        cdef size_t new_len
        for i in range(self.iov_len):
            old_type = self._buffs[i].type

            if self._iov[i].buffer.value is NULL:
                if self._iov[i].buffer.length == 0:
                    new_val = None
                else:
                    new_len = self._iov[i].buffer.length
                    new_val = b'\x00' * new_len
            else:
                new_len = self._iov[i].buffer.length
                new_val = (<char*>self._iov[i].buffer.value)[:new_len]

            alloc = False
            if self._iov[i].type & GSS_IOV_BUFFER_FLAG_ALLOCATE:
                alloc = True

            # NB(directxman12): GSSAPI (at least in MIT krb5) doesn't
            # unset the allocate flag (because it's an "input flag",
            # so this needs to come second and be separate
            if self._iov[i].type & GSS_IOV_BUFFER_FLAG_ALLOCATED:
                alloc = None

            self._buffs[i] = IOVBuffer(old_type, alloc, new_val)

        self.c_changed = False

    def __getitem__(IOV self, ind):
        if self.c_changed:
            self._recreate_python_values()

        return self._buffs[ind]

    def __len__(IOV self):
        if self.c_changed:
            self._recreate_python_values()

        return len(self._buffs)

    def __iter__(IOV self):
        if self.c_changed:
            self._recreate_python_values()

        for val in self._buffs:
            yield val

    def __contains__(IOV self, item):
        if self.c_changed:
            self._recreate_python_values()

        return item in self._buffs

    def __reversed__(IOV self):
        if self.c_changed:
            self._recreate_python_values()

        for val in reversed(self._buffs):
            yield val

    def index(IOV self, value):
        for i, v in enumerate(self):
            if v == value:
                return i

        raise ValueError

    def count(IOV self, value):
        return sum(1 for v in self if v == value)

    def __repr__(IOV self):
        if self.c_changed:
            self._recreate_python_values()

        return "<{module}.{name} {buffs}>".format(
            module=type(self).__module__, name=type(self).__name__,
            buffs=repr(self._buffs))

    def __str__(IOV self):
        buff_strs = []
        for buff in self:
            type_val = str(buff.type).split('.')[1].upper()
            if buff.value is None:
                auto_alloc = buff.allocate
                if auto_alloc:
                    buff_strs.append(type_val + "(allocate)")
                else:
                    buff_strs.append(type_val + "(empty)")
            else:
                if buff.allocate is None:
                    alloc_str = ", allocated"
                else:
                    alloc_str = ""
                buff_strs.append("{0}({1!r}{2})".format(type_val,
                                                        buff.value, alloc_str))

        return "<IOV {0}>".format(' | '.join(buff_strs))

    def __dealloc__(IOV self):
        cdef OM_uint32 tmp_min_stat
        cdef int i
        if self._iov is not NULL:
            gss_release_iov_buffer(&tmp_min_stat, self._iov, self.iov_len)

            for i in range(self.iov_len):
                if self._iov[i].buffer.value is not NULL:
                    free(self._iov[i].buffer.value)

            free(self._iov)


def wrap_iov(SecurityContext context not None, IOV message not None,
             confidential=True, qop=None):
    cdef int conf_req = confidential
    cdef gss_qop_t qop_req = qop if qop is not None else GSS_C_QOP_DEFAULT
    cdef int conf_used

    cdef gss_iov_buffer_desc *res_arr = message.__cvalue__()

    cdef OM_uint32 maj_stat, min_stat

    with nogil:
        maj_stat = gss_wrap_iov(&min_stat, context.raw_ctx, conf_req, qop_req,
                                &conf_used, res_arr, message.iov_len)

    if maj_stat == GSS_S_COMPLETE:
        message.c_changed = True
        return <bint>conf_used
    else:
        raise GSSError(maj_stat, min_stat)


def unwrap_iov(SecurityContext context not None, IOV message not None):
    cdef int conf_used
    cdef gss_qop_t qop_used
    cdef gss_iov_buffer_desc *res_arr = message.__cvalue__()

    cdef OM_uint32 maj_stat, min_stat

    with nogil:
        maj_stat = gss_unwrap_iov(&min_stat, context.raw_ctx, &conf_used,
                                  &qop_used, res_arr, message.iov_len)

    if maj_stat == GSS_S_COMPLETE:
        message.c_changed = True
        return IOVUnwrapResult(<bint>conf_used, qop_used)
    else:
        raise GSSError(maj_stat, min_stat)


def wrap_iov_length(SecurityContext context not None, IOV message not None,
                    confidential=True, qop=None):
    cdef int conf_req = confidential
    cdef gss_qop_t qop_req = qop if qop is not None else GSS_C_QOP_DEFAULT
    cdef int conf_used

    cdef gss_iov_buffer_desc *res_arr = message.__cvalue__()

    cdef OM_uint32 maj_stat, min_stat

    with nogil:
        maj_stat = gss_wrap_iov_length(&min_stat, context.raw_ctx,
                                       conf_req, qop_req,
                                       &conf_used, res_arr, message.iov_len)

    if maj_stat == GSS_S_COMPLETE:
        message.c_changed = True
        return <bint>conf_used
    else:
        raise GSSError(maj_stat, min_stat)