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


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")


# if __name__ == "__main__":
#     salt = "my_secret_salt_string"
#     cipher = PassCipher(salt)
#     original_text = "这是一段需要加密的敏感信息"
#     encrypted_text = cipher.encrypt(original_text)
#     print(f"加密后的文本: {encrypted_text}")
#     decrypted_text = cipher.decrypt(encrypted_text)
#     print(f"解密后的文本: {decrypted_text}")
#     print(f"加解密是否成功: {original_text == decrypted_text}")