import base64
import hashlib
from Crypto.Cipher import AES
from Crypto import Random
from Crypto.Util.Padding import pad, unpad
from codes.common import clibs


class PassCipher:
    def __init__(self, salt):
        salt = salt.encode("utf-8")
        self.key = hashlib.sha256(salt).digest()

    def encrypt(self, plaintext):
        plaintext = plaintext.encode("utf-8")
        iv = Random.new().read(AES.block_size)
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        ciphertext = cipher.encrypt(pad(plaintext, AES.block_size))
        return base64.b64encode(iv + ciphertext).decode("utf-8")

    def decrypt(self, ciphertext):
        ciphertext = base64.b64decode(ciphertext)
        iv = ciphertext[:AES.block_size]
        ciphertext = ciphertext[AES.block_size:]
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        plaintext = unpad(cipher.decrypt(ciphertext), AES.block_size)
        return plaintext.decode("utf-8")

    @staticmethod
    def gen_salt(text: str):
        key = ""
        passwd = {idx: char for idx, char in enumerate(text * 4)}
        for idx in range(32):
            char_i = 0 if ord(passwd[idx]) - clibs.account["code_dict"][idx] < 0 else ord(passwd[idx]) - clibs.account["code_dict"][idx]
            key += chr(char_i)
        salt = base64.urlsafe_b64encode(key.encode()).decode()
        return salt