diff --git a/lumni/src/apps/builtin/llm/prompt/src/chat/db/encryption/mod.rs b/lumni/src/apps/builtin/llm/prompt/src/chat/db/encryption/mod.rs index e696b2c..d3d0ed2 100644 --- a/lumni/src/apps/builtin/llm/prompt/src/chat/db/encryption/mod.rs +++ b/lumni/src/apps/builtin/llm/prompt/src/chat/db/encryption/mod.rs @@ -2,7 +2,7 @@ use std::fs; use std::path::PathBuf; use base64::engine::general_purpose; -use base64::{Engine, Engine as _}; +use base64::Engine; use lumni::api::error::{ApplicationError, EncryptionError}; use ring::aead; use ring::rand::{SecureRandom, SystemRandom}; diff --git a/lumni/src/apps/builtin/llm/prompt/src/chat/db/user_profiles/mod.rs b/lumni/src/apps/builtin/llm/prompt/src/chat/db/user_profiles/mod.rs index 72284f9..bab7504 100644 --- a/lumni/src/apps/builtin/llm/prompt/src/chat/db/user_profiles/mod.rs +++ b/lumni/src/apps/builtin/llm/prompt/src/chat/db/user_profiles/mod.rs @@ -90,40 +90,55 @@ impl UserProfileDbHandler { new_settings: &JsonValue, ) -> Result<(), ApplicationError> { let mut db = self.db.lock().await; - db.process_queue_with_result(|tx| { - let current_data: Option<(String, Option)> = tx + // First, check if the profile exists and get its current data and is_default status + let current_data: Option<(String, Option, bool)> = tx .query_row( - "SELECT options, ssh_key_hash FROM user_profiles WHERE \ - name = ?", + "SELECT options, ssh_key_hash, is_default FROM \ + user_profiles WHERE name = ?", params![profile_name], - |row| Ok((row.get(0)?, row.get(1)?)), + |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)), ) .optional() .map_err(|e| ApplicationError::DatabaseError(e.to_string()))?; - let (merged_settings, ssh_key_hash) = - if let Some((current_json, existing_hash)) = current_data { - let current: JsonValue = + let (merged_settings, ssh_key_hash, is_default) = + if let Some((current_json, existing_hash, is_default)) = + current_data + { + let mut current: JsonValue = serde_json::from_str(¤t_json).map_err(|e| { ApplicationError::InvalidInput(format!( "Invalid JSON: {}", e )) })?; - let merged = self.merge_settings(¤t, new_settings)?; - (merged, existing_hash) + + // Merge settings, handling deletions + if let Some(current_obj) = current.as_object_mut() { + if let Some(new_obj) = new_settings.as_object() { + for (key, value) in new_obj { + if value.is_null() { + current_obj.remove(key); // Remove the key if the new value is null + } else { + current_obj + .insert(key.clone(), value.clone()); // Otherwise, update or add the key-value pair + } + } + } + } + + (current, existing_hash, is_default) } else { let ssh_key_hash = self .encryption_handler .as_ref() .and_then(|_| self.calculate_ssh_key_hash().ok()); - (new_settings.clone(), ssh_key_hash) + (new_settings.clone(), ssh_key_hash, false) // New profiles are not default by default }; let processed_settings = self.process_settings(&merged_settings, true, false)?; - let json_string = serde_json::to_string(&processed_settings) .map_err(|e| { ApplicationError::InvalidInput(format!( @@ -132,10 +147,11 @@ impl UserProfileDbHandler { )) })?; + // Use INSERT OR REPLACE, but explicitly set the is_default status tx.execute( "INSERT OR REPLACE INTO user_profiles (name, options, \ - ssh_key_hash) VALUES (?, ?, ?)", - params![profile_name, json_string, ssh_key_hash], + ssh_key_hash, is_default) VALUES (?, ?, ?, ?)", + params![profile_name, json_string, ssh_key_hash, is_default], ) .map_err(|e| ApplicationError::DatabaseError(e.to_string()))?; @@ -144,29 +160,6 @@ impl UserProfileDbHandler { .map_err(ApplicationError::from) } - fn merge_settings( - &self, - current: &JsonValue, - new: &JsonValue, - ) -> Result { - match (current, new) { - (JsonValue::Object(current_obj), JsonValue::Object(new_obj)) => { - let mut merged = current_obj.clone(); - for (key, value) in new_obj { - merged.insert( - key.clone(), - self.merge_settings( - current_obj.get(key).unwrap_or(&JsonValue::Null), - value, - )?, - ); - } - Ok(JsonValue::Object(merged)) - } - (_, new) => Ok(new.clone()), - } - } - fn process_settings( &self, value: &JsonValue, diff --git a/lumni/src/apps/builtin/llm/prompt/src/cli/subcommands/profile.rs b/lumni/src/apps/builtin/llm/prompt/src/cli/subcommands/profile.rs index 8b613ba..21bc3f9 100644 --- a/lumni/src/apps/builtin/llm/prompt/src/cli/subcommands/profile.rs +++ b/lumni/src/apps/builtin/llm/prompt/src/cli/subcommands/profile.rs @@ -1,4 +1,5 @@ use clap::{Arg, ArgAction, ArgGroup, ArgMatches, Command}; +use libc::kCCCallSequenceError; use lumni::api::error::ApplicationError; use serde_json::{Map, Value as JsonValue}; @@ -8,101 +9,138 @@ use crate::external as lumni; pub fn create_profile_subcommand() -> Command { Command::new("profile") .about("Manage user profiles") - .arg(Arg::new("name").help("Name of the profile").index(1)) - .arg( - Arg::new("set") - .long("set") - .short('s') - .help("Set profile values") - .num_args(2) - .value_names(["KEY", "VALUE"]) - .action(ArgAction::Append), - ) - .arg( - Arg::new("secure") - .long("secure") - .help("Mark the previous value as secure (to be encrypted)") - .action(ArgAction::SetTrue) - .requires("set"), - ) - .arg( - Arg::new("get") - .long("get") - .short('g') - .help("Get a specific profile value") - .num_args(1) - .value_name("KEY"), - ) + .subcommand(create_list_subcommand()) + .subcommand(create_show_subcommand()) + .subcommand(create_set_subcommand()) + .subcommand(create_get_subcommand()) + .subcommand(create_del_subcommand()) + .subcommand(create_rm_subcommand()) + .subcommand(create_set_default_subcommand()) + .subcommand(create_show_default_subcommand()) +} + +fn create_list_subcommand() -> Command { + Command::new("list").about("List all profiles") +} + +fn create_show_subcommand() -> Command { + Command::new("show") + .about("Show profile settings") + .arg(Arg::new("name").help("Name of the profile")) .arg( Arg::new("show-decrypted") .long("show-decrypted") .help("Show decrypted values instead of masked values") .action(ArgAction::SetTrue), ) +} + +fn create_set_subcommand() -> Command { + Command::new("set") + .about("Set a profile value") + .arg(Arg::new("name").help("Name of the profile")) + .arg(Arg::new("key").help("Key to set")) + .arg(Arg::new("value").help("Value to set")) .arg( - Arg::new("show") - .long("show") - .help("Show all profile values") - .action(ArgAction::SetTrue), - ) - .arg( - Arg::new("delete") - .long("delete") - .short('d') - .help("Delete the profile") + Arg::new("secure") + .long("secure") + .help("Mark the value as secure (to be encrypted)") .action(ArgAction::SetTrue), ) +} + +fn create_get_subcommand() -> Command { + Command::new("get") + .about("Get a specific profile value") + .arg(Arg::new("name").help("Name of the profile")) + .arg(Arg::new("key").help("Key to get")) .arg( - Arg::new("default") - .long("default") - .help("Set as the default profile") + Arg::new("show-decrypted") + .long("show-decrypted") + .help("Show decrypted value instead of masked value") .action(ArgAction::SetTrue), ) +} + +fn create_del_subcommand() -> Command { + Command::new("del") + .about("Delete a key from a profile") + .arg(Arg::new("name").help("Name of the profile")) + .arg(Arg::new("key").help("Key to delete")) +} + +fn create_rm_subcommand() -> Command { + // Renamed from create_delete_subcommand + Command::new("rm") + .about("Remove a profile") + .arg(Arg::new("name").help("Name of the profile to remove")) +} + +fn create_set_default_subcommand() -> Command { + Command::new("set-default") + .about("Set a profile as default") + .arg(Arg::new("name").help("Name of the profile")) +} + +fn create_show_default_subcommand() -> Command { + Command::new("show-default") + .about("Show the default profile") .arg( - Arg::new("list") - .long("list") - .short('l') - .help("List all profiles") + Arg::new("show-decrypted") + .long("show-decrypted") + .help("Show decrypted values instead of masked values") .action(ArgAction::SetTrue), ) - .group( - ArgGroup::new("profile_group") - .args(["set", "get", "show", "delete", "default", "list"]) - .required(false) - .multiple(false), - ) } pub async fn handle_profile_subcommand( profile_matches: &ArgMatches, db_handler: &mut UserProfileDbHandler, ) -> Result<(), ApplicationError> { - if profile_matches.get_flag("list") { - let profiles = db_handler.list_profiles().await?; - println!("Available profiles:"); - for profile in profiles { - println!(" {}", profile); + match profile_matches.subcommand() { + Some(("list", _)) => { + let profiles = db_handler.list_profiles().await?; + let default_profile = db_handler.get_default_profile().await?; + println!("Available profiles:"); + for profile in profiles { + if Some(&profile) == default_profile.as_ref() { + println!(" {} (default)", profile); + } else { + println!(" {}", profile); + } + } } - return Ok(()); - } - let profile_name = profile_matches.get_one::("name"); - - if profile_matches.contains_id("set") { - if let Some(profile_name) = profile_name { - let mut settings = JsonValue::Object(Map::new()); - let values: Vec<&str> = profile_matches - .get_many::("set") - .unwrap() - .map(AsRef::as_ref) - .collect(); + Some(("show", show_matches)) => { + if show_matches.contains_id("name") { + eprintln!("Name: {:?}", show_matches.get_one::("name")); + let profile_name = + show_matches.get_one::("name").unwrap(); + let show_decrypted = show_matches.get_flag("show-decrypted"); + let settings = db_handler + .get_profile_settings(profile_name, !show_decrypted) + .await?; + println!("Profile '{}' settings:", profile_name); + for (key, value) in settings.as_object().unwrap() { + println!(" {}: {}", key, value); + } + } else { + create_show_subcommand().print_help()?; + } + } - let mut i = 0; - while i < values.len() { - let key = values[i]; - let value = values[i + 1]; - let is_secure = profile_matches.get_flag("secure"); + Some(("set", set_matches)) => { + if set_matches.contains_id("name") + && set_matches.contains_id("key") + && set_matches.contains_id("value") + { + let profile_name = + set_matches.get_one::("name").unwrap(); + let key = set_matches.get_one::("key").unwrap(); + let value = set_matches.get_one::("value").unwrap(); + let is_secure = set_matches.get_flag("secure"); + let mut settings = JsonValue::Object(Map::new()); if is_secure { settings[key.to_string()] = JsonValue::Object(Map::from_iter(vec![ @@ -116,84 +154,99 @@ pub async fn handle_profile_subcommand( settings[key.to_string()] = JsonValue::String(value.to_string()); } - i += 2; - } - db_handler.create_or_update(profile_name, &settings).await?; - println!("Profile '{}' updated.", profile_name); - } else { - println!("Error: Profile name is required for set operation"); + db_handler.create_or_update(profile_name, &settings).await?; + println!( + "Profile '{}' updated. Key '{}' set.", + profile_name, key + ); + } else { + create_set_subcommand().print_help()?; + } } - } else if let Some(key) = profile_matches.get_one::("get") { - if let Some(profile_name) = profile_name { - let show_decrypted = profile_matches.get_flag("show-decrypted"); - let settings = db_handler - .get_profile_settings(profile_name, !show_decrypted) - .await?; - if let Some(value) = settings.get(key) { - println!("{}: {}", key, value); + + Some(("get", get_matches)) => { + if get_matches.contains_id("name") && get_matches.contains_id("key") + { + let profile_name = + get_matches.get_one::("name").unwrap(); + let key = get_matches.get_one::("key").unwrap(); + let show_decrypted = get_matches.get_flag("show-decrypted"); + let settings = db_handler + .get_profile_settings(profile_name, !show_decrypted) + .await?; + if let Some(value) = settings.get(key) { + println!("{}: {}", key, value); + } else { + println!( + "Key '{}' not found in profile '{}'", + key, profile_name + ); + } } else { + create_get_subcommand().print_help()?; + } + } + + Some(("del", del_matches)) => { + if del_matches.contains_id("name") && del_matches.contains_id("key") + { + let profile_name = + del_matches.get_one::("name").unwrap(); + let key = del_matches.get_one::("key").unwrap(); + + let mut settings = JsonValue::Object(Map::new()); + settings[key.to_string()] = JsonValue::Null; // Null indicates deletion + + db_handler.create_or_update(profile_name, &settings).await?; println!( - "Key '{}' not found in profile '{}'", + "Key '{}' deleted from profile '{}'.", key, profile_name ); + } else { + create_del_subcommand().print_help()?; } - } else { - println!("Error: Profile name is required for get operation"); } - } else if profile_matches.get_flag("show") { - if let Some(profile_name) = profile_name { - let show_decrypted = profile_matches.get_flag("show-decrypted"); - let settings = db_handler - .get_profile_settings(profile_name, !show_decrypted) - .await?; - println!("Profile '{}' settings:", profile_name); - for (key, value) in settings.as_object().unwrap() { - println!(" {}: {}", key, value); + + Some(("rm", rm_matches)) => { + if rm_matches.contains_id("name") { + let profile_name = + rm_matches.get_one::("name").unwrap(); + db_handler.delete_profile(profile_name).await?; + println!("Profile '{}' removed.", profile_name); + } else { + create_rm_subcommand().print_help()?; } - } else { - println!("Error: Profile name is required for show operation"); - } - } else if profile_matches.get_flag("delete") { - if let Some(profile_name) = profile_name { - db_handler.delete_profile(profile_name).await?; - println!("Profile '{}' deleted.", profile_name); - } else { - println!("Error: Profile name is required for delete operation"); } - } else if profile_matches.get_flag("default") { - if let Some(profile_name) = profile_name { - db_handler.set_default_profile(profile_name).await?; - println!("Profile '{}' set as default.", profile_name); - } else { - println!("Error: Profile name is required to set as default"); + + Some(("set-default", default_matches)) => { + if default_matches.contains_id("name") { + let profile_name = + default_matches.get_one::("name").unwrap(); + db_handler.set_default_profile(profile_name).await?; + println!("Profile '{}' set as default.", profile_name); + } else { + create_set_default_subcommand().print_help()?; + } } - } else if profile_name.is_some() { - // If a profile name is provided but no action is specified, show that profile - let profile_name = profile_name.unwrap(); - let show_decrypted = profile_matches.get_flag("show-decrypted"); - let settings = match db_handler - .get_profile_settings(profile_name, !show_decrypted) - .await - { - Ok(settings) => settings, - Err(ApplicationError::DatabaseError(_)) => { - // Profile doesn't exist, create it with empty settings - let empty_settings = JsonValue::Object(Map::new()); - db_handler - .create_or_update(profile_name, &empty_settings) - .await?; - empty_settings + + Some(("show-default", show_default_matches)) => { + if let Some(default_profile) = db_handler.get_default_profile().await? { + println!("Default profile: {}", default_profile); + let show_decrypted = show_default_matches.get_flag("show-decrypted"); + let settings = db_handler.get_profile_settings(&default_profile, !show_decrypted).await?; + println!("Settings:"); + for (key, value) in settings.as_object().unwrap() { + println!(" {}: {}", key, value); + } + } else { + println!("No default profile set."); } - Err(e) => return Err(e), - }; - println!("Profile '{}' settings:", profile_name); - for (key, value) in settings.as_object().unwrap() { - println!(" {}: {}", key, value); } - } else { - let mut profile_command = create_profile_subcommand(); - profile_command.print_help()?; + + _ => { + create_profile_subcommand().print_help()?; + } } Ok(())