# Copyright (c) Huawei Technologies Co., Ltd. 2026-2026. All rights reserved.
# MindIE is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         `http://license.coscl.org.cn/MulanPSL2`
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

import base64
import logging
import os
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes

from olc.encryptor.encryptor import Encryptor

logger = logging.getLogger(__name__)


class AesEncryptor(Encryptor):
    """
    AES 加密器实现，使用 AES-CBC 模式和 PBKDF2 密钥派生。

    特点:
    - 使用 AES-256-CBC 加密算法
    - 通过 PBKDF2HMAC 从密码派生密钥
    - 每次加密使用随机 salt 和 IV，确保相同明文加密结果不同
    - 加密数据格式: [salt(16B)][IV(16B)][ciphertext][PKCS7填充]

    Example:
        >>> encryptor = AesEncryptor("my_password")
        >>> encrypted = encryptor.encrypt("敏感数据")
        >>> decrypted = encryptor.decrypt(encrypted)
        >>> decrypted
        '敏感数据'
    """

    def __init__(self, password: str, salt: bytes = None):
        """
        初始化 AES 加密器。

        Args:
            password: 用于派生加密密钥的密码字符串
            salt: 可选的盐值，用于密钥派生。如果不提供，会自动生成随机盐值。
                  盐值会被包含在加密结果中，解密时会自动提取。
        """
        self.password = password.encode('utf-8')
        # 如果没有提供盐值，生成16字节的随机盐
        self.salt = salt if salt else os.urandom(16)
        # 派生 AES 密钥
        self.key = self._derive_key()

    def _derive_key(self) -> bytes:
        """
        使用 PBKDF2HMAC 从密码派生 AES 密钥。

        通过多次迭代的哈希运算，将用户密码转换为安全的加密密钥，
        防止暴力破解和彩虹表攻击。

        Returns:
            bytes: 32字节的 AES-256 密钥
        """
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),  # 使用 SHA-256 哈希算法
            length=32,  # 生成 256位 (32字节) 密钥用于 AES-256
            salt=self.salt,  # 盐值，防止彩虹表攻击
            iterations=100000,  # 迭代次数，增加破解难度
            backend=default_backend()
        )
        return kdf.derive(self.password)

    def encrypt(self, raw: str) -> str:
        """
        加密字符串。

        Args:
            raw: 要加密的明文字符串

        Returns:
            str: Base64 编码的加密结果，格式为: salt + IV + 密文

        Note:
            每次加密会生成新的随机 IV，因此相同明文多次加密结果不同。
        """
        if not raw:
            return raw

        # 生成16字节的随机初始化向量 (IV)
        iv = os.urandom(16)
        
        # 创建 AES-CBC 加密器
        cipher = Cipher(algorithms.AES(self.key), modes.CBC(iv), backend=default_backend())
        encryptor = cipher.encryptor()

        # 使用 PKCS7 填充，确保明文长度是16字节的倍数
        padder = padding.PKCS7(128).padder()
        padded_data = padder.update(raw.encode('utf-8')) + padder.finalize()

        # 执行加密
        ciphertext = encryptor.update(padded_data) + encryptor.finalize()

        # 组合 salt、IV 和密文，然后进行 Base64 编码
        encrypted_data = self.salt + iv + ciphertext
        return base64.b64encode(encrypted_data).decode('utf-8')

    def decrypt(self, raw: str) -> str:
        """
        解密字符串。

        Args:
            raw: Base64 编码的加密字符串

        Returns:
            str: 解密后的明文字符串

        Raises:
            Exception: 解密失败时抛出异常（如密码错误、数据损坏等）
        """
        try:
            if not raw:
                return raw

            # 从 Base64 解码
            encrypted_data = base64.b64decode(raw)

            # 从加密数据中提取 salt、IV 和密文
            salt = encrypted_data[:16]  # 前16字节是盐值
            iv = encrypted_data[16:32]  # 接下来16字节是 IV
            ciphertext = encrypted_data[32:]  # 剩余部分是密文

            # 使用提取的 salt 重新派生密钥
            kdf = PBKDF2HMAC(
                algorithm=hashes.SHA256(),
                length=32,
                salt=salt,
                iterations=100000,
                backend=default_backend()
            )
            key = kdf.derive(self.password)

            # 创建 AES-CBC 解密器
            cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=default_backend())
            decryptor = cipher.decryptor()

            # 执行解密
            padded_data = decryptor.update(ciphertext) + decryptor.finalize()

            # 移除 PKCS7 填充
            unpadder = padding.PKCS7(128).unpadder()
            data = unpadder.update(padded_data) + unpadder.finalize()
            return data.decode('utf-8')
        except Exception:
            logger.error(f"AES decrypt error, value: {raw}")
            return raw