"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Monday, 26th Aug., 2024

    OBJECTIVE:

        To provide an easy way to hash inputs.

    REFERENCES:

        1) Book: Full Stack Python Security - Dennis Byrne

    DOWNLOADS:

        N/A

"""


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   IMPORT                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# For hashing:
import hashlib
import hmac
import secrets
from bcrypt import hashpw, gensalt

# To work with buffers:
import io


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                           MACROS / ONE-TIME INIT                                              ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                VARIABLES                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   CLASSES                                                     ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


class Hasher:

    def __init__(
            self,
            algorithm = hashlib.sha256,
            key = None
    ):

        """
        hashes data or a message or a file. Uses HMAC if 'key' is specified, else performs simple hashing.
        :param algorithm: The algorithm to use. SHA256 by default.
        :param key: Specify this as either a string or as bytes to use HMAC. Leave as null for simple hashing.
        """

        # Initialize the hasher:
        self.__hasher = None,
        self.__algorithm = algorithm
        self.__hmac_key = key
        if self.__hmac_key is not None:
            self.__hmac_key = self.__hmac_key.encode("utf-8") if isinstance(self.__hmac_key, str) else self.__hmac_key
        self.reset()

    @staticmethod
    def generate_key(byte_count = 32, url_safe = False):

        """
        A mechanism to generate key/salt values.
        NOTE: For a proper salt for passwords, I recommend using "generate_salt" method. It's far better.
        :param byte_count: The number of bytes to have in the key. The hex output (as string) will have 2x the
            characters.
        :param url_safe: Set to True if you need the generated output to be a part of a URL.
        :return: The generate key/salt.
        """

        return secrets.token_urlsafe(byte_count) if url_safe else secrets.token_hex(byte_count)

    @staticmethod
    def generate_salt():

        """
        Generate a salt to use while hashing things like passwords.
        :return: The salt as bytes.
        """

        return gensalt()

    @staticmethod
    def hash_password(password, salt):

        """
        Hashes a password with the given salt.
        :param password: The password to hash, either as a string or as bytes.
        :param salt: The salt to hash the password with, either as a string or as bytes.
        :return: The hashed string.
        """

        return hashpw(
            password = password.encode("utf-8") if isinstance(password, str) else password,
            salt = salt.encode("utf-8") if isinstance(salt, str) else salt
        )

    @staticmethod
    def compare_hashes(hash_0, hash_1):

        """
        Compares two hashes to see if they match.
        Comparison is done in constant time to avoid timing-based side-channel attacks.
        :param hash_0: One of the hashes to compare.
        :param hash_1: The other hash to compare.
        :return: True if they match, else False.
        """

        return hmac.compare_digest(hash_0, hash_1)

    def reset(self):

        """
        Resets the hasher by removing all the data that was fed into it.
        :return: None.
        """

        if self.__hmac_key is not None:
            self.__hasher = hmac.new(
                key = self.__hmac_key,
                digestmod = self.__algorithm
            )
        else: self.__hasher = self.__algorithm()

    def update(self, data):

        """
        Adds data to the hash to update it.
        :param data: The data to be hashed.
        :return: None.
        """

        data = data.encode("utf-8") if isinstance(data, str) else data
        self.__hasher.update(data)

    def digest(self):

        """
        Returns the hexadecimal representation of the hash as a string.
        :return: The hexadecimal representation of the hash as a string
        """

        return self.__hasher.digest()

    def hexdigest(self):

        """
        Returns the hexadecimal representation of the hash as a string.
        :return: The hexadecimal representation of the hash as a string
        """

        return self.__hasher.hexdigest()

    def hash_message(self, message, as_hex = True):

        """
        Hashes one message and returns the result, and then resets the instance.
        :param message: The data you want to hash.
        :param as_hex: Invokes 'hexdigest' if True, else 'digest'.
        :return: The hash of the message in either hexadecimal string form or binary form.
        """

        self.update(message)
        hash_result = self.hexdigest() if as_hex else self.digest()
        self.reset()
        return hash_result

    def hash_file(self, file, chunk_size = 4096, as_hex = True):

        """
        Hashes one file and returns the result, and then resets the instance.
        :param file: The file you want to hash either as a path or as some buffer (like io.BytesIO).
        :param chunk_size: The size of data (in bytes) that you would like to pick at one time.
        :param as_hex: Invokes 'hexdigest' if True, else 'digest'.
        :return: The hash of the file in either hexadecimal string form or binary form.
        """

        # In case the file was given as a io.BytesIO buffer:
        if isinstance(file, io.BytesIO):
            file.seek(0)
            while True:
                chunk = file.read(chunk_size)
                if not chunk: break
                self.update(chunk)

        # In case the file was given as a path:
        else:
            with open(file, "rb") as f:
                for chunk in iter(lambda: f.read(chunk_size), b""):
                    self.update(chunk)

        # Now we capture the results, reset the instance, and return the result:
        hash_result = self.hexdigest() if as_hex else self.digest()
        self.reset()
        return hash_result


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                               MAIN PROGRAM                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


if __name__ == "__main__":

    print(
        Hasher.hash_password(
            password = "mic test, mic test, 123",
            salt = Hasher.generate_salt()
        )
    )
