portkey/

vault.rs

use anyhow::{Context, Result};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sodiumoxide::crypto::pwhash::argon2id13;
use sodiumoxide::crypto::secretbox;
use std::fs;
use std::io::Write;
use std::os::unix::fs::PermissionsExt;
use std::path::PathBuf;
use uuid::Uuid;

use crate::crypto::{generate_salt, MasterKey};
use crate::models::{Server, VaultData};

#[derive(Debug, Serialize, Deserialize)]
pub struct VaultFile {
    pub salt: argon2id13::Salt,
    pub nonce: secretbox::Nonce,
    pub ciphertext: Vec<u8>,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
}

pub struct Vault {
    data_path: PathBuf,
    master_key: Option<MasterKey>,
    data: Option<VaultData>,
}

impl Vault {
    pub fn new() -> Result<Self> {
        let data_dir = dirs::data_dir()
            .context("Failed to find data directory")?
            .join("portkey");
        let data_path = data_dir.join("vault.dat");
        Self::new_at(data_path)
    }

    pub fn new_at(data_path: PathBuf) -> Result<Self> {
        if let Some(parent) = data_path.parent() {
            fs::create_dir_all(parent)?;
        }

        Ok(Self {
            data_path,
            master_key: None,
            data: None,
        })
    }

    pub fn exists(&self) -> bool {
        self.data_path.exists()
    }

    pub fn unlock(&mut self, password: Option<&str>) -> Result<()> {
        if !self.exists() {
            return Err(anyhow::anyhow!("Vault does not exist"));
        }

        let vault_file = self.load_vault_file()?;

        // Try to decrypt with password if provided
        if let Some(password) = password {
            let master_key = MasterKey::from_password(password, &vault_file.salt)?;

            // Check if this looks like encrypted data by attempting decryption
            let decrypted_data = master_key.decrypt(&vault_file.ciphertext, &vault_file.nonce)?;
            let vault_data: VaultData = serde_json::from_slice(&decrypted_data)
                .context("Failed to deserialize vault data")?;

            self.master_key = Some(master_key);
            self.data = Some(vault_data);
        } else {
            // No password provided, assume unencrypted vault
            let vault_data: VaultData = serde_json::from_slice(&vault_file.ciphertext)
                .context("Failed to deserialize vault data - try providing a password")?;

            self.master_key = None;
            self.data = Some(vault_data);
        }

        Ok(())
    }

    pub fn create(&mut self, password: Option<&str>) -> Result<()> {
        if self.exists() {
            return Err(anyhow::anyhow!("Vault already exists"));
        }

        let vault_data = VaultData::new();
        let serialized = serde_json::to_vec(&vault_data)?;

        let vault_file = if let Some(password) = password {
            // Password-protected vault
            let salt = generate_salt();
            let master_key = MasterKey::from_password(password, &salt)?;
            let (nonce, ciphertext) = master_key.encrypt(&serialized);

            VaultFile {
                salt,
                nonce,
                ciphertext,
                created_at: Utc::now(),
                updated_at: Utc::now(),
            }
        } else {
            // Unencrypted vault (no password)
            let salt = generate_salt(); // Still use salt for consistency
            let nonce = secretbox::gen_nonce();

            VaultFile {
                salt,
                nonce,
                ciphertext: serialized, // Store data unencrypted
                created_at: Utc::now(),
                updated_at: Utc::now(),
            }
        };

        self.save_vault_file(&vault_file)?;

        if let Some(password) = password {
            let master_key = MasterKey::from_password(password, &vault_file.salt)?;
            self.master_key = Some(master_key);
        }
        self.data = Some(vault_data);

        Ok(())
    }

    pub fn is_unlocked(&self) -> bool {
        self.data.is_some()
    }

    pub fn add_server(&mut self, server: Server) -> Result<()> {
        self.ensure_unlocked()?;

        let data = self.data.as_mut().unwrap();
        data.add_server(server);

        self.save()?;
        Ok(())
    }

    pub fn remove_server(&mut self, id: &uuid::Uuid) -> Result<bool> {
        self.ensure_unlocked()?;

        let data = self.data.as_mut().unwrap();
        let removed = data.remove_server(id);

        if removed {
            self.save()?;
        }

        Ok(removed)
    }

    pub fn list_servers(&self) -> Result<&Vec<Server>> {
        self.ensure_unlocked()?;

        Ok(&self.data.as_ref().unwrap().servers)
    }

    pub fn find_server(&self, id: &uuid::Uuid) -> Result<Option<&Server>> {
        self.ensure_unlocked()?;

        Ok(self.data.as_ref().unwrap().find_server(id))
    }

    pub fn replace_server(&mut self, server: Server) -> Result<bool> {
        self.ensure_unlocked()?;
        let data = self.data.as_mut().unwrap();
        let replaced = data.replace_server(server);
        if replaced {
            self.save()?;
        }
        Ok(replaced)
    }

    pub fn vault_path(&self) -> &PathBuf {
        &self.data_path
    }

    fn ensure_unlocked(&self) -> Result<()> {
        if !self.is_unlocked() {
            return Err(anyhow::anyhow!("Vault is locked"));
        }
        Ok(())
    }

    fn load_vault_file(&self) -> Result<VaultFile> {
        let content = fs::read(&self.data_path)?;
        let vault_file: VaultFile = serde_json::from_slice(&content)?;
        Ok(vault_file)
    }

    fn save_vault_file(&self, vault_file: &VaultFile) -> Result<()> {
        let content = serde_json::to_vec(vault_file)?;

        let parent = self
            .data_path
            .parent()
            .ok_or_else(|| anyhow::anyhow!("Vault path has no parent directory"))?;
        fs::create_dir_all(parent)?;

        let file_name = self
            .data_path
            .file_name()
            .ok_or_else(|| anyhow::anyhow!("Vault path has no file name"))?
            .to_string_lossy();
        let temp_path = parent.join(format!(".{file_name}.{}.tmp", Uuid::new_v4()));

        let mut file = fs::OpenOptions::new()
            .create_new(true)
            .write(true)
            .open(&temp_path)?;

        let mut perms = file.metadata()?.permissions();
        perms.set_mode(0o600); // Read/write for owner only
        file.set_permissions(perms)?;

        file.write_all(&content)?;
        file.sync_all()?;
        drop(file);

        if let Err(error) = fs::rename(&temp_path, &self.data_path) {
            let _ = fs::remove_file(&temp_path);
            return Err(error.into());
        }

        Ok(())
    }

    fn save(&mut self) -> Result<()> {
        let data = self.data.as_ref().unwrap();
        let serialized = serde_json::to_vec(data)?;

        let vault_file = if let Some(master_key) = &self.master_key {
            // Encrypted vault: reuse existing salt to keep key derivation stable
            let existing = self.load_vault_file().ok();
            let salt = existing
                .as_ref()
                .map(|f| f.salt)
                .unwrap_or_else(generate_salt);

            let (nonce, ciphertext) = master_key.encrypt(&serialized);
            VaultFile {
                salt,
                nonce,
                ciphertext,
                created_at: existing.map(|f| f.created_at).unwrap_or_else(Utc::now),
                updated_at: Utc::now(),
            }
        } else {
            // Unencrypted vault
            let salt = generate_salt();
            let nonce = secretbox::gen_nonce();

            VaultFile {
                salt,
                nonce,
                ciphertext: serialized, // Store unencrypted
                created_at: self
                    .load_vault_file()
                    .map(|f| f.created_at)
                    .unwrap_or_else(|_| Utc::now()),
                updated_at: Utc::now(),
            }
        };

        self.save_vault_file(&vault_file)?;
        Ok(())
    }
}