hms/crates/erp-health/src/crypto.rs

use aes_gcm::aead::Aead;
use aes_gcm::{Aes256Gcm, KeyInit, Nonce};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use hmac::{Hmac, Mac};
use sha2::Sha256;
use zeroize::Zeroizing;

use erp_core::error::{AppError, AppResult};

type HmacSha256 = Hmac<Sha256>;

const CIPHER_VERSION: &str = "v1";

pub struct HealthCrypto {
    aes_key: Zeroizing<[u8; 32]>,
    hmac_key: Zeroizing<[u8; 32]>,
}

impl HealthCrypto {
    pub fn from_keys(aes_key_hex: &str, hmac_key_hex: &str) -> AppResult<Self> {
        let aes_key = hex::decode(aes_key_hex)
            .map_err(|e| AppError::Internal(format!("AES key hex decode failed: {}", e)))?;
        let hmac_key = hex::decode(hmac_key_hex)
            .map_err(|e| AppError::Internal(format!("HMAC key hex decode failed: {}", e)))?;
        if aes_key.len() != 32 || hmac_key.len() != 32 {
            return Err(AppError::Internal(
                "Encryption keys must be 32 bytes each".into(),
            ));
        }
        let aes = Zeroizing::new({
            let mut a = [0u8; 32];
            a.copy_from_slice(&aes_key);
            a
        });
        let hmac = Zeroizing::new({
            let mut h = [0u8; 32];
            h.copy_from_slice(&hmac_key);
            h
        });
        Ok(Self {
            aes_key: aes,
            hmac_key: hmac,
        })
    }

    /// Dev fallback: derive deterministic keys from a single dev string.
    /// DO NOT use in production.
    pub fn dev_default() -> Self {
        use sha2::Digest;
        let aes_key = <Sha256 as Digest>::digest(b"erp-health-aes-dev-key-DO-NOT-USE-IN-PROD");
        let hmac_key = <Sha256 as Digest>::digest(b"erp-health-hmac-dev-key-DO-NOT-USE-IN-PROD");
        let aes = Zeroizing::new({
            let mut a = [0u8; 32];
            a.copy_from_slice(&aes_key);
            a
        });
        let hmac = Zeroizing::new({
            let mut h = [0u8; 32];
            h.copy_from_slice(&hmac_key);
            h
        });
        Self {
            aes_key: aes,
            hmac_key: hmac,
        }
    }

    pub fn encrypt(&self, plaintext: &str) -> AppResult<String> {
        let cipher = Aes256Gcm::new_from_slice(&*self.aes_key)
            .map_err(|e| AppError::Internal(format!("AES init failed: {}", e)))?;
        let nonce_bytes = uuid::Uuid::now_v7();
        let nonce = Nonce::from_slice(&nonce_bytes.as_bytes()[..12]);
        let ciphertext = cipher
            .encrypt(nonce, plaintext.as_bytes())
            .map_err(|e| AppError::Internal(format!("Encryption failed: {}", e)))?;
        let mut combined = nonce_bytes.as_bytes()[..12].to_vec();
        combined.extend_from_slice(&ciphertext);
        Ok(format!("{}|{}", CIPHER_VERSION, BASE64.encode(&combined)))
    }

    pub fn decrypt(&self, encoded: &str) -> AppResult<String> {
        // 剥离版本前缀（兼容旧格式无前缀）
        let b64 = if let Some(pos) = encoded.find('|') {
            &encoded[pos + 1..]
        } else {
            encoded
        };
        let combined = BASE64
            .decode(b64)
            .map_err(|e| AppError::Internal(format!("Base64 decode failed: {}", e)))?;
        if combined.len() < 12 {
            return Err(AppError::Internal("Ciphertext too short".into()));
        }
        let (nonce_bytes, ciphertext) = combined.split_at(12);
        let cipher = Aes256Gcm::new_from_slice(&*self.aes_key)
            .map_err(|e| AppError::Internal(format!("AES init failed: {}", e)))?;
        let plaintext = cipher
            .decrypt(Nonce::from_slice(nonce_bytes), ciphertext)
            .map_err(|e| AppError::Internal(format!("Decryption failed: {}", e)))?;
        String::from_utf8(plaintext)
            .map_err(|e| AppError::Internal(format!("UTF-8 decode failed: {}", e)))
    }

    pub fn hmac_hash(&self, value: &str) -> String {
        let mut mac = <HmacSha256 as hmac::Mac>::new_from_slice(&*self.hmac_key)
            .expect("HMAC key length is valid");
        mac.update(value.as_bytes());
        hex::encode(mac.finalize().into_bytes())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_crypto() -> HealthCrypto {
        HealthCrypto::dev_default()
    }

    #[test]
    fn encrypt_decrypt_roundtrip() {
        let crypto = test_crypto();
        let plaintext = "110101199001011234";
        let encrypted = crypto.encrypt(plaintext).unwrap();
        let decrypted = crypto.decrypt(&encrypted).unwrap();
        assert_eq!(plaintext, decrypted);
    }

    #[test]
    fn encrypt_produces_different_ciphertexts() {
        let crypto = test_crypto();
        let plaintext = "110101199001011234";
        let e1 = crypto.encrypt(plaintext).unwrap();
        let e2 = crypto.encrypt(plaintext).unwrap();
        assert_ne!(e1, e2); // 不同 nonce 导致不同密文
    }

    #[test]
    fn decrypt_wrong_key_fails() {
        let crypto1 = HealthCrypto::dev_default();
        let hex_key = "00".repeat(32); // 64 个 0
        let crypto2 = HealthCrypto::from_keys(&hex_key, &hex_key).unwrap();
        let encrypted = crypto1.encrypt("test").unwrap();
        assert!(crypto2.decrypt(&encrypted).is_err());
    }

    #[test]
    fn hmac_hash_deterministic() {
        let crypto = test_crypto();
        let hash1 = crypto.hmac_hash("110101199001011234");
        let hash2 = crypto.hmac_hash("110101199001011234");
        assert_eq!(hash1, hash2);
    }

    #[test]
    fn hmac_hash_different_inputs() {
        let crypto = test_crypto();
        let h1 = crypto.hmac_hash("123456789012345678");
        let h2 = crypto.hmac_hash("987654321098765432");
        assert_ne!(h1, h2);
    }

    #[test]
    fn encrypt_empty_string() {
        let crypto = test_crypto();
        let encrypted = crypto.encrypt("").unwrap();
        let decrypted = crypto.decrypt(&encrypted).unwrap();
        assert_eq!("", decrypted);
    }

    #[test]
    fn decrypt_too_short_fails() {
        let crypto = test_crypto();
        let short = BASE64.encode(b"short");
        assert!(crypto.decrypt(&format!("v1|{}", short)).is_err());
    }

    #[test]
    fn from_keys_invalid_hex() {
        let result = HealthCrypto::from_keys("not-hex", "not-hex");
        assert!(result.is_err());
    }

    #[test]
    fn from_keys_wrong_length() {
        let result = HealthCrypto::from_keys("ab", "cd");
        assert!(result.is_err());
    }

    #[test]
    fn encrypt_has_version_prefix() {
        let crypto = test_crypto();
        let encrypted = crypto.encrypt("test").unwrap();
        assert!(encrypted.starts_with("v1|"), "密文应以 v1| 开头");
    }

    #[test]
    fn decrypt_legacy_no_prefix() {
        let crypto = test_crypto();
        // 手动构造旧格式密文（无版本前缀）
        let cipher = Aes256Gcm::new_from_slice(&*crypto.aes_key).unwrap();
        let nonce_bytes = uuid::Uuid::now_v7();
        let nonce = Nonce::from_slice(&nonce_bytes.as_bytes()[..12]);
        let ciphertext = cipher.encrypt(nonce, b"legacy-test".as_ref()).unwrap();
        let mut combined = nonce_bytes.as_bytes()[..12].to_vec();
        combined.extend_from_slice(&ciphertext);
        let legacy = BASE64.encode(&combined);

        // 旧格式（无 v1| 前缀）仍可解密
        let decrypted = crypto.decrypt(&legacy).unwrap();
        assert_eq!("legacy-test", decrypted);
    }
}