58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
#!/usr/bin/env python3
|
|
|
|
"""
|
|
AES Encryption methods
|
|
"""
|
|
|
|
__author__ = '@tormakris'
|
|
__copyright__ = "Copyright 2020, onSpot Team"
|
|
__module_name__ = "aes_encrypt"
|
|
__version__text__ = "1"
|
|
|
|
import base64
|
|
import pickle
|
|
|
|
from Crypto.Cipher import AES
|
|
|
|
from flaskaddons.fred import flaskred
|
|
|
|
|
|
class AESCrypto:
|
|
def __init__(self, encoded_secret_key: str, padding_character: bytes = '{'.encode('ascii')):
|
|
self.padding_character = padding_character
|
|
self.encoded_secret_key = encoded_secret_key
|
|
|
|
def encrypt_message(self, private_msg: str) -> tuple:
|
|
secret_key = base64.b64decode(self.encoded_secret_key)
|
|
cipher = AES.new(secret_key, AES.MODE_EAX)
|
|
ciphertext, tag = cipher.encrypt_and_digest(private_msg.encode('UTF-8'))
|
|
return cipher.nonce, ciphertext, tag
|
|
|
|
def decrypt_message(self, nonce: bytes, encoded_encrypted_msg: bytes, tag: bytes) -> str:
|
|
secret_key = base64.b64decode(self.encoded_secret_key)
|
|
cipher = AES.new(secret_key, AES.MODE_EAX, nonce)
|
|
msg = cipher.decrypt_and_verify(encoded_encrypted_msg, tag).decode('UTF-8')
|
|
return msg
|
|
|
|
|
|
class EncryptedUserRedis:
|
|
|
|
def __init__(self, encoded_secret_key: str):
|
|
self.aes = AESCrypto(encoded_secret_key)
|
|
|
|
def store(self, user: dict) -> None:
|
|
nonce, ciphertext, tag = self.aes.encrypt_message(user['password'])
|
|
user['nonce'] = nonce
|
|
user['ciphertext'] = ciphertext
|
|
user['tag'] = tag
|
|
user.pop('password', None)
|
|
flaskred.set(user['name'], pickle.dumps(user))
|
|
|
|
def load(self, username: str) -> dict:
|
|
encrypteddict = pickle.loads(flaskred.get(username))
|
|
|
|
plaindict = {"name": encrypteddict['name'],
|
|
"password": self.aes.decrypt_message(encrypteddict['nonce'], encrypteddict['ciphertext'],
|
|
encrypteddict['tag'])}
|
|
return plaindict
|