ska-referee-test/ref-test/common/security/__init__.py
viveksantayana 288ecb60e1 Finished client result API.
Need to work on adjustment user codes and server email notifications.
2021-11-30 18:06:24 +00:00

60 lines
1.9 KiB
Python

from os import environ, path
from cryptography.fernet import Fernet
def generate_keyfile():
with open('./common/security/.encryption.key', 'wb') as keyfile:
key = Fernet.generate_key()
keyfile.write(key)
def load_key():
with open('./common/security/.encryption.key', 'rb') as keyfile:
key = keyfile.read()
return key
def check_keyfile_exists():
return path.isfile('./common/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)