from Crypto.Cipher import AES

BLOCK_SIZE = AES.block_size


class AESCipher:

    def __init__(self, key, iv):
        try:
            key = bytes.fromhex(key)
        except:
            key = bytes(key, encoding='utf-8')
        iv = bytes(iv[0:16], encoding='utf-8')
        self.cipher = AES.new(key=key, mode=AES.MODE_CBC, IV=iv)

    @staticmethod
    def pad(b):
        return b + (BLOCK_SIZE - len(b) % BLOCK_SIZE) * chr(BLOCK_SIZE - len(b) % BLOCK_SIZE).encode()

    @staticmethod
    def unpad(b):
        return b[:-ord(b[len(b) - 1:])]

    def encrypt(self, text):
        text = self.pad(text)
        encrypted_text = self.cipher.encrypt(text)
        return encrypted_text

    def decrypt(self, encrypted_text):
        decrypted_text = self.cipher.decrypt(encrypted_text)
        return self.unpad(decrypted_text)