oxigdal-security 0.1.7

//! Data encryption at rest.

use crate::encryption::{EncryptedData, EncryptionAlgorithm, EncryptionMetadata};
use crate::error::{Result, SecurityError};
use aes_gcm::{
    Aes256Gcm, Nonce,
    aead::{Aead, KeyInit},
};
use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
use chacha20poly1305::ChaCha20Poly1305;

/// Encryptor for data at rest.
pub struct AtRestEncryptor {
    algorithm: EncryptionAlgorithm,
    key: Vec<u8>,
    key_id: String,
}

impl AtRestEncryptor {
    /// Create a new encryptor with the given algorithm and key.
    pub fn new(algorithm: EncryptionAlgorithm, key: Vec<u8>, key_id: String) -> Result<Self> {
        // Validate key length
        let required_length = match algorithm {
            EncryptionAlgorithm::Aes256Gcm => 32,
            EncryptionAlgorithm::ChaCha20Poly1305 => 32,
        };

        if key.len() != required_length {
            return Err(SecurityError::encryption(format!(
                "Invalid key length: expected {}, got {}",
                required_length,
                key.len()
            )));
        }

        Ok(Self {
            algorithm,
            key,
            key_id,
        })
    }

    /// Generate a random key for the given algorithm.
    ///
    /// # Errors
    ///
    /// Returns [`SecurityError::Encryption`] if the OS RNG fails to fill
    /// the key buffer (e.g. a sandboxed/restricted environment lacking
    /// `/dev/urandom` or the `getrandom` syscall).
    pub fn generate_key(algorithm: EncryptionAlgorithm) -> Result<Vec<u8>> {
        let mut key = vec![
            0u8;
            match algorithm {
                EncryptionAlgorithm::Aes256Gcm => 32,
                EncryptionAlgorithm::ChaCha20Poly1305 => 32,
            }
        ];
        getrandom::fill(&mut key)
            .map_err(|e| SecurityError::encryption(format!("Failed to generate key: {}", e)))?;
        Ok(key)
    }

    /// Encrypt data.
    pub fn encrypt(&self, plaintext: &[u8], aad: Option<&[u8]>) -> Result<EncryptedData> {
        match self.algorithm {
            EncryptionAlgorithm::Aes256Gcm => self.encrypt_aes_gcm(plaintext, aad),
            EncryptionAlgorithm::ChaCha20Poly1305 => self.encrypt_chacha(plaintext, aad),
        }
    }

    /// Decrypt data.
    pub fn decrypt(&self, encrypted: &EncryptedData) -> Result<Vec<u8>> {
        // Verify algorithm matches
        if encrypted.metadata.algorithm != self.algorithm {
            return Err(SecurityError::decryption(format!(
                "Algorithm mismatch: expected {:?}, got {:?}",
                self.algorithm, encrypted.metadata.algorithm
            )));
        }

        // Verify key ID matches
        if encrypted.metadata.key_id != self.key_id {
            return Err(SecurityError::decryption(format!(
                "Key ID mismatch: expected {}, got {}",
                self.key_id, encrypted.metadata.key_id
            )));
        }

        match self.algorithm {
            EncryptionAlgorithm::Aes256Gcm => self.decrypt_aes_gcm(encrypted),
            EncryptionAlgorithm::ChaCha20Poly1305 => self.decrypt_chacha(encrypted),
        }
    }

    fn encrypt_aes_gcm(&self, plaintext: &[u8], aad: Option<&[u8]>) -> Result<EncryptedData> {
        let cipher = Aes256Gcm::new_from_slice(&self.key)
            .map_err(|e| SecurityError::encryption(format!("Failed to create cipher: {}", e)))?;

        // Generate random nonce
        let mut nonce_bytes = [0u8; 12];
        getrandom::fill(&mut nonce_bytes)
            .map_err(|e| SecurityError::encryption(format!("Failed to generate nonce: {}", e)))?;
        let nonce = Nonce::from(nonce_bytes);

        // Encrypt with optional AAD
        let ciphertext = if let Some(aad_data) = aad {
            cipher
                .encrypt(
                    &nonce,
                    aes_gcm::aead::Payload {
                        msg: plaintext,
                        aad: aad_data,
                    },
                )
                .map_err(|e| SecurityError::encryption(format!("Encryption failed: {}", e)))?
        } else {
            cipher
                .encrypt(&nonce, plaintext)
                .map_err(|e| SecurityError::encryption(format!("Encryption failed: {}", e)))?
        };

        let metadata = EncryptionMetadata::new(
            self.algorithm,
            self.key_id.clone(),
            nonce_bytes.to_vec(),
            aad.map(|a| a.to_vec()),
        );

        Ok(EncryptedData::new(ciphertext, metadata))
    }

    fn decrypt_aes_gcm(&self, encrypted: &EncryptedData) -> Result<Vec<u8>> {
        let cipher = Aes256Gcm::new_from_slice(&self.key)
            .map_err(|e| SecurityError::decryption(format!("Failed to create cipher: {}", e)))?;

        if encrypted.metadata.iv.len() != 12 {
            return Err(SecurityError::decryption(format!(
                "Invalid nonce length: expected 12, got {}",
                encrypted.metadata.iv.len()
            )));
        }

        let nonce = Nonce::try_from(encrypted.metadata.iv.as_slice())
            .map_err(|_| SecurityError::decryption("Invalid nonce bytes"))?;

        let plaintext = if let Some(ref aad) = encrypted.metadata.aad {
            cipher
                .decrypt(
                    &nonce,
                    aes_gcm::aead::Payload {
                        msg: &encrypted.ciphertext,
                        aad,
                    },
                )
                .map_err(|e| SecurityError::decryption(format!("Decryption failed: {}", e)))?
        } else {
            cipher
                .decrypt(&nonce, encrypted.ciphertext.as_ref())
                .map_err(|e| SecurityError::decryption(format!("Decryption failed: {}", e)))?
        };

        Ok(plaintext)
    }

    fn encrypt_chacha(&self, plaintext: &[u8], aad: Option<&[u8]>) -> Result<EncryptedData> {
        let cipher = ChaCha20Poly1305::new_from_slice(&self.key)
            .map_err(|e| SecurityError::encryption(format!("Failed to create cipher: {}", e)))?;

        // Generate random nonce
        let mut nonce_bytes = [0u8; 12];
        getrandom::fill(&mut nonce_bytes)
            .map_err(|e| SecurityError::encryption(format!("Failed to generate nonce: {}", e)))?;
        let nonce = chacha20poly1305::Nonce::from(nonce_bytes);

        // Encrypt with optional AAD
        let ciphertext = if let Some(aad_data) = aad {
            cipher
                .encrypt(
                    &nonce,
                    chacha20poly1305::aead::Payload {
                        msg: plaintext,
                        aad: aad_data,
                    },
                )
                .map_err(|e| SecurityError::encryption(format!("Encryption failed: {}", e)))?
        } else {
            cipher
                .encrypt(&nonce, plaintext)
                .map_err(|e| SecurityError::encryption(format!("Encryption failed: {}", e)))?
        };

        let metadata = EncryptionMetadata::new(
            self.algorithm,
            self.key_id.clone(),
            nonce_bytes.to_vec(),
            aad.map(|a| a.to_vec()),
        );

        Ok(EncryptedData::new(ciphertext, metadata))
    }

    fn decrypt_chacha(&self, encrypted: &EncryptedData) -> Result<Vec<u8>> {
        let cipher = ChaCha20Poly1305::new_from_slice(&self.key)
            .map_err(|e| SecurityError::decryption(format!("Failed to create cipher: {}", e)))?;

        if encrypted.metadata.iv.len() != 12 {
            return Err(SecurityError::decryption(format!(
                "Invalid nonce length: expected 12, got {}",
                encrypted.metadata.iv.len()
            )));
        }

        let nonce = chacha20poly1305::Nonce::try_from(encrypted.metadata.iv.as_slice())
            .map_err(|_| SecurityError::decryption("Invalid nonce bytes"))?;

        let plaintext = if let Some(ref aad) = encrypted.metadata.aad {
            cipher
                .decrypt(
                    &nonce,
                    chacha20poly1305::aead::Payload {
                        msg: &encrypted.ciphertext,
                        aad,
                    },
                )
                .map_err(|e| SecurityError::decryption(format!("Decryption failed: {}", e)))?
        } else {
            cipher
                .decrypt(&nonce, encrypted.ciphertext.as_ref())
                .map_err(|e| SecurityError::decryption(format!("Decryption failed: {}", e)))?
        };

        Ok(plaintext)
    }

    /// Encrypt data in place (overwrites input buffer).
    pub fn encrypt_in_place(
        &self,
        buffer: &mut Vec<u8>,
        aad: Option<&[u8]>,
    ) -> Result<EncryptionMetadata> {
        let encrypted = self.encrypt(buffer, aad)?;
        buffer.clear();
        buffer.extend_from_slice(&encrypted.ciphertext);
        Ok(encrypted.metadata)
    }

    /// Get the algorithm used by this encryptor.
    pub fn algorithm(&self) -> EncryptionAlgorithm {
        self.algorithm
    }

    /// Get the key ID.
    pub fn key_id(&self) -> &str {
        &self.key_id
    }
}

/// Field-level encryptor for encrypting specific fields in structured data.
pub struct FieldEncryptor {
    encryptor: AtRestEncryptor,
}

impl FieldEncryptor {
    /// Create a new field encryptor.
    pub fn new(encryptor: AtRestEncryptor) -> Self {
        Self { encryptor }
    }

    /// Encrypt a string field.
    pub fn encrypt_string(&self, value: &str) -> Result<String> {
        let encrypted = self.encryptor.encrypt(value.as_bytes(), None)?;
        let json = serde_json::to_string(&encrypted)?;
        Ok(BASE64.encode(json))
    }

    /// Decrypt a string field.
    pub fn decrypt_string(&self, encrypted: &str) -> Result<String> {
        let json = BASE64
            .decode(encrypted)
            .map_err(|e| SecurityError::decryption(format!("Base64 decode failed: {}", e)))?;
        let encrypted_data: EncryptedData = serde_json::from_slice(&json)?;
        let plaintext = self.encryptor.decrypt(&encrypted_data)?;
        String::from_utf8(plaintext)
            .map_err(|e| SecurityError::decryption(format!("UTF-8 decode failed: {}", e)))
    }

    /// Encrypt a JSON-serializable value.
    pub fn encrypt_json<T: serde::Serialize>(&self, value: &T) -> Result<String> {
        let json = serde_json::to_vec(value)?;
        let encrypted = self.encryptor.encrypt(&json, None)?;
        let encrypted_json = serde_json::to_string(&encrypted)?;
        Ok(BASE64.encode(encrypted_json))
    }

    /// Decrypt a JSON-serializable value.
    pub fn decrypt_json<T: serde::de::DeserializeOwned>(&self, encrypted: &str) -> Result<T> {
        let json = BASE64
            .decode(encrypted)
            .map_err(|e| SecurityError::decryption(format!("Base64 decode failed: {}", e)))?;
        let encrypted_data: EncryptedData = serde_json::from_slice(&json)?;
        let plaintext = self.encryptor.decrypt(&encrypted_data)?;
        serde_json::from_slice(&plaintext).map_err(SecurityError::from)
    }
}

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

    #[test]
    fn test_generate_key_returns_result_with_correct_length_and_uniqueness() {
        // Regression test: `generate_key` used to be infallible (`Vec<u8>`)
        // and panicked via `.expect("getrandom failed")` on OS RNG failure.
        // It must now be fallible and surface a typed `SecurityError`
        // instead of panicking.
        let key1 = AtRestEncryptor::generate_key(EncryptionAlgorithm::Aes256Gcm)
            .expect("generate_key should succeed under normal conditions");
        let key2 = AtRestEncryptor::generate_key(EncryptionAlgorithm::Aes256Gcm)
            .expect("generate_key should succeed under normal conditions");

        assert_eq!(key1.len(), 32);
        assert_eq!(key2.len(), 32);
        // Keys should be independently random, not deterministic/reused.
        assert_ne!(key1, key2);

        let chacha_key = AtRestEncryptor::generate_key(EncryptionAlgorithm::ChaCha20Poly1305)
            .expect("generate_key should succeed under normal conditions");
        assert_eq!(chacha_key.len(), 32);
    }

    #[test]
    fn test_aes_gcm_encryption() {
        let key = AtRestEncryptor::generate_key(EncryptionAlgorithm::Aes256Gcm)
            .expect("Failed to generate key");
        let encryptor =
            AtRestEncryptor::new(EncryptionAlgorithm::Aes256Gcm, key, "test-key".to_string())
                .expect("Failed to create encryptor");

        let plaintext = b"Hello, World!";
        let encrypted = encryptor
            .encrypt(plaintext, None)
            .expect("Encryption failed");

        assert_ne!(encrypted.ciphertext, plaintext);
        assert_eq!(encrypted.metadata.algorithm, EncryptionAlgorithm::Aes256Gcm);

        let decrypted = encryptor.decrypt(&encrypted).expect("Decryption failed");
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_aes_gcm_with_aad() {
        let key = AtRestEncryptor::generate_key(EncryptionAlgorithm::Aes256Gcm)
            .expect("Failed to generate key");
        let encryptor =
            AtRestEncryptor::new(EncryptionAlgorithm::Aes256Gcm, key, "test-key".to_string())
                .expect("Failed to create encryptor");

        let plaintext = b"Hello, World!";
        let aad = b"additional data";
        let encrypted = encryptor
            .encrypt(plaintext, Some(aad))
            .expect("Encryption failed");

        let decrypted = encryptor.decrypt(&encrypted).expect("Decryption failed");
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_chacha_encryption() {
        let key = AtRestEncryptor::generate_key(EncryptionAlgorithm::ChaCha20Poly1305)
            .expect("Failed to generate key");
        let encryptor = AtRestEncryptor::new(
            EncryptionAlgorithm::ChaCha20Poly1305,
            key,
            "test-key".to_string(),
        )
        .expect("Failed to create encryptor");

        let plaintext = b"Hello, World!";
        let encrypted = encryptor
            .encrypt(plaintext, None)
            .expect("Encryption failed");

        assert_ne!(encrypted.ciphertext, plaintext);
        assert_eq!(
            encrypted.metadata.algorithm,
            EncryptionAlgorithm::ChaCha20Poly1305
        );

        let decrypted = encryptor.decrypt(&encrypted).expect("Decryption failed");
        assert_eq!(decrypted, plaintext);
    }

    #[test]
    fn test_field_encryptor_string() {
        let key = AtRestEncryptor::generate_key(EncryptionAlgorithm::Aes256Gcm)
            .expect("Failed to generate key");
        let encryptor =
            AtRestEncryptor::new(EncryptionAlgorithm::Aes256Gcm, key, "test-key".to_string())
                .expect("Failed to create encryptor");
        let field_encryptor = FieldEncryptor::new(encryptor);

        let original = "sensitive data";
        let encrypted = field_encryptor
            .encrypt_string(original)
            .expect("Encryption failed");

        assert_ne!(encrypted, original);

        let decrypted = field_encryptor
            .decrypt_string(&encrypted)
            .expect("Decryption failed");
        assert_eq!(decrypted, original);
    }

    #[test]
    fn test_encrypt_in_place() {
        let key = AtRestEncryptor::generate_key(EncryptionAlgorithm::Aes256Gcm)
            .expect("Failed to generate key");
        let encryptor =
            AtRestEncryptor::new(EncryptionAlgorithm::Aes256Gcm, key, "test-key".to_string())
                .expect("Failed to create encryptor");

        let mut buffer = b"Hello, World!".to_vec();
        let original = buffer.clone();

        let metadata = encryptor
            .encrypt_in_place(&mut buffer, None)
            .expect("Encryption failed");

        assert_ne!(buffer, original);

        let encrypted = EncryptedData::new(buffer, metadata);
        let decrypted = encryptor.decrypt(&encrypted).expect("Decryption failed");
        assert_eq!(decrypted, original);
    }
}