zclaw_openfang/crates/zclaw-saas/src/crypto.rs

//! 通用加密工具 (AES-256-GCM)
//!
//! 提供 API Key、TOTP secret 等敏感数据的加密/解密。
//! 存储格式: `enc:<base64(nonce(12 bytes) || ciphertext)>`

use aes_gcm::aead::{Aead, KeyInit, OsRng};
use aes_gcm::aead::rand_core::RngCore;
use aes_gcm::{Aes256Gcm, Nonce};
use crate::error::{SaasError, SaasResult};

/// 启动时迁移所有旧格式 TOTP secret（明文或固定 nonce → 随机 nonce `enc:` 格式）
///
/// 查找 `totp_secret IS NOT NULL AND totp_secret != '' AND totp_secret NOT LIKE 'enc:%'` 的行，
/// 用当前 AES-256-GCM 密钥加密后写回。
pub async fn migrate_legacy_totp_secrets(pool: &sqlx::PgPool, enc_key: &[u8; 32]) -> anyhow::Result<u32> {
    let rows: Vec<(String, String)> = sqlx::query_as(
        "SELECT id, totp_secret FROM accounts WHERE totp_secret IS NOT NULL AND totp_secret != '' AND totp_secret NOT LIKE 'enc:%'"
    )
    .fetch_all(pool)
    .await?;

    let count = rows.len() as u32;
    for (account_id, plaintext_secret) in &rows {
        let encrypted = encrypt_value(plaintext_secret, enc_key)?;
        sqlx::query("UPDATE accounts SET totp_secret = $1 WHERE id = $2")
            .bind(&encrypted)
            .bind(account_id)
            .execute(pool)
            .await?;
    }
    if count > 0 {
        tracing::info!("Migrated {} legacy TOTP secrets to encrypted format", count);
    }
    Ok(count)
}

/// 加密值的前缀标识
pub const ENCRYPTED_PREFIX: &str = "enc:";

/// AES-256-GCM nonce 长度 (12 字节)
const NONCE_SIZE: usize = 12;

/// 加密明文值 (AES-256-GCM, 随机 nonce)
///
/// 返回格式: `enc:<base64(nonce(12 bytes) || ciphertext)>`
/// 每次加密使用随机 nonce，相同明文产生不同密文。
pub fn encrypt_value(plaintext: &str, key: &[u8; 32]) -> SaasResult<String> {
    let cipher = Aes256Gcm::new_from_slice(key)
        .map_err(|e| SaasError::Encryption(format!("加密初始化失败: {}", e)))?;

    let mut nonce_bytes = [0u8; NONCE_SIZE];
    OsRng.fill_bytes(&mut nonce_bytes);
    let nonce = Nonce::from_slice(&nonce_bytes);

    let ciphertext = cipher.encrypt(nonce, plaintext.as_bytes())
        .map_err(|e| SaasError::Encryption(format!("加密失败: {}", e)))?;

    let mut combined = nonce_bytes.to_vec();
    combined.extend_from_slice(&ciphertext);

    Ok(format!("{}{}", ENCRYPTED_PREFIX, data_encoding::BASE64.encode(&combined)))
}

/// 解密 `enc:` 前缀的加密值
///
/// 仅支持新格式 (随机 nonce)，不支持旧格式 (固定 nonce)。
/// 旧格式数据应通过一次性迁移函数转换。
pub fn decrypt_value(encrypted: &str, key: &[u8; 32]) -> SaasResult<String> {
    let encoded = encrypted.strip_prefix(ENCRYPTED_PREFIX)
        .ok_or_else(|| SaasError::Encryption("加密值格式无效 (缺少 enc: 前缀)".into()))?;

    let raw = data_encoding::BASE64.decode(encoded.as_bytes())
        .map_err(|_| SaasError::Encryption("加密值 Base64 解码失败".into()))?;

    if raw.len() <= NONCE_SIZE {
        return Err(SaasError::Encryption("加密值数据不完整".into()));
    }

    let cipher = Aes256Gcm::new_from_slice(key)
        .map_err(|e| SaasError::Encryption(format!("解密初始化失败: {}", e)))?;

    let (nonce_bytes, ciphertext) = raw.split_at(NONCE_SIZE);
    let nonce = Nonce::from_slice(nonce_bytes);

    let plaintext = cipher.decrypt(nonce, ciphertext)
        .map_err(|_| SaasError::Encryption("解密失败 (密钥可能已变更)".into()))?;

    String::from_utf8(plaintext)
        .map_err(|_| SaasError::Encryption("解密后数据无效 UTF-8".into()))
}

/// 检查值是否已加密 (以 `enc:` 开头)
pub fn is_encrypted(value: &str) -> bool {
    value.starts_with(ENCRYPTED_PREFIX)
}

/// 批量迁移: 将旧的固定 nonce 加密值重新加密为随机 nonce 格式
///
/// 输入为旧格式 (固定 nonce `zclaw_totp_nce`) 加密的 base64 数据，
/// 输出为新格式 `enc:<base64(random_nonce || ciphertext)>`。
pub fn re_encrypt_from_legacy(legacy_base64: &str, legacy_key: &[u8; 32], new_key: &[u8; 32]) -> SaasResult<String> {
    // 先用旧 nonce 解密
    let cipher = Aes256Gcm::new_from_slice(legacy_key)
        .map_err(|e| SaasError::Encryption(format!("解密初始化失败: {}", e)))?;

    let raw = data_encoding::BASE64.decode(legacy_base64.as_bytes())
        .or_else(|_| data_encoding::BASE32.decode(legacy_base64.as_bytes()))
        .map_err(|_| SaasError::Encryption("旧格式 Base64/Base32 解码失败".into()))?;

    // 尝试新格式 (前 12 字节为 nonce)
    if raw.len() > NONCE_SIZE {
        let (nonce_bytes, ciphertext) = raw.split_at(NONCE_SIZE);
        let nonce = Nonce::from_slice(nonce_bytes);
        if let Ok(plaintext_bytes) = cipher.decrypt(nonce, ciphertext) {
            let plaintext = String::from_utf8(plaintext_bytes)
                .map_err(|_| SaasError::Encryption("旧格式解密后数据无效".into()))?;
            return encrypt_value(&plaintext, new_key);
        }
    }

    // 回退到旧格式: 固定 nonce
    let legacy_nonce = Nonce::from_slice(b"zclaw_totp_nce");
    let plaintext_bytes = cipher.decrypt(legacy_nonce, raw.as_ref())
        .map_err(|_| SaasError::Encryption("旧格式解密失败".into()))?;
    let plaintext = String::from_utf8(plaintext_bytes)
        .map_err(|_| SaasError::Encryption("旧格式解密后数据无效".into()))?;

    encrypt_value(&plaintext, new_key)
}