ska-referee-test/ref-test/common/security/__init__.py

from os import environ, 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):
    if not check_keyfile_exists():
        generate_keyfile()
    _encryption_key = load_key()
    fernet = Fernet(_encryption_key)
    if type(input) == str:
        input = input.encode()
        output = fernet.encrypt(input)
        return output.decode()
    if type(input) == dict:
        output = {}
        for key,value in input.items():
            if type(value) == dict:
                output[key] = encrypt(value)
            else:
                value = value.encode()
                output[key] = fernet.encrypt(value)
                output[key] = output[key].decode()
        return output

def decrypt(input):
    if not check_keyfile_exists():
        raise EncryptionKeyMissing
    _encryption_key = load_key()
    fernet = Fernet(_encryption_key)
    if type(input) == str:
        input = input.encode()
        output = fernet.decrypt(input)
        return output.decode()
    if type(input) == dict:
        output = {}
        for key, value in input.items():
            if type(value) == dict:
                output[key] = decrypt(value)
            else:
                value = value.encode()
                output[key] = fernet.decrypt(value)
                output[key] = output[key].decode()
        return output

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