# Copyright 2013-2019 Donald Stufft and individual contributors
#
# 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.

from __future__ import absolute_import, division, print_function

from six import integer_types

from nacl import exceptions as exc
from nacl._sodium import ffi, lib
from nacl.exceptions import ensure


crypto_generichash_BYTES = lib.crypto_generichash_blake2b_bytes()
crypto_generichash_BYTES_MIN = lib.crypto_generichash_blake2b_bytes_min()
crypto_generichash_BYTES_MAX = lib.crypto_generichash_blake2b_bytes_max()
crypto_generichash_KEYBYTES = lib.crypto_generichash_blake2b_keybytes()
crypto_generichash_KEYBYTES_MIN = lib.crypto_generichash_blake2b_keybytes_min()
crypto_generichash_KEYBYTES_MAX = lib.crypto_generichash_blake2b_keybytes_max()
crypto_generichash_SALTBYTES = lib.crypto_generichash_blake2b_saltbytes()
crypto_generichash_PERSONALBYTES = \
    lib.crypto_generichash_blake2b_personalbytes()
crypto_generichash_STATEBYTES = lib.crypto_generichash_statebytes()

_OVERLONG = '{0} length greater than {1} bytes'
_TOOBIG = '{0} greater than {1}'


def _checkparams(digest_size, key, salt, person):
    """Check hash paramters"""
    ensure(isinstance(key, bytes),
           'Key must be a bytes sequence',
           raising=exc.TypeError)

    ensure(isinstance(salt, bytes),
           'Salt must be a bytes sequence',
           raising=exc.TypeError)

    ensure(isinstance(person, bytes),
           'Person must be a bytes sequence',
           raising=exc.TypeError)

    ensure(isinstance(digest_size, integer_types),
           'Digest size must be an integer number',
           raising=exc.TypeError)

    ensure(digest_size <= crypto_generichash_BYTES_MAX,
           _TOOBIG.format("Digest_size", crypto_generichash_BYTES_MAX),
           raising=exc.ValueError)

    ensure(len(key) <= crypto_generichash_KEYBYTES_MAX,
           _OVERLONG.format("Key", crypto_generichash_KEYBYTES_MAX),
           raising=exc.ValueError)

    ensure(len(salt) <= crypto_generichash_SALTBYTES,
           _OVERLONG.format("Salt", crypto_generichash_SALTBYTES),
           raising=exc.ValueError)

    ensure(len(person) <= crypto_generichash_PERSONALBYTES,
           _OVERLONG.format("Person", crypto_generichash_PERSONALBYTES),
           raising=exc.ValueError)


def generichash_blake2b_salt_personal(data,
                                      digest_size=crypto_generichash_BYTES,
                                      key=b'', salt=b'', person=b''):
    """One shot hash interface

    :param data: the input data to the hash function
    :param digest_size: must be at most
                        :py:data:`.crypto_generichash_BYTES_MAX`;
                        the default digest size is
                        :py:data:`.crypto_generichash_BYTES`
    :type digest_size: int
    :param key: must be at most
                :py:data:`.crypto_generichash_KEYBYTES_MAX` long
    :type key: bytes
    :param salt: must be at most
                 :py:data:`.crypto_generichash_SALTBYTES` long;
                 will be zero-padded if needed
    :type salt: bytes
    :param person: must be at most
                   :py:data:`.crypto_generichash_PERSONALBYTES` long:
                   will be zero-padded if needed
    :type person: bytes
    :return: digest_size long digest
    :rtype: bytes
    """

    _checkparams(digest_size, key, salt, person)

    ensure(isinstance(data, bytes),
           'Input data must be a bytes sequence',
           raising=exc.TypeError)

    digest = ffi.new("unsigned char[]", digest_size)

    # both _salt and _personal must be zero-padded to the correct length
    _salt = ffi.new("unsigned char []", crypto_generichash_SALTBYTES)
    _person = ffi.new("unsigned char []", crypto_generichash_PERSONALBYTES)

    ffi.memmove(_salt, salt, len(salt))
    ffi.memmove(_person, person, len(person))

    rc = lib.crypto_generichash_blake2b_salt_personal(digest, digest_size,
                                                      data, len(data),
                                                      key, len(key),
                                                      _salt, _person)
    ensure(rc == 0, 'Unexpected failure',
           raising=exc.RuntimeError)

    return ffi.buffer(digest, digest_size)[:]


class Blake2State(object):
    """
    Python-level wrapper for the crypto_generichash_blake2b state buffer
    """
    __slots__ = ['_statebuf', 'digest_size']

    def __init__(self, digest_size):
        self._statebuf = ffi.new("unsigned char[]",
                                 crypto_generichash_STATEBYTES)
        self.digest_size = digest_size

    def __reduce__(self):
        """
        Raise the same exception as hashlib's blake implementation
        on copy.copy()
        """
        raise TypeError("can't pickle {} objects".format(
            self.__class__.__name__))

    def copy(self):
        _st = self.__class__(self.digest_size)
        ffi.memmove(_st._statebuf,
                    self._statebuf, crypto_generichash_STATEBYTES)
        return _st


def generichash_blake2b_init(key=b'', salt=b'',
                             person=b'',
                             digest_size=crypto_generichash_BYTES):
    """
    Create a new initialized blake2b hash state

    :param key: must be at most
                :py:data:`.crypto_generichash_KEYBYTES_MAX` long
    :type key: bytes
    :param salt: must be at most
                 :py:data:`.crypto_generichash_SALTBYTES` long;
                 will be zero-padded if needed
    :type salt: bytes
    :param person: must be at most
                   :py:data:`.crypto_generichash_PERSONALBYTES` long:
                   will be zero-padded if needed
    :type person: bytes
    :param digest_size: must be at most
                        :py:data:`.crypto_generichash_BYTES_MAX`;
                        the default digest size is
                        :py:data:`.crypto_generichash_BYTES`
    :type digest_size: int
    :return: a initialized :py:class:`.Blake2State`
    :rtype: object
    """

    _checkparams(digest_size, key, salt, person)

    state = Blake2State(digest_size)

    # both _salt and _personal must be zero-padded to the correct length
    _salt = ffi.new("unsigned char []", crypto_generichash_SALTBYTES)
    _person = ffi.new("unsigned char []", crypto_generichash_PERSONALBYTES)

    ffi.memmove(_salt, salt, len(salt))
    ffi.memmove(_person, person, len(person))

    rc = lib.crypto_generichash_blake2b_init_salt_personal(state._statebuf,
                                                           key, len(key),
                                                           digest_size,
                                                           _salt, _person)
    ensure(rc == 0, 'Unexpected failure',
           raising=exc.RuntimeError)

    return state


def generichash_blake2b_update(state, data):
    """Update the blake2b hash state

    :param state: a initialized Blake2bState object as returned from
                     :py:func:`.crypto_generichash_blake2b_init`
    :type state: :py:class:`.Blake2State`
    :param data:
    :type data: bytes
    """

    ensure(isinstance(state, Blake2State),
           'State must be a Blake2State object',
           raising=exc.TypeError)

    ensure(isinstance(data, bytes),
           'Input data must be a bytes sequence',
           raising=exc.TypeError)

    rc = lib.crypto_generichash_blake2b_update(state._statebuf,
                                               data, len(data))
    ensure(rc == 0, 'Unexpected failure',
           raising=exc.RuntimeError)


def generichash_blake2b_final(state):
    """Finalize the blake2b hash state and return the digest.

    :param state: a initialized Blake2bState object as returned from
                     :py:func:`.crypto_generichash_blake2b_init`
    :type state: :py:class:`.Blake2State`
    :return: the blake2 digest of the passed-in data stream
    :rtype: bytes
    """

    ensure(isinstance(state, Blake2State),
           'State must be a Blake2State object',
           raising=exc.TypeError)

    _digest = ffi.new("unsigned char[]", crypto_generichash_BYTES_MAX)
    rc = lib.crypto_generichash_blake2b_final(state._statebuf,
                                              _digest, state.digest_size)

    ensure(rc == 0, 'Unexpected failure',
           raising=exc.RuntimeError)
    return ffi.buffer(_digest, state.digest_size)[:]