from Crypto.Cipher import AES
from Crypto import Random
import hashlib
from base64 import b64encode, b64decode

class AEScipher(object):
    def __init__(self, key):
        self.block_size = AES.block_size
        self.key = hashlib.sha256(key.encode()).digest()

    def encrypt(self, plain_text):
        plain_text = self.__pad(plain_text)
        iv = Random.new().read(self.block_size)
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        encrypted_text = cipher.encrypt(plain_text)
        return b64encode(iv + encrypted_text).decode("utf-8")

    def decrypt(self, encrypted_text):
        encrypted_text = b64decode(encrypted_text)
        iv = encrypted_text[:self.block_size]
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        plain_text = cipher.decrypt(encrypted_text[self.block_size:])
        return self.__unpad(plain_text)

    def __pad(self, plain_text):
        print("inside _pad")
        if isinstance(plain_text, str):
            print("encoding in base64")
            plain_text = plain_text.encode('utf-8')
        print("go go")
        number_of_bytes_to_pad = self.block_size - len(plain_text) % self.block_size
        padding = bytes([number_of_bytes_to_pad] * number_of_bytes_to_pad)
        return plain_text + padding

    @staticmethod
    def __unpad(plain_text):
        if isinstance(plain_text, bytes):
            return plain_text[:-plain_text[-1]]
        return plain_text[:-ord(plain_text[-1])]