from os import path
from cryptography.fernet import Fernet

def generate_keyfile():
    with open('./security/.encryption.key', 'wb') as keyfile:
        key = Fernet.generate_key()
        keyfile.write(key)

def load_key():
    with open('./security/.encryption.key', 'rb') as keyfile:
        key = keyfile.read()
        return key

def check_keyfile_exists():
    return path.isfile('./security/.encryption.key')

def encrypt(input:str):
    input = input.encode()
    if not check_keyfile_exists():
        generate_keyfile()
    _encryption_key = load_key()
    fernet = Fernet(_encryption_key)
    output = fernet.encrypt(input)
    return output.decode()

def decrypt(input):
    if not check_keyfile_exists():
        raise EncryptionKeyMissing
    input = input.encode()
    _encryption_key = load_key()
    fernet = Fernet(_encryption_key)
    output = fernet.decrypt(input)
    return output.decode()

class EncryptionKeyMissing(Exception):
    def __init__(self, message='There is no encryption keyfile.'):
        self.message = message
        super().__init__(self.message)