activity-tracker/src/storage/encryption.rs

/// Encryption utilities using AES-256-GCM with PBKDF2 key derivation
use aes_gcm::{
    aead::{Aead, KeyInit, OsRng},
    Aes256Gcm, Nonce,
};
use pbkdf2::{password_hash::SaltString, pbkdf2_hmac};
use rand::RngCore;
use sha2::Sha512;
use crate::error::{AppError, Result};

const NONCE_SIZE: usize = 12; // GCM recommended nonce size
const SALT_SIZE: usize = 16;
const KEY_SIZE: usize = 32; // 256 bits
const PBKDF2_ITERATIONS: u32 = 100_000;

pub struct Encryptor {
    cipher: Aes256Gcm,
}

impl Encryptor {
    /// Create new encryptor from user password
    pub fn from_password(password: &str) -> Result<Self> {
        // Generate random salt
        let salt = SaltString::generate(&mut OsRng);
        let key = Self::derive_key(password, salt.as_bytes())?;

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

        Ok(Self { cipher })
    }

    /// Create encryptor from password and salt (for decryption)
    pub fn from_password_and_salt(password: &str, salt: &[u8]) -> Result<Self> {
        let key = Self::derive_key(password, salt)?;

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

        Ok(Self { cipher })
    }

    /// Derive encryption key from password using PBKDF2-HMAC-SHA512
    fn derive_key(password: &str, salt: &[u8]) -> Result<Vec<u8>> {
        let mut key = vec![0u8; KEY_SIZE];
        pbkdf2_hmac::<Sha512>(
            password.as_bytes(),
            salt,
            PBKDF2_ITERATIONS,
            &mut key,
        );
        Ok(key)
    }

    /// Encrypt data with AES-256-GCM
    /// Format: [salt (16B)][nonce (12B)][ciphertext]
    pub fn encrypt(&self, plaintext: &[u8]) -> Result<Vec<u8>> {
        // Generate random nonce
        let mut nonce_bytes = [0u8; NONCE_SIZE];
        OsRng.fill_bytes(&mut nonce_bytes);
        let nonce = Nonce::from_slice(&nonce_bytes);

        // Encrypt
        let ciphertext = self.cipher
            .encrypt(nonce, plaintext)
            .map_err(|e| AppError::Encryption(format!("Encryption failed: {}", e)))?;

        // Generate salt for storage
        let mut salt = vec![0u8; SALT_SIZE];
        OsRng.fill_bytes(&mut salt);

        // Combine salt + nonce + ciphertext
        let mut result = Vec::with_capacity(SALT_SIZE + NONCE_SIZE + ciphertext.len());
        result.extend_from_slice(&salt);
        result.extend_from_slice(&nonce_bytes);
        result.extend_from_slice(&ciphertext);

        Ok(result)
    }

    /// Decrypt data
    /// Expected format: [salt (16B)][nonce (12B)][ciphertext]
    pub fn decrypt(&self, encrypted: &[u8]) -> Result<Vec<u8>> {
        if encrypted.len() < SALT_SIZE + NONCE_SIZE {
            return Err(AppError::Encryption("Invalid encrypted data size".to_string()));
        }

        // Extract nonce and ciphertext (skip salt for now)
        let nonce_start = SALT_SIZE;
        let nonce = Nonce::from_slice(&encrypted[nonce_start..nonce_start + NONCE_SIZE]);
        let ciphertext = &encrypted[nonce_start + NONCE_SIZE..];

        // Decrypt
        let plaintext = self.cipher
            .decrypt(nonce, ciphertext)
            .map_err(|e| AppError::Encryption(format!("Decryption failed: {}", e)))?;

        Ok(plaintext)
    }

    /// Encrypt if data is provided
    pub fn encrypt_optional(&self, data: Option<&[u8]>) -> Result<Option<Vec<u8>>> {
        match data {
            Some(d) => Ok(Some(self.encrypt(d)?)),
            None => Ok(None),
        }
    }

    /// Decrypt if data is provided
    pub fn decrypt_optional(&self, data: Option<&[u8]>) -> Result<Option<Vec<u8>>> {
        match data {
            Some(d) => Ok(Some(self.decrypt(d)?)),
            None => Ok(None),
        }
    }
}

/// Generate secure salt
pub fn generate_salt() -> Vec<u8> {
    let mut salt = vec![0u8; SALT_SIZE];
    OsRng.fill_bytes(&mut salt);
    salt
}

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

    #[test]
    fn test_encryption_decryption() {
        let password = "test_password_123";
        let encryptor = Encryptor::from_password(password).unwrap();

        let plaintext = b"Hello, World! This is a test message.";
        let encrypted = encryptor.encrypt(plaintext).unwrap();
        let decrypted = encryptor.decrypt(&encrypted).unwrap();

        assert_eq!(plaintext.to_vec(), decrypted);
        assert_ne!(plaintext.to_vec(), encrypted); // Should be different
    }

    #[test]
    fn test_encryption_with_empty_data() {
        let encryptor = Encryptor::from_password("password").unwrap();
        let encrypted = encryptor.encrypt(b"").unwrap();
        let decrypted = encryptor.decrypt(&encrypted).unwrap();
        assert_eq!(decrypted.len(), 0);
    }
}