# Author: Hubert Kario (c) 2019
# see LICENCE file for legal information regarding use of this file

"""
HMAC module that works in FIPS mode.

Note that this makes this code FIPS non-compliant!
"""

# Because we are extending the hashlib module, we need to import all its
# fields to suppport the same uses
from . import tlshashlib
from .compat import compatHMAC
try:
    from hmac import compare_digest
    __all__ = ["new", "compare_digest", "HMAC"]
except ImportError:
    __all__ = ["new", "HMAC"]


class HMAC(object):
    """Hacked version of HMAC that works in FIPS mode even with MD5."""

    def __init__(self, key, msg=None, digestmod=None):
        """
        Initialise the HMAC and hash first portion of data.

        msg: data to hash
        digestmod: name of hash or object that be used as a hash and be cloned
        """
        self.key = key
        if digestmod is None:
            digestmod = 'md5'
        if callable(digestmod):
            digestmod = digestmod()
        if not hasattr(digestmod, 'digest_size'):
            digestmod = tlshashlib.new(digestmod)
        self.block_size = digestmod.block_size
        self.digest_size = digestmod.digest_size
        self.digestmod = digestmod
        if len(key) > self.block_size:
            k_hash = digestmod.copy()
            k_hash.update(compatHMAC(key))
            key = k_hash.digest()
        if len(key) < self.block_size:
            key = key + b'\x00' * (self.block_size - len(key))
        key = bytearray(key)
        ipad = bytearray(b'\x36' * self.block_size)
        opad = bytearray(b'\x5c' * self.block_size)
        i_key = bytearray(i ^ j for i, j in zip(key, ipad))
        self._o_key = bytearray(i ^ j for i, j in zip(key, opad))
        self._context = digestmod.copy()
        self._context.update(compatHMAC(i_key))
        if msg:
            self._context.update(compatHMAC(msg))

    def update(self, msg):
        self._context.update(compatHMAC(msg))

    def digest(self):
        i_digest = self._context.digest()
        o_hash = self.digestmod.copy()
        o_hash.update(compatHMAC(self._o_key))
        o_hash.update(compatHMAC(i_digest))
        return o_hash.digest()

    def copy(self):
        new = HMAC.__new__(HMAC)
        new.key = self.key
        new.digestmod = self.digestmod
        new.block_size = self.block_size
        new.digest_size = self.digest_size
        new._o_key = self._o_key
        new._context = self._context.copy()
        return new


def new(*args, **kwargs):
    """General constructor that works in FIPS mode."""
    return HMAC(*args, **kwargs)
