#!/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