eric / rustwarden

use crate::error::AppError;
use aes_gcm::{
    AeadCore, Aes256Gcm, Nonce,
    aead::{Aead, KeyInit, OsRng},
};
use base64::{Engine, engine::general_purpose::STANDARD as B64};
use std::env;

const ENCRYPTED_PREFIX: &str = "enc:";

/// Service providing AES-256-GCM encryption and decryption for sensitive settings.
pub struct CryptoService;

impl CryptoService {
    fn get_key() -> Result<[u8; 32], AppError> {
        let key_hex = env::var("ENCRYPTION_KEY")
            .or_else(|_| env::var("OAUTH_ENCRYPTION_KEY"))
            .map_err(|_| {
            AppError::Internal(
                "ENCRYPTION_KEY (or OAUTH_ENCRYPTION_KEY) is not set. A 64-character hex key is required for encryption.".to_string(),
            )
        })?;

        let bytes = hex::decode(key_hex.trim())
            .map_err(|e| AppError::Internal(format!("ENCRYPTION_KEY is not valid hex: {}", e)))?;

        if bytes.len() != 32 {
            return Err(AppError::Internal(format!(
                "ENCRYPTION_KEY must be exactly 64 hex characters (32 bytes), got {} bytes",
                bytes.len()
            )));
        }

        let mut key = [0u8; 32];
        key.copy_from_slice(&bytes);
        Ok(key)
    }

    /// Encrypt a plaintext string. Returns `enc:<base64(nonce || ciphertext)>`.
    pub fn encrypt(plaintext: &str) -> Result<String, AppError> {
        let key_bytes = Self::get_key()?;
        let cipher = Aes256Gcm::new_from_slice(&key_bytes)
            .map_err(|e| AppError::Internal(format!("Failed to create cipher: {}", e)))?;

        let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
        let ciphertext = cipher
            .encrypt(&nonce, plaintext.as_bytes())
            .map_err(|e| AppError::Internal(format!("Encryption failed: {}", e)))?;

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

        Ok(format!("{}{}", ENCRYPTED_PREFIX, B64.encode(&combined)))
    }

    /// Decrypt a value previously encrypted by [`Self::encrypt`]. Handles both
    /// `enc:` prefixed (encrypted) and plain text (legacy/migration) values.
    pub fn decrypt(value: &str) -> Result<String, AppError> {
        let encoded = match value.strip_prefix(ENCRYPTED_PREFIX) {
            Some(s) => s,
            None => return Ok(value.to_string()),
        };

        let key_bytes = Self::get_key()?;
        let cipher = Aes256Gcm::new_from_slice(&key_bytes)
            .map_err(|e| AppError::Internal(format!("Failed to create cipher: {}", e)))?;

        let combined = B64
            .decode(encoded)
            .map_err(|e| AppError::Internal(format!("Failed to decode encrypted value: {}", e)))?;

        if combined.len() < 12 {
            return Err(AppError::Internal("Encrypted value too short".to_string()));
        }

        let (nonce_bytes, ciphertext) = combined.split_at(12);
        let nonce_array: [u8; 12] = nonce_bytes
            .try_into()
            .map_err(|_| AppError::Internal("Invalid nonce length".to_string()))?;
        let nonce = Nonce::from(nonce_array);

        let plaintext = cipher.decrypt(&nonce, ciphertext).map_err(|_| {
            AppError::Internal("Decryption failed – wrong key or corrupted data".to_string())
        })?;

        String::from_utf8(plaintext)
            .map_err(|e| AppError::Internal(format!("Decrypted value is not valid UTF-8: {}", e)))
    }

    /// Returns true if an encryption key is configured.
    pub fn is_available() -> bool {
        env::var("ENCRYPTION_KEY").is_ok() || env::var("OAUTH_ENCRYPTION_KEY").is_ok()
    }

    /// Returns true if the value is already encrypted (has the `enc:` prefix).
    pub fn is_encrypted(value: &str) -> bool {
        value.starts_with(ENCRYPTED_PREFIX)
    }
}

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

    /// Roundtrip encryption / decryption with a fixed test key.
    ///
    /// Mutating `OAUTH_ENCRYPTION_KEY` via `env::set_var` is **process-wide**.
    /// Other tests in this binary (notably in `totp_service`) also write
    /// `ENCRYPTION_KEY`, and any concurrent reader of either variable can
    /// observe the wrong value mid-flight. The previous "race is benign"
    /// comment was incorrect: an intermittent failure was reported on
    /// 2026-05-18 when this test ran alongside the broader workspace suite.
    ///
    /// The `#[serial(encryption_key)]` attribute (from `serial_test`)
    /// serialises every test annotated with the same key, eliminating the
    /// race without requiring `--test-threads=1`.
    #[test]
    #[serial_test::serial(encryption_key)]
    fn test_roundtrip() {
        // SAFETY: env mutation is serialised across the suite by
        // `#[serial(encryption_key)]`, so no other thread reads the var
        // while this block runs.
        unsafe {
            env::set_var(
                "OAUTH_ENCRYPTION_KEY",
                "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef",
            );
        }
        let secret = "my-super-secret-client-key";
        let encrypted = CryptoService::encrypt(secret).unwrap();
        assert!(encrypted.starts_with("enc:"));
        let decrypted = CryptoService::decrypt(&encrypted).unwrap();
        assert_eq!(decrypted, secret);
    }

    #[test]
    fn test_plaintext_passthrough() {
        let plain = "not-encrypted-value";
        let result = CryptoService::decrypt(plain).unwrap();
        assert_eq!(result, plain);
    }
}