sal/src/vault/kvs/store.rs

//! Implementation of a simple key-value store using the filesystem.

use crate::vault::kvs::error::{KvsError, Result};
use crate::vault::symmetric::implementation::{
    decrypt_symmetric, derive_key_from_password, encrypt_symmetric,
};
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};

/// A key-value pair.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KvPair {
    pub key: String,
    pub value: String,
}

/// A common trait for key-value store implementations.
///
/// This trait defines the operations that all key-value stores must support,
/// regardless of the underlying implementation.
pub trait KVStore: Clone {
    /// Stores a value with the given key.
    fn set<K, V>(&self, key: K, value: &V) -> Result<()>
    where
        K: ToString,
        V: Serialize;

    /// Retrieves a value for the given key.
    fn get<K, V>(&self, key: K) -> Result<V>
    where
        K: ToString,
        V: DeserializeOwned;

    /// Deletes a value for the given key.
    fn delete<K>(&self, key: K) -> Result<()>
    where
        K: ToString;

    /// Checks if a key exists in the store.
    fn contains<K>(&self, key: K) -> Result<bool>
    where
        K: ToString;

    /// Lists all keys in the store.
    fn keys(&self) -> Result<Vec<String>>;

    /// Clears all key-value pairs from the store.
    fn clear(&self) -> Result<()>;

    /// Gets the name of the store.
    fn name(&self) -> &str;

    /// Gets whether the store is encrypted.
    fn is_encrypted(&self) -> bool;
}

/// A simple key-value store.
///
/// This implementation uses the filesystem to store key-value pairs.
#[derive(Clone)]
pub struct KvStore {
    /// The name of the store
    name: String,
    /// The path to the store file
    path: PathBuf,
    /// In-memory cache of the store data
    data: Arc<Mutex<HashMap<String, String>>>,
    /// Whether the store is encrypted
    encrypted: bool,
    /// The password for encryption (if encrypted)
    password: Option<String>,
}

/// Gets the path to the key-value store directory.
pub fn get_store_path() -> PathBuf {
    let home_dir = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
    home_dir.join(".hero-vault").join("kvs")
}

/// Creates a new key-value store with the given name.
///
/// # Arguments
///
/// * `name` - The name of the store
/// * `encrypted` - Whether to encrypt the store
/// * `password` - The password for encryption (required if encrypted is true)
///
/// # Returns
///
/// A new `KvStore` instance
pub fn create_store(name: &str, encrypted: bool, password: Option<&str>) -> Result<KvStore> {
    // Check if password is provided when encryption is enabled
    if encrypted && password.is_none() {
        return Err(KvsError::Other(
            "Password required for encrypted store".to_string(),
        ));
    }

    // Create the store directory if it doesn't exist
    let store_dir = get_store_path();
    if !store_dir.exists() {
        fs::create_dir_all(&store_dir)?;
    }

    // Create the store file path
    let store_path = store_dir.join(format!("{}.json", name));

    // Create an empty store
    let store = KvStore {
        name: name.to_string(),
        path: store_path,
        data: Arc::new(Mutex::new(HashMap::new())),
        encrypted,
        password: password.map(|s| s.to_string()),
    };

    // Save the empty store
    store.save()?;

    Ok(store)
}

/// Opens an existing key-value store with the given name.
///
/// # Arguments
///
/// * `name` - The name of the store
/// * `password` - The password for decryption (required if the store is encrypted)
///
/// # Returns
///
/// The opened `KvStore` instance
pub fn open_store(name: &str, password: Option<&str>) -> Result<KvStore> {
    // Get the store file path
    let store_dir = get_store_path();
    let store_path = store_dir.join(format!("{}.json", name));

    // Check if the store exists
    if !store_path.exists() {
        return Err(KvsError::StoreNotFound(name.to_string()));
    }

    // Read the store file
    let file_content = fs::read_to_string(&store_path)?;

    // Check if the file is encrypted (simple heuristic)
    let is_encrypted = !file_content.starts_with('{');

    // If encrypted, we need a password
    if is_encrypted && password.is_none() {
        return Err(KvsError::Other(
            "Password required for encrypted store".to_string(),
        ));
    }

    // Parse the store data
    let data: HashMap<String, String> = if is_encrypted {
        // Decrypt the file content
        let password = password.unwrap();
        let encrypted_data: Vec<u8> = serde_json::from_str(&file_content)?;
        let key = derive_key_from_password(password);
        let decrypted_data = decrypt_symmetric(&key, &encrypted_data)?;
        let decrypted_str = String::from_utf8(decrypted_data)
            .map_err(|e| KvsError::Deserialization(e.to_string()))?;
        serde_json::from_str(&decrypted_str)?
    } else {
        serde_json::from_str(&file_content)?
    };

    // Create the store
    let store = KvStore {
        name: name.to_string(),
        path: store_path,
        data: Arc::new(Mutex::new(data)),
        encrypted: is_encrypted,
        password: password.map(|s| s.to_string()),
    };

    Ok(store)
}

/// Deletes a key-value store with the given name.
///
/// # Arguments
///
/// * `name` - The name of the store to delete
///
/// # Returns
///
/// `Ok(())` if the operation was successful
pub fn delete_store(name: &str) -> Result<()> {
    // Get the store file path
    let store_dir = get_store_path();
    let store_path = store_dir.join(format!("{}.json", name));

    // Check if the store exists
    if !store_path.exists() {
        return Err(KvsError::StoreNotFound(name.to_string()));
    }

    // Delete the store file
    fs::remove_file(store_path)?;

    Ok(())
}

/// Lists all available key-value stores.
///
/// # Returns
///
/// A vector of store names
pub fn list_stores() -> Result<Vec<String>> {
    // Get the store directory
    let store_dir = get_store_path();
    if !store_dir.exists() {
        return Ok(Vec::new());
    }

    // List all JSON files in the directory
    let mut stores = Vec::new();
    for entry in fs::read_dir(store_dir)? {
        let entry = entry?;
        let path = entry.path();
        if path.is_file() && path.extension().map_or(false, |ext| ext == "json") {
            if let Some(name) = path.file_stem() {
                if let Some(name_str) = name.to_str() {
                    stores.push(name_str.to_string());
                }
            }
        }
    }

    Ok(stores)
}

impl KvStore {
    /// Saves the store to disk.
    fn save(&self) -> Result<()> {
        // Get the store data
        let data = self.data.lock().unwrap();

        // Serialize the data
        let serialized = serde_json::to_string(&*data)?;

        // Write to file
        if self.encrypted {
            if let Some(password) = &self.password {
                // Encrypt the data
                let key = derive_key_from_password(password);
                let encrypted_data = encrypt_symmetric(&key, serialized.as_bytes())?;
                let encrypted_json = serde_json::to_string(&encrypted_data)?;
                fs::write(&self.path, encrypted_json)?;
            } else {
                return Err(KvsError::Other(
                    "Password required for encrypted store".to_string(),
                ));
            }
        } else {
            fs::write(&self.path, serialized)?;
        }

        Ok(())
    }
}

// Implement the KVStore trait for the standard KvStore
impl KVStore for KvStore {
    /// Stores a value with the given key.
    fn set<K, V>(&self, key: K, value: &V) -> Result<()>
    where
        K: ToString,
        V: Serialize,
    {
        let key_str = key.to_string();
        let serialized = serde_json::to_string(value)?;

        // Update in-memory data
        {
            let mut data = self.data.lock().unwrap();
            data.insert(key_str, serialized);
        }

        // Save to disk
        self.save()?;

        Ok(())
    }

    /// Retrieves a value for the given key.
    fn get<K, V>(&self, key: K) -> Result<V>
    where
        K: ToString,
        V: DeserializeOwned,
    {
        let key_str = key.to_string();
        let data = self.data.lock().unwrap();

        match data.get(&key_str) {
            Some(serialized) => {
                let value: V = serde_json::from_str(serialized)?;
                Ok(value)
            }
            None => Err(KvsError::KeyNotFound(key_str)),
        }
    }

    /// Deletes a value for the given key.
    fn delete<K>(&self, key: K) -> Result<()>
    where
        K: ToString,
    {
        let key_str = key.to_string();

        // Update in-memory data
        {
            let mut data = self.data.lock().unwrap();
            if data.remove(&key_str).is_none() {
                return Err(KvsError::KeyNotFound(key_str));
            }
        }

        // Save to disk
        self.save()?;

        Ok(())
    }

    /// Checks if a key exists in the store.
    fn contains<K>(&self, key: K) -> Result<bool>
    where
        K: ToString,
    {
        let key_str = key.to_string();
        let data = self.data.lock().unwrap();

        Ok(data.contains_key(&key_str))
    }

    /// Lists all keys in the store.
    fn keys(&self) -> Result<Vec<String>> {
        let data = self.data.lock().unwrap();

        Ok(data.keys().cloned().collect())
    }

    /// Clears all key-value pairs from the store.
    fn clear(&self) -> Result<()> {
        // Update in-memory data
        {
            let mut data = self.data.lock().unwrap();
            data.clear();
        }

        // Save to disk
        self.save()?;

       Ok(())
    }

    /// Gets the name of the store.
    fn name(&self) -> &str {
        &self.name
    }

    /// Gets whether the store is encrypted.
    fn is_encrypted(&self) -> bool {
        self.encrypted
    }
}