from sqlalchemy import create_engine, Column, Integer, String, TypeDecorator
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from cryptography.fernet import Fernet
import base64

# 生成加密密钥 (实际应用中应该安全地存储和管理这个密钥)
key = Fernet.generate_key()
cipher_suite = Fernet(key)

Base = declarative_base()


class EncryptedString(TypeDecorator):
    """自动加密/解密的字符串类型"""

    impl = String  # 底层数据库类型

    def process_bind_param(self, value, dialect):
        """在写入数据库前加密数据"""
        if value is not None:
            # 如果是字符串则编码为bytes
            if isinstance(value, str):
                value = value.encode()
            # 加密并返回base64编码的字符串(便于存储)
            return base64.b64encode(cipher_suite.encrypt(value)).decode('utf-8')
        return value

    def process_result_value(self, value, dialect):
        """从数据库读取后解密数据"""
        if value is not None:
            # 解密base64编码的字符串
            return cipher_suite.decrypt(base64.b64decode(value)).decode('utf-8')
        return value