"""Ciphers module.
Provide Asymmetric and Symmetric cipher functions.\
"""
from enum import IntEnum

from Crypto import Random
from Crypto.Cipher import PKCS1_v1_5, AES
from Crypto.PublicKey import RSA
from Crypto.Random import random as rand
from Crypto.Util import Counter


class AESBlockCipher:

    class Size(IntEnum):
        bit_128 = 1
        bit_192 = 2
        bit_256 = 3

    supported_mode = [AES.MODE_CBC, AES.MODE_CTR]

    def __init__(self, keysize=Size.bit_256, mode=AES.MODE_CBC, secret=None, salt=None):
        """Generate AES key according to the keysize and secret.
        keysize -- int, AES keysize in byte (16, 24, 32).
        secret -- bytes, only used for Password-Based AES cipher.
        """

        if mode not in self.supported_mode:
            raise None  # TODO raise error

        self.aes = None
        self.MODE = mode

        if keysize not in list(map(int, self.Size)):
            raise Exception("Invalid key size.")

        keysize = int((128 + 64 * (keysize - 1)) / 8)

        if not secret:
            """Generate randome key"""
            self.key = Random.get_random_bytes(keysize)
        else:  # Password-Based
            """Construct AES key from user secret with argon2 hash function"""
            import argon2

            self.salt = Random.get_random_bytes(16) if salt is None else salt
            time_cost = 2
            memory_cost = 512
            parallelism = 2
            hash_len = keysize

            self.key = argon2.low_level.hash_secret_raw(secret, self.salt, time_cost,
                                                        memory_cost, parallelism,
                                                        hash_len, argon2.low_level.Type.I)

    def _setup_cipher(self, iv=None):
        # 16-byte, 128-bit initial vector/value for both CBC and CTR mode
        iv = Random.get_random_bytes(AES.block_size) if not iv else iv
        if self.MODE == AES.MODE_CBC:
            return iv, AES.new(self.key, AES.MODE_CBC, iv)
        if self.MODE == AES.MODE_CTR:
            ctr = Counter.new(128, initial_value=int.from_bytes(iv, byteorder="big"))
            return iv, AES.new(self.key, AES.MODE_CTR, counter=ctr)

    @staticmethod
    def pad(b):
        return b + \
               ((AES.block_size - len(b) % AES.block_size) *
                chr(AES.block_size - len(b) % AES.block_size)).encode()

    @staticmethod
    def unpad(b):
        return b[:-b[-1]]

    # Encrypt a chuck of bytes
    def encrypt(self, b, last=False):
        iv = bytes()
        if not self.aes:  # first time, setup iv
            iv, self.aes = self._setup_cipher()

        if not last:
            return iv + self.aes.encrypt(b)
        else:
            # if last chunk, pad the chunk
            ret = self.aes.encrypt(self.pad(b))
            self.aes = None
            return iv + ret  # return iv also since it could be the only chunk

    def decrypt(self, b, last=False):
        if not self.aes:  # first time, retrieve iv
            iv = b[:AES.block_size]  # both modes have the same iv size
            b = b[len(iv):]
            _, self.aes = self._setup_cipher(iv)

        if not last:
            return self.aes.decrypt(b)
        else:
            # if last chunk, unpad the chunk
            dec = self.aes.decrypt(b)
            return self.unpad(dec)


class RSACipher:
    rsa = None
    cipher = None

    class Size(IntEnum):
        bit_1024 = 1
        bit_2048 = 2
        bit_4096 = 3
        bit_8192 = 4

    def import_key(self, key, passphrase=None):
        # if key is bytes, load it directly
        if isinstance(key, bytes):
            self.rsa = RSA.importKey(key, passphrase)
        else:  # else read it as file
            with open(key, "rb") as f:
                self.rsa = RSA.importKey(f.read(), passphrase)

        self.cipher = PKCS1_v1_5.new(self.rsa)

    def export_key(self, keyfile=None, format="PEM", passphrase=None, public_only=False):
        if not hasattr(self, "rsa"):
            raise Exception("Missing RSA cipher...")

        # if has keyfile, write to key file
        if keyfile:
            with open(keyfile, "wb") as f:
                if public_only:
                    pubkey = self.rsa.publickey()
                    f.write(pubkey.exportKey(format=format))
                else:
                    f.write(self.rsa.exportKey(format=format, passphrase=passphrase))
        else:  # else return bytes
            if public_only:
                pubkey = self.rsa.publickey()
                return pubkey.exportKey(format=format)
            else:
                return self.rsa.exportKey(format=format, passphrase=passphrase)

    def generate(self, keysize=Size.bit_4096, progress_func=None):
        self.progress_func = progress_func

        if keysize not in list(map(int, self.Size)):
            raise Exception("Invalid key size.")

        keysize = {1: 1024, 2: 2048, 3: 4096, 4: 8192}[keysize]
        self.rsa = RSA.generate(keysize, progress_func=self.gen_prog)

        self.cipher = PKCS1_v1_5.new(self.rsa)

        self.gen_prog("Finished, key size: " + str(RSACipher.Size(keysize)))

    # progress function for generating key
    def gen_prog(self, s):
        d = {"p,q": "Generating p and q primes...",
             "u": "Generating public exponent...",
             "d": "Generating private exponent..."}
        if self.progress_func:
            if s in d.keys():
                self.progress_func(d[s])
            else:
                self.progress_func(s)

    # encrypt with rsa key, data size should not exceed n - 11
    def encrypt(self, b):
        if not hasattr(self, "rsa"):
            raise Exception("Missing RSA cipher...")

        return self.cipher.encrypt(b)

    def decrypt(self, b):
        if not hasattr(self, "rsa"):
            raise Exception("Missing RSA cipher...")

        return self.cipher.decrypt(b, Random.get_random_bytes(rand.randint(0, self.rsa.size())))
