aibot 1.0.2

#![allow(deprecated)]
//! # AES-128-ECB 加解密工具
//!
//! iLink 协议使用 **AES-128-ECB** + **PKCS7 padding** 对所有媒体文件进行加密。
//!
//! ## AES 密钥格式
//!
//! iLink 协议中 `aes_key` 字段可能有多种编码，本模块的 [`parse_aes_key`] 函数自动识别：
//!
//! 1. **32 位十六进制字符串** —— `image_item.aes_key` 的常见格式
//! 2. **Base64 编码的 16 字节原始密钥**
//! 3. **Base64 编码的 32 位十六进制字符串** —— 文件/语音/视频的常见格式
//! 4. **原始 16/24/32 字节** —— 透传方式
//!
//! ## 示例
//!
//! ```rust
//! use aibot::crypto::{parse_aes_key, aes_ecb_decrypt, aes_ecb_encrypt};
//!
//! let key_hex = "0123456789abcdef0123456789abcdef";
//! let key = parse_aes_key(key_hex).unwrap();
//! let data = b"Hello, World!";
//!
//! let encrypted = aes_ecb_encrypt(data, &key).unwrap();
//! let decrypted = aes_ecb_decrypt(&encrypted, key_hex).unwrap();
//! assert_eq!(decrypted, data);
//! ```

use crate::error::{BotError, BotResult};
use aes::cipher::{BlockCipherDecrypt, BlockCipherEncrypt, KeyInit};
use aes::Aes128;
use cipher::Block;
use xbw::crypt::{
    base64::{Base64, Engine},
    hex,
};

/// AES 块大小（16 字节 = 128 位）。
const BLOCK_SIZE: usize = 16;

/// 解析 iLink AES 密钥为 16 字节原始密钥。
///
/// 自动检测编码格式（对标官方 TypeScript 的 `parseAesKey`）：
///
/// | 步骤 | 检测逻辑 | 成功条件 |
/// |------|----------|----------|
/// | 1 | 32/48/64 位十六进制字符串 → hex 解码 | 结果为 16 字节 |
/// | 2 | Base64 解码 | 结果为 16 字节 |
/// | 3 | Base64 解码 → 32 位十六进制 → 再次 hex 解码 | 结果为 16 字节 |
/// | 4 | 直接作为原始字节 | 结果为 16 字节 |
///
/// 如果所有格式都未能得到 16 字节密钥，返回 [`BotError::InvalidData`]。
pub fn parse_aes_key(key_str: &str) -> BotResult<[u8; 16]> {
    let raw = key_str.trim();

    // --- 格式 1：十六进制字符串 ------------------------------------------------
    if [32, 48, 64].contains(&raw.len()) && raw.chars().all(|c| c.is_ascii_hexdigit()) {
        let bytes = hex::decode(raw)?;
        if bytes.len() == 16 {
            return Ok(arraycopy(&bytes));
        }
    }

    // --- 格式 2：Base64 编码 --------------------------------------------------
    let decoded = Base64
        .decode(raw)
        .unwrap_or_else(|_| raw.as_bytes().to_vec());

    // 2a: Base64 解码后直接是 16 字节
    if decoded.len() == 16 {
        return Ok(arraycopy(&decoded));
    }

    // 2b: Base64 解码后是 32 字节十六进制字符串 → 再次 hex 解码
    if decoded.len() == 32 && decoded.iter().all(|&b| b.is_ascii_hexdigit()) {
        let bytes = hex::decode(&decoded)?;
        if bytes.len() == 16 {
            return Ok(arraycopy(&bytes));
        }
    }

    // --- 所有格式均失败 ---------------------------------------------------------
    Err(BotError::invalid_data(format!(
        "无法解析 AES 密钥: raw={:?} len={}",
        &raw.chars().take(20).collect::<String>(),
        raw.len()
    )))
}

/// 使用 AES-128-ECB + PKCS7 解密数据。
///
/// ## 参数
///
/// * `data` — 待解密的加密字节数据（长度必须是 16 的倍数）。
/// * `aes_key` — AES 密钥字符串（支持多种格式，见 [`parse_aes_key`]）。
///
/// ## 返回值
///
/// 解密并去除 PKCS7 padding 后的原始数据。
///
/// ## 错误
///
/// - 如果 `aes_key` 无法解析为 16 字节密钥 → [`BotError::InvalidData`]
/// - 如果 `data` 长度不是 16 的倍数 → [`BotError::InvalidData`]
/// - 如果 PKCS7 padding 无效 → [`BotError::InvalidData`]
pub fn aes_ecb_decrypt(data: &[u8], aes_key: &str) -> BotResult<Vec<u8>> {
    let key = parse_aes_key(aes_key)?;
    let cipher = Aes128::new_from_slice(&key)
        .map_err(|e| BotError::crypto(format!("AES 密钥初始化失败: {:?}", e)))?;

    if data.len() % BLOCK_SIZE != 0 {
        return Err(BotError::invalid_data(format!(
            "加密数据长度 {} 不是块大小 {} 的整数倍",
            data.len(),
            BLOCK_SIZE
        )));
    }

    let mut buf = data.to_vec();

    // ECB 模式：逐块解密
    for chunk in buf.chunks_mut(BLOCK_SIZE) {
        let mut block = Block::<Aes128>::from_mut_slice(chunk);
        cipher.decrypt_block(&mut block);
    }

    // 移除 PKCS7 padding
    let pad_byte = buf.last().copied().unwrap_or(0) as usize;
    if pad_byte == 0 || pad_byte > BLOCK_SIZE {
        return Err(BotError::invalid_data(format!(
            "无效的 PKCS7 填充字节: {}",
            pad_byte
        )));
    }
    // 校验所有填充字节是否正确
    for &b in &buf[buf.len() - pad_byte..] {
        if b as usize != pad_byte {
            return Err(BotError::invalid_data(format!(
                "PKCS7 填充校验失败: 期望 {pad_byte}，实际得到 {b}"
            )));
        }
    }
    buf.truncate(buf.len() - pad_byte);

    Ok(buf)
}

/// 使用 AES-128-ECB + PKCS7 加密数据。
///
/// ## 参数
///
/// * `data` — 待加密的原始字节数据。
/// * `key` — 16 字节的 AES 密钥（已解析好的原始密钥）。
///
/// ## 返回值
///
/// 加密后的数据（长度自动补齐到 16 的倍数）。
pub fn aes_ecb_encrypt(data: &[u8], key: &[u8; 16]) -> BotResult<Vec<u8>> {
    let cipher = Aes128::new_from_slice(key)
        .map_err(|e| BotError::crypto(format!("AES 密钥初始化失败: {:?}", e)))?;

    // PKCS7 填充
    let padding = BLOCK_SIZE - (data.len() % BLOCK_SIZE);
    let padded_len = data.len() + padding;
    let mut buf = vec![0u8; padded_len];
    buf[..data.len()].copy_from_slice(data);
    buf[data.len()..].fill(padding as u8);

    // ECB 模式：逐块加密
    for chunk in buf.chunks_mut(BLOCK_SIZE) {
        let mut block = Block::<Aes128>::from_mut_slice(chunk);
        cipher.encrypt_block(&mut block);
    }

    Ok(buf)
}

// ============================================================================
// 内部辅助函数
// ============================================================================

/// 将切片拷贝为定长 16 字节数组。
fn arraycopy(src: &[u8]) -> [u8; 16] {
    let mut dst = [0u8; 16];
    dst.copy_from_slice(src);
    dst
}