# NOTE: Cryptography import
from Crypto.Cipher import AES

# NOTE: Generic imports
import traceback
import logging


logger = logging.getLogger("__main__." + __name__)


class SteelPetal:
    """
    Cryptography utility for encrypting files
    """

    def __init__(self, key: str, nonce=None, testData=None):
        """
        Parameters
        ----------
        key: str
            User's plaintext key

        nonce
            Cryptographic artifact we can use to reinitialize cryptographic operations
        """
        try:
            if nonce is None:
                self.cipher = AES.new(self.pad(key).encode("utf-8"), AES.MODE_GCM)
                self.nonce = self.cipher.nonce
            else:
                self.cipher = AES.new(
                    self.pad(key).encode("utf-8"), AES.MODE_GCM, nonce=nonce
                )
                self.nonce = nonce
            if testData is not None:
                try:
                    self.cipher.decrypt(testData)
                except Exception:
                    logger.error(traceback.format_exc())
                    return False
        except Exception:
            logger.error(traceback.format_exc())

    def pad(self, key: str):
        """
        Pad key to make it usable

        key: str
            User's plain text key
        """
        BS = AES.block_size
        key = key + (BS - len(key) % BS) * chr(BS - len(key) % BS)
        return key

    def encrypt(self, data: bytes):
        """
        Encrypt binary data

        data: bytes
            Data to encrypt
        """
        try:
            return self.cipher.encrypt_and_digest(data)
        except Exception:
            logger.error(traceback.format_exec())
            return False

    def decrypt(self, data: bytes):
        """
        Decrypt encrypted binary data

        data: bytes
            Data to decrypt
        """
        try:
            return self.cipher.decrypt(data)
        except Exception:
            logger.error(traceback.format_exec())
            return False