""" This module holds the Repository class """
import base64
import os
import sys
import json
import zlib
import bcrypt

from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import constant_time


class ChecksumException(Exception):
    """
    Class for raising Checksum Mismatch Exception
    This is raised during load() if checksums did not match
    """
    pass


class AuthenticationException(Exception):
    """
    Class for raising Authentication Failure Exception
    This is raised during load() if password is wrong
    """
    pass


class Repository:
    """
    Class representing a repository.
    Mimics a dict, but with context management and custom storage mechanism
    """
    # Repository configuration options
    # This is not None so that we can call .get() with a default return
    # And get a valid result in case the caller never bothered to send us the config
    config = {}

    def __init__(self, name, password):
        """
        Constructor for new repository object

        Keyword arguments:
            name -- Name of repository to open
            password -- Password used to secure repository
        """
        # Repository name shouldn't contain a space, to avoid file system issues
        assert " " not in name
        self.__name = name  # Repository is identified by the name
        self.__passwd = password    # The user password is stored because we need it later
        # The repo is stored in a .repo file under the directory specified by the config
        self.__file = Repository.config.get('storage', 'repositories') + '/' + name + '.repo'
        self.__metadata = None  # Metadata is either loaded from file or created by the first call to load()
        self.__data = None  # None because we haven't decrypted it yet

    @staticmethod
    def __generate_nonce():
        """
        Generate new nonce
        Return generated nonce
        """
        noncesize = int(Repository.config.get('nonce.size', 512))    # Get the nonce size
        nonce = base64.b64encode(os.urandom(noncesize))         # Generate a new nonce
        return nonce

    @staticmethod
    def __derive_repokey(nonce, salt=None):
        """
        Derive repokey from nonce
        Return derived repository password (repokey) and the salt
        """
        size = 32
        times = int(Repository.config.get('key.iterations', 10000))        # Get the hashing iteration count
        if not salt:        # No salt provided, generate one
            salt = os.urandom(size)
        # Create a KDF instance
        kdf = PBKDF2HMAC(                 # Using PBKDF2
            algorithm=hashes.SHA256(),    # HMAC SHA256
            length=size,                  # Length of key (Has to be 32 to be a Fernet key)
            salt=salt,                    # Salt to use
            iterations=times,             # Iterated
            backend=default_backend()     # OpenSSL is the default backend usually
        )
        key = base64.urlsafe_b64encode(kdf.derive(nonce))     # Derive a password
        return key, salt  # Return the generated password and salt used

    def __generate_newkey(self):
        """
        Generate new nonce and derive key from it.
        Return (nonce, key, salt)
        """
        nonce = self.__generate_nonce()             # Generate a new nonce
        key, salt = self.__derive_repokey(nonce)    # Derive the repository key from it
        return nonce, key, salt   # Return the generation parameters and resulting key

    def __derive_noncekey(self, salt=None):
        """
        Derive key from password
        Return derived noncekey (used to encrypt nonce)
        """
        size = 32
        times = int(Repository.config.get('noncekey.iterations', 100000))    # Number of times to hash
        # If no salt is given, use a newly generated one
        if not salt:
            salt = os.urandom(size)
        # Create a KDF instance
        kdf = PBKDF2HMAC(                 # Using PBKDF2
            algorithm=hashes.SHA256(),    # HMAC SHA256
            length=size,                  # Length of key (Has to be 32 to be a Fernet key)
            salt=salt,                    # Salt to use
            iterations=times,             # Iterated
            backend=default_backend()     # OpenSSL is the default backend usually
        )
        key = base64.urlsafe_b64encode(kdf.derive(bytes(self.__passwd, 'utf-8')))     # Derive a new key
        return key, salt      # Return the key and the salt used

    def __encrypt_nonce(self, nonce):
        """
        Encrypt the nonce by calling __derive_noncekey() to get the key and salt
        And then encrypt it to get encrypted nonce
        Return encrypted nonce (enonce, salt)
        """
        noncekey, psalt = self.__derive_noncekey()  # Derive the repository key
        # Encrypt the nonce with the nonce key
        fern = Fernet(noncekey)
        enonce = fern.encrypt(nonce)
        return enonce, psalt  # Return the encrypted nonce and the salt

    def __decrypt_nonce(self, enonce, salt):
        """
        Decrypt the nonce using the given salt
        And then decrypt it
        Return the nonce
        """
        noncekey, _ = self.__derive_noncekey(salt=salt)     # Derive the nonce key
        # Decrypt the nonce
        fern = Fernet(noncekey)
        nonce = fern.decrypt(enonce)
        return nonce    # Return plaintext nonce

    def __calc_data_checksum(self, data=None):
        """
        Calculate the checksum of the data in the repo
        Return checksum
        """
        # Whirlpool instance used to calculate the checksum
        digest = hashes.Hash(hashes.Whirlpool(), backend=default_backend())
        # If no data provided, use the repository data
        if not data:
            data = bytes(json.dumps(self.__data), 'utf-8')
        # Calculate and return the checksum
        digest.update(data)
        return digest.finalize()

    def __encrypt_data(self, repokey):
        """
        Encrypt the data in the repository
        Return (encrypted data, checksum)
        """
        fern = Fernet(repokey)
        data = bytes(json.dumps(self.__data), 'utf-8')      # Put data into a serialized format
        edata = fern.encrypt(data)
        checksum = self.__calc_data_checksum()      # Calculate the checksum
        return edata, checksum        # Return the encrypted data and checksum

    def __decrypt_data(self, edata, checksum, repokey):
        """
        Decrypt the data given the repokey (verifying the integrity)
        Return plaintext data
        """
        fern = Fernet(repokey)
        data = fern.decrypt(edata)      # Decrypt the data
        # Check the checksum
        if self.__calc_data_checksum(data) != checksum:
            print("Checksum did not match", file=sys.stderr)
            raise ChecksumException
        return data     # Return the decrypted data

    def __encrypt(self):
        """
        Re-Encrypt the repository:
            - Generate new nonce
            - Use nonce to derive new key
            - Encrypt nonce using derived key from users password
            - Store it
            - Use new key to encrypt repository.
        """
        nonce, repokey, noncesalt = self.__generate_newkey()        # Nonce, repokey and salt used in derivation
        enonce, psalt = self.__encrypt_nonce(nonce)         # Encrypted nonce
        edata, checksum = self.__encrypt_data(repokey)      # get data and checksum encrypted
        # Metadata updates
        # Encrypted Nonce
        self.__metadata["EncryptedNonce"] = base64.urlsafe_b64encode(enonce).decode('utf-8')
        # Salt used to derive from the nonce
        self.__metadata["NonceSalt"] = base64.urlsafe_b64encode(noncesalt).decode('utf-8')
        # Salt used to derive from the password
        self.__metadata["PasswordSalt"] = base64.urlsafe_b64encode(psalt).decode('utf-8')
        # Repository data checksum
        self.__metadata["Checksum"] = base64.urlsafe_b64encode(checksum).decode('utf-8')
        return edata

    def __decrypt(self, edata):
        """
        Decrypt the repository
            - Use passwd and psalt to derive noncekey
            - Use nonce key to decrypt nonce
            - Use decrypted nonce and noncesalt to derive repokey
            - Use repokey to decrypt the edata and verify the checksum
            Return the decrypted data
        """
        # Gather Metadata into variables
        enonce = base64.urlsafe_b64decode(self.__metadata["EncryptedNonce"])    # Encrypted Nonce
        noncesalt = base64.urlsafe_b64decode(self.__metadata["NonceSalt"])      # Salt used to derive from the nonce
        psalt = base64.urlsafe_b64decode(self.__metadata["PasswordSalt"])       # Salt used to derive from the password
        checksum = base64.urlsafe_b64decode(self.__metadata["Checksum"])        # Repository data checksum
        # Derive the noncekey first, then decrypt the nonce
        nonce = self.__decrypt_nonce(enonce, salt=psalt)
        # Derive the repokey from the nonce
        repokey, _ = self.__derive_repokey(nonce, salt=noncesalt)
        # Decrypt the repository data using the repokey
        data = self.__decrypt_data(edata, checksum, repokey=repokey).decode('utf-8')
        return json.loads(data)

    def __compress(self, edata):
        """
        Compress the repository contents into byte array
        Return compressed byte array
        """
        # Dump the data and metadata into a serializable format
        data = {"Data": edata.decode('utf-8'), "Metadata":self.__metadata}
        # Dump into a JSON string, and then compress using gzip
        return zlib.compress(bytes(json.dumps(data), 'utf-8'))

    def __decompress(self, filedata):
        """
        Decompress the supplied byte array into repository contents
        Return the encrypted data, loads Metadata into class member
        """
        # Decompress data fro gzipped format
        data = zlib.decompress(filedata)
        data = data.decode('utf-8')     # Decode from bytearray
        # Load into dict
        data = json.loads(data)
        assert all(k in data for k in ["Data", "Metadata"])        # Check for correctness
        self.__metadata = data["Metadata"]      # Update the Metadata
        return bytes(data["Data"], 'utf-8')     # Return the encrypted data

    def load(self):
        """
        Decrypt the password repository.
            - Load the file into memory.
            - Read the Data and Metadata
            - Verify the password
            - Decrypt the repository
            - Load data
            Return nothing
        """
        # If file does not exist, then this is a new repo creation in a with statement
        # In this case we just create the file (empty) and return
        if not os.path.exists(self.__file):
            self.__metadata = dict()
            self.__data = dict()
            return
        # Load the binary data into bytearray
        with open(self.__file, 'rb') as datafile:
            filedata = datafile.read()
        # Decompress and load metadata
        edata = self.__decompress(filedata)
        # Check if password matched
        if self.__check_credentials():
            self.__data = self.__decrypt(edata)     # Decrypt and load the data
        else:
            raise AuthenticationException

    def store(self):
        """
        Store the repo back in the file.
            - Call __encrypt to get the data to store and Metadata updates
            - Store that in the file
            Return nothing, update Metadata
        """
        if self.__metadata is None:
            print("Repository was not loaded", file=sys.stderr)
            self.__metadata = dict()
        # Get the data encrypted, update the metadata
        edata = self.__encrypt()
        # Add the credentials
        # This is onl needed for new repositories
        # But let's do this regardless
        self.__store_credentials()
        # Compress the data,
        filedata = self.__compress(edata)
        # Open repository file and dump the new repository contents
        with open(self.__file, 'wb') as datafile:
            datafile.write(filedata)
        # Set data to None, means the repository is no longer usable
        # One must call load() again to make it usable
        self.__data = None

    def __store_credentials(self):
        """ Save the users password bcrypt hash in the Metadata """
        # Calculate the hash of the password
        pwhash = bcrypt.hashpw(bytes(self.__passwd, 'utf-8'), bcrypt.gensalt()).decode('utf-8')
        self.__metadata["PasswordHash"] = pwhash        # Update the metadata

    def __check_credentials(self):
        """ Check if user password is valid - we only proceed to decryption if this passphrases. """
        assert "PasswordHash" in self.__metadata        # Metadata must contain the password hash
        pwhash = bytes(self.__metadata["PasswordHash"], 'utf-8')        # Get the hash
        # Check if hashes matched
        if constant_time.bytes_eq(pwhash, bcrypt.hashpw(bytes(self.__passwd, 'utf-8'), pwhash)):
            return True
        return False

    def is_decrypted(self):
        """ Return true if this repo is currently decrypted """
        # If data is None, the repo has not been load()-ed yet
        return self.__data is not None

    def __getitem__(self, key):
        """
        Get password of account from repository

        Keyword arguments:
            key -- Account name

        Returns:
            Dict containing associated data including password
        """
        # Repository must be decrypted for this to work
        assert self.is_decrypted()
        if key in self.__data:
            return self.__data[key]
        else:
            raise KeyError

    def __setitem__(self, key, item):
        """
        Add new password to repository

        Keyword arguments:
            key -- Account name
            item -- Password of account + associated info
        """
        # Repository must be decrypted for this to work
        assert self.is_decrypted()
        # Get maximum number of accounts a repository can hold
        max_len = int(Repository.config.get('accounts.max', 100))
        if len(self) >= max_len:
            # Too many accounts
            raise ValueError("Too many accounts in repository")
        # Key must be a string
        assert isinstance(key, str)
        # Item must be an object
        if not item:
            raise ValueError("None Type for value")
        # Item must be a dictionary
        assert isinstance(item, dict)
        # The dictionary must have a "Password" field (This is a __Password__ manager after all)
        assert "Password" in item
        # No duplicate key overwriting
        if key in self.__data:
            raise KeyError("Duplicate key")
        self.__data[key] = item

    def __delitem__(self, key):
        """
        Remove a password from this repository

        Keyword arguments:
            key -- Account name
        """
        # Repository must be decrypted for this to work
        assert self.is_decrypted()
        if key not in self.__data:
            raise KeyError("Key not present")
        del self.__data[key]

    def __repr__(self):
        """ Return description string, does not contain any data."""
        return "Encrypted repository: {} loaded from {}".format(self.__name, self.__file)

    def __str__(self):
        """ Return description string, does not caontain any data."""
        return self.__name

    def __len__(self):
        """ Return number of accounts in this repository if decrypted """
        # Repository must be decrypted for this to work
        assert self.is_decrypted()
        return len(self.__data)     # number of entries in data is the length of the repository

    def __enter__(self):
        """ Context management load() """
        self.load()     # Load from file
        return self     # Return self instance to enable operations

    def __exit__(self, etype, value, traceback):
        """ Context management store() """
        self.store()    # Store the repository in the file

    def clear(self):
        """ Delete all accounts if decrypted """
        # Repository must be decrypted for this to work
        assert self.is_decrypted()
        self.__data.clear()     # Clear all data from the repository

    def keys(self):
        """ Return tuple of all accounts in repository. """
        # Repository must be decrypted for this to work
        assert self.is_decrypted()
        return self.__data.keys()      # Return the list of keys

    def values(self):
        """ We will not let users get access to all the values in a single function call. """
        # This is probably not safe
        raise NotImplementedError

    def items(self):
        """ If decrypted, return iterator to all items in the underlying dict. """
        # Repository must be decrypted for this to work
        assert self.is_decrypted()
        return self.__data.items()     # Return all items

    def __contains__(self, key):
        """ Check if repo contains any account with this name """
        # Repository must be decrypted for this to work
        assert self.is_decrypted()
        return key in self.__data       # If key is in __data it is in the repository

    def __iter__(self):
        """ Return generator to all accounts in repo, if decrypted. """
        # Repository must be decrypted for this to work
        for account in self.__data:
            assert self.is_decrypted()
            yield account

if __name__ == '__main__':
    pass
