trackWork 0.15.0

A terminal-based time tracking application for managing work sessions
use anyhow::{anyhow, Result};
use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;

use aes_gcm::aead::{Aead, KeyInit};
use aes_gcm::{Aes256Gcm, Nonce};
use argon2::Argon2;
use serde::{Deserialize, Serialize};
use zeroize::Zeroize;

const MAGIC: &[u8; 4] = b"TWE\x01";
const SALT_LEN: usize = 16;
const NONCE_LEN: usize = 12;

#[derive(Debug, Serialize, Deserialize)]
struct SecretsFile {
    #[serde(default)]
    secrets: HashMap<String, String>,
}

pub struct SecretsManager {
    encryption_key: Option<[u8; 32]>,
    salt: Option<[u8; 16]>,
}

impl SecretsManager {
    pub fn new() -> Self {
        SecretsManager {
            encryption_key: None,
            salt: None,
        }
    }

    /// Check if the secrets file exists and has the encrypted magic header.
    pub fn has_encrypted_file(&self) -> bool {
        let path = match Self::secrets_file_path() {
            Ok(p) => p,
            Err(_) => return false,
        };
        if !path.exists() {
            return false;
        }
        match fs::read(&path) {
            Ok(data) => data.len() >= 4 && &data[..4] == MAGIC,
            Err(_) => false,
        }
    }

    /// Whether a decryption key is currently loaded.
    pub fn is_unlocked(&self) -> bool {
        self.encryption_key.is_some()
    }

    /// Derive key from passphrase + salt, decrypt the file to verify, store key+salt.
    pub fn unlock(&mut self, passphrase: &str) -> Result<()> {
        let path = Self::secrets_file_path()?;
        let data = fs::read(&path)?;

        if data.len() < 4 + SALT_LEN + NONCE_LEN {
            return Err(anyhow!("Encrypted file too short"));
        }
        if &data[..4] != MAGIC {
            return Err(anyhow!("Not an encrypted secrets file"));
        }

        let salt: [u8; 16] = data[4..20].try_into().unwrap();
        let nonce_bytes: [u8; 12] = data[20..32].try_into().unwrap();
        let ciphertext = &data[32..];

        let key = Self::derive_key(passphrase, &salt)?;
        let cipher = Aes256Gcm::new_from_slice(&key)
            .map_err(|e| anyhow!("Failed to create cipher: {}", e))?;
        let nonce = Nonce::from_slice(&nonce_bytes);

        // Decrypt to verify the passphrase is correct
        cipher
            .decrypt(nonce, ciphertext)
            .map_err(|_| anyhow!("Wrong passphrase"))?;

        self.encryption_key = Some(key);
        self.salt = Some(salt);
        Ok(())
    }

    /// Set a new passphrase: generate new salt, derive key, re-encrypt all secrets.
    pub fn set_passphrase(&mut self, passphrase: &str) -> Result<()> {
        // Read current secrets (plaintext or decrypted)
        let secrets_data = self.read_secrets_file()?;

        // Generate new salt
        let mut salt = [0u8; SALT_LEN];
        use rand::RngCore;
        rand::thread_rng().fill_bytes(&mut salt);

        let key = Self::derive_key(passphrase, &salt)?;
        self.encryption_key = Some(key);
        self.salt = Some(salt);

        // Re-write (will now encrypt)
        self.write_secrets_file(&secrets_data)?;
        Ok(())
    }

    /// Verify old passphrase, then set new passphrase.
    pub fn change_passphrase(&mut self, old: &str, new: &str) -> Result<()> {
        // Verify old passphrase by attempting unlock
        let path = Self::secrets_file_path()?;
        let data = fs::read(&path)?;

        if data.len() < 4 + SALT_LEN + NONCE_LEN || &data[..4] != MAGIC {
            return Err(anyhow!("No encrypted file to change passphrase for"));
        }

        let salt: [u8; 16] = data[4..20].try_into().unwrap();
        let nonce_bytes: [u8; 12] = data[20..32].try_into().unwrap();
        let ciphertext = &data[32..];

        let old_key = Self::derive_key(old, &salt)?;
        let cipher = Aes256Gcm::new_from_slice(&old_key)
            .map_err(|e| anyhow!("Failed to create cipher: {}", e))?;
        let nonce = Nonce::from_slice(&nonce_bytes);

        cipher
            .decrypt(nonce, ciphertext)
            .map_err(|_| anyhow!("Wrong old passphrase"))?;

        // Old passphrase verified; now set new one
        self.encryption_key = Some(old_key);
        self.salt = Some(salt);
        self.set_passphrase(new)
    }

    pub fn get(&self, key: &str) -> Result<Option<String>> {
        let data = self.read_secrets_file()?;
        Ok(data.secrets.get(key).cloned())
    }

    pub fn set(&self, key: &str, value: &str) -> Result<()> {
        let mut data = self.read_secrets_file()?;
        data.secrets.insert(key.to_string(), value.to_string());
        self.write_secrets_file(&data)
    }

    pub fn delete(&self, key: &str) -> Result<()> {
        let mut data = self.read_secrets_file()?;
        data.secrets.remove(key);
        self.write_secrets_file(&data)
    }

    /// Delete the secrets file entirely and reset encryption state.
    pub fn delete_secrets_file(&mut self) -> Result<()> {
        let path = Self::secrets_file_path()?;
        if path.exists() {
            fs::remove_file(&path)?;
        }
        self.encryption_key = None;
        self.salt = None;
        Ok(())
    }

    fn secrets_file_path() -> Result<PathBuf> {
        let home = dirs::home_dir().ok_or_else(|| anyhow!("No home directory"))?;
        Ok(home.join(".timetrack.secrets"))
    }

    fn derive_key(passphrase: &str, salt: &[u8; 16]) -> Result<[u8; 32]> {
        let mut key = [0u8; 32];
        Argon2::default()
            .hash_password_into(passphrase.as_bytes(), salt, &mut key)
            .map_err(|e| anyhow!("Key derivation failed: {}", e))?;
        Ok(key)
    }

    fn read_secrets_file(&self) -> Result<SecretsFile> {
        let path = Self::secrets_file_path()?;
        if !path.exists() {
            return Ok(SecretsFile {
                secrets: HashMap::new(),
            });
        }

        let raw = fs::read(&path)?;

        // Check for encrypted format
        if raw.len() >= 4 && &raw[..4] == MAGIC {
            let key = self
                .encryption_key
                .ok_or_else(|| anyhow!("Secrets file is encrypted but no key loaded"))?;

            if raw.len() < 4 + SALT_LEN + NONCE_LEN {
                return Err(anyhow!("Encrypted file too short"));
            }

            let nonce_bytes: [u8; 12] = raw[20..32].try_into().unwrap();
            let ciphertext = &raw[32..];

            let cipher = Aes256Gcm::new_from_slice(&key)
                .map_err(|e| anyhow!("Failed to create cipher: {}", e))?;
            let nonce = Nonce::from_slice(&nonce_bytes);

            let plaintext = cipher
                .decrypt(nonce, ciphertext)
                .map_err(|_| anyhow!("Decryption failed"))?;

            let contents = String::from_utf8(plaintext)?;
            let data: SecretsFile = toml::from_str(&contents)?;
            Ok(data)
        } else {
            // Legacy plaintext TOML
            let contents = String::from_utf8(raw)?;
            let data: SecretsFile = toml::from_str(&contents)?;
            Ok(data)
        }
    }

    fn write_secrets_file(&self, data: &SecretsFile) -> Result<()> {
        let path = Self::secrets_file_path()?;
        let contents = toml::to_string_pretty(data)?;

        if let (Some(key), Some(salt)) = (self.encryption_key, self.salt) {
            // Encrypt
            let mut nonce_bytes = [0u8; NONCE_LEN];
            use rand::RngCore;
            rand::thread_rng().fill_bytes(&mut nonce_bytes);

            let cipher = Aes256Gcm::new_from_slice(&key)
                .map_err(|e| anyhow!("Failed to create cipher: {}", e))?;
            let nonce = Nonce::from_slice(&nonce_bytes);

            let ciphertext = cipher
                .encrypt(nonce, contents.as_bytes())
                .map_err(|e| anyhow!("Encryption failed: {}", e))?;

            let mut output = Vec::with_capacity(4 + SALT_LEN + NONCE_LEN + ciphertext.len());
            output.extend_from_slice(MAGIC);
            output.extend_from_slice(&salt);
            output.extend_from_slice(&nonce_bytes);
            output.extend_from_slice(&ciphertext);

            fs::write(&path, &output)?;
        } else {
            // Plaintext
            fs::write(&path, &contents)?;
        }

        // chmod 600 on Unix
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            let perms = fs::Permissions::from_mode(0o600);
            fs::set_permissions(&path, perms)?;
        }

        Ok(())
    }
}

impl Drop for SecretsManager {
    fn drop(&mut self) {
        if let Some(ref mut key) = self.encryption_key {
            key.zeroize();
        }
    }
}