owlen/crates/owlen-core/src/encryption.rs

use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;

use aes_gcm::{
    aead::{Aead, KeyInit},
    Aes256Gcm, Nonce,
};
use anyhow::{bail, Context, Result};
use ring::digest;
use ring::rand::{SecureRandom, SystemRandom};
use serde::{Deserialize, Serialize};
use serde_json::Value as JsonValue;

pub struct EncryptedStorage {
    cipher: Aes256Gcm,
    storage_path: PathBuf,
}

#[derive(Serialize, Deserialize)]
struct EncryptedData {
    nonce: [u8; 12],
    ciphertext: Vec<u8>,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct VaultData {
    pub master_key: Vec<u8>,
    #[serde(default)]
    pub settings: HashMap<String, JsonValue>,
}

pub struct VaultHandle {
    storage: EncryptedStorage,
    pub data: VaultData,
}

impl VaultHandle {
    pub fn master_key(&self) -> &[u8] {
        &self.data.master_key
    }

    pub fn settings(&self) -> &HashMap<String, JsonValue> {
        &self.data.settings
    }

    pub fn settings_mut(&mut self) -> &mut HashMap<String, JsonValue> {
        &mut self.data.settings
    }

    pub fn persist(&self) -> Result<()> {
        self.storage.store(&self.data)
    }
}

impl EncryptedStorage {
    pub fn new(storage_path: PathBuf, password: &str) -> Result<Self> {
        let digest = digest::digest(&digest::SHA256, password.as_bytes());
        let cipher = Aes256Gcm::new_from_slice(digest.as_ref())
            .map_err(|_| anyhow::anyhow!("Invalid key length for AES-256"))?;

        if let Some(parent) = storage_path.parent() {
            fs::create_dir_all(parent).context("Failed to ensure storage directory exists")?;
        }

        Ok(Self {
            cipher,
            storage_path,
        })
    }

    pub fn store<T: Serialize>(&self, data: &T) -> Result<()> {
        let json = serde_json::to_vec(data).context("Failed to serialize data")?;

        let nonce = generate_nonce()?;
        let nonce_ref = Nonce::from_slice(&nonce);

        let ciphertext = self
            .cipher
            .encrypt(nonce_ref, json.as_ref())
            .map_err(|e| anyhow::anyhow!("Encryption failed: {}", e))?;

        let encrypted_data = EncryptedData { nonce, ciphertext };
        let encrypted_json = serde_json::to_vec(&encrypted_data)?;

        fs::write(&self.storage_path, encrypted_json).context("Failed to write encrypted data")?;

        Ok(())
    }

    pub fn load<T: for<'de> Deserialize<'de>>(&self) -> Result<T> {
        let encrypted_json =
            fs::read(&self.storage_path).context("Failed to read encrypted data")?;

        let encrypted_data: EncryptedData =
            serde_json::from_slice(&encrypted_json).context("Failed to parse encrypted data")?;

        let nonce_ref = Nonce::from_slice(&encrypted_data.nonce);
        let plaintext = self
            .cipher
            .decrypt(nonce_ref, encrypted_data.ciphertext.as_ref())
            .map_err(|e| anyhow::anyhow!("Decryption failed: {}", e))?;

        let data: T =
            serde_json::from_slice(&plaintext).context("Failed to deserialize decrypted data")?;

        Ok(data)
    }

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

    pub fn delete(&self) -> Result<()> {
        if self.exists() {
            fs::remove_file(&self.storage_path).context("Failed to delete encrypted storage")?;
        }
        Ok(())
    }

    pub fn verify_password(&self) -> Result<()> {
        if !self.exists() {
            return Ok(());
        }

        let encrypted_json =
            fs::read(&self.storage_path).context("Failed to read encrypted data")?;

        if encrypted_json.is_empty() {
            return Ok(());
        }

        let encrypted_data: EncryptedData =
            serde_json::from_slice(&encrypted_json).context("Failed to parse encrypted data")?;

        let nonce_ref = Nonce::from_slice(&encrypted_data.nonce);
        self.cipher
            .decrypt(nonce_ref, encrypted_data.ciphertext.as_ref())
            .map(|_| ())
            .map_err(|e| anyhow::anyhow!("Decryption failed: {}", e))
    }
}

pub fn prompt_password(prompt: &str) -> Result<String> {
    let password = rpassword::prompt_password(prompt)
        .map_err(|e| anyhow::anyhow!("Failed to read password: {e}"))?;
    if password.is_empty() {
        bail!("Password cannot be empty");
    }
    Ok(password)
}

pub fn prompt_new_password() -> Result<String> {
    loop {
        let first = prompt_password("Enter new master password: ")?;
        let confirm = prompt_password("Confirm master password: ")?;
        if first == confirm {
            return Ok(first);
        }
        println!("Passwords did not match. Please try again.");
    }
}

pub fn unlock_with_password(storage_path: PathBuf, password: &str) -> Result<VaultHandle> {
    let storage = EncryptedStorage::new(storage_path, password)?;
    let data = load_or_initialize_vault(&storage)?;
    Ok(VaultHandle { storage, data })
}

pub fn unlock_interactive(storage_path: PathBuf) -> Result<VaultHandle> {
    if storage_path.exists() {
        for attempt in 0..3 {
            let password = prompt_password("Enter master password: ")?;
            match unlock_with_password(storage_path.clone(), &password) {
                Ok(handle) => return Ok(handle),
                Err(err) => {
                    println!("Failed to unlock vault: {err}");
                    if attempt == 2 {
                        return Err(err);
                    }
                }
            }
        }
        bail!("Failed to unlock encrypted storage after multiple attempts");
    } else {
        println!(
            "No encrypted storage found at {}. Initializing a new vault.",
            storage_path.display()
        );
        let password = prompt_new_password()?;
        let storage = EncryptedStorage::new(storage_path, &password)?;
        let data = VaultData {
            master_key: generate_master_key()?,
            ..Default::default()
        };
        storage.store(&data)?;
        Ok(VaultHandle { storage, data })
    }
}

fn load_or_initialize_vault(storage: &EncryptedStorage) -> Result<VaultData> {
    match storage.load::<VaultData>() {
        Ok(data) => {
            if data.master_key.len() != 32 {
                bail!(
                    "Corrupted vault: master key has invalid length ({}). \
                    Expected 32 bytes for AES-256. Vault cannot be recovered.",
                    data.master_key.len()
                );
            }
            Ok(data)
        }
        Err(err) => {
            if storage.exists() {
                return Err(err);
            }
            let data = VaultData {
                master_key: generate_master_key()?,
                ..Default::default()
            };
            storage.store(&data)?;
            Ok(data)
        }
    }
}

fn generate_master_key() -> Result<Vec<u8>> {
    let mut key = vec![0u8; 32];
    SystemRandom::new()
        .fill(&mut key)
        .map_err(|_| anyhow::anyhow!("Failed to generate master key"))?;
    Ok(key)
}

fn generate_nonce() -> Result<[u8; 12]> {
    let mut nonce = [0u8; 12];
    let rng = SystemRandom::new();
    rng.fill(&mut nonce)
        .map_err(|_| anyhow::anyhow!("Failed to generate nonce"))?;
    Ok(nonce)
}