add ability to manage ssh-keys via profile, add profile helper

lumni/src/apps/builtin/llm/prompt/src/app.rs

@@ -169,6 +169,7 @@ pub async fn run_cli(
         env.get_config_dir().expect("Config directory not defined");
     let sqlite_file = config_dir.join("chat.db");
 
+    // TODO: None uses default ssh key path, add support for custom ssh key path
     let encryption_handler =
         EncryptionHandler::new_from_path(None)?.map(Arc::new);
 

lumni/src/apps/builtin/llm/prompt/src/chat/db/connector.rs

@@ -2,11 +2,12 @@ use std::collections::VecDeque;
 use std::path::PathBuf;
 use std::sync::{Arc, Mutex};
 
-use lumni::api::error::{ApplicationError, EncryptionError};
+use lumni::api::error::ApplicationError;
 use rusqlite::{params, Error as SqliteError, Transaction};
 
 use crate::external as lumni;
 
+#[derive(Debug)]
 pub struct DatabaseConnector {
     connection: rusqlite::Connection,
     operation_queue: Arc<Mutex<VecDeque<String>>>,

lumni/src/apps/builtin/llm/prompt/src/chat/db/encryption/mod.rs

@@ -13,6 +13,7 @@ use rsa::pkcs8::{
     LineEnding,
 };
 use rsa::{BigUint, RsaPrivateKey, RsaPublicKey};
+use sha2::{Digest, Sha256};
 
 use crate::external as lumni;
 
@@ -338,4 +339,14 @@ impl EncryptionHandler {
 
         Ok(decrypted_data.to_vec())
     }
+
+    pub fn get_private_key_hash(
+        private_key_path: &PathBuf,
+    ) -> Result<String, ApplicationError> {
+        let file_content = fs::read(private_key_path)
+            .map_err(|e| ApplicationError::IOError(e))?;
+        let mut hasher = Sha256::new();
+        hasher.update(&file_content);
+        Ok(format!("{:x}", hasher.finalize()))
+    }
 }

lumni/src/apps/builtin/llm/prompt/src/chat/db/schema.sql

@@ -12,6 +12,14 @@ CREATE TABLE user_profiles (
     is_default INTEGER DEFAULT 0
 );
 
+CREATE TABLE encryption_keys (
+    id INTEGER PRIMARY KEY,
+    name TEXT UNIQUE NOT NULL,
+    file_path TEXT NOT NULL,
+    sha256_hash TEXT NOT NULL,
+    key_type TEXT NOT NULL
+);
+
 CREATE TABLE models (
     identifier TEXT PRIMARY KEY,
     info TEXT, -- JSON string

lumni/src/apps/builtin/llm/prompt/src/chat/db/user_profiles/mod.rs

@@ -1,3 +1,4 @@
+use std::path::PathBuf;
 use std::sync::Arc;
 
 use base64::engine::general_purpose;
@@ -12,6 +13,7 @@ use super::connector::{DatabaseConnector, DatabaseOperationError};
 use super::encryption::EncryptionHandler;
 use crate::external as lumni;
 
+#[derive(Debug)]
 pub struct UserProfileDbHandler {
     profile_name: Option<String>,
     db: Arc<TokioMutex<DatabaseConnector>>,
@@ -39,6 +41,17 @@ impl UserProfileDbHandler {
         self.profile_name = Some(profile_name);
     }
 
+    pub fn get_encryption_handler(&self) -> Option<&Arc<EncryptionHandler>> {
+        self.encryption_handler.as_ref()
+    }
+
+    pub fn set_encryption_handler(
+        &mut self,
+        encryption_handler: Arc<EncryptionHandler>,
+    ) {
+        self.encryption_handler = Some(encryption_handler);
+    }
+
     fn encrypt_value(
         &self,
         content: &str,
@@ -210,49 +223,6 @@ impl UserProfileDbHandler {
         })?
     }
 
-    fn fetch_and_process_settings(
-        &self,
-        tx: &Transaction,
-        profile_name: &str,
-        mask_encrypted: bool,
-    ) -> Result<JsonValue, DatabaseOperationError> {
-        let (json_string, ssh_key_hash) =
-            self.fetch_profile_data(tx, profile_name)?;
-        let settings: JsonValue = self.parse_json(&json_string)?;
-
-        if let Some(hash) = ssh_key_hash {
-            self.verify_ssh_key_hash(Some(&hash))
-                .map_err(DatabaseOperationError::ApplicationError)?;
-        }
-
-        self.process_settings(&settings, false, mask_encrypted)
-            .map_err(DatabaseOperationError::ApplicationError)
-    }
-
-    fn fetch_profile_data(
-        &self,
-        tx: &Transaction,
-        profile_name: &str,
-    ) -> Result<(String, Option<String>), DatabaseOperationError> {
-        tx.query_row(
-            "SELECT options, ssh_key_hash FROM user_profiles WHERE name = ?",
-            params![profile_name],
-            |row| Ok((row.get(0)?, row.get(1)?)),
-        )
-        .map_err(DatabaseOperationError::SqliteError)
-    }
-
-    fn parse_json(
-        &self,
-        json_string: &str,
-    ) -> Result<JsonValue, DatabaseOperationError> {
-        serde_json::from_str(json_string).map_err(|e| {
-            DatabaseOperationError::ApplicationError(
-                ApplicationError::InvalidInput(format!("Invalid JSON: {}", e)),
-            )
-        })
-    }
-
     fn process_settings(
         &self,
         value: &JsonValue,
@@ -294,35 +264,6 @@ impl UserProfileDbHandler {
         }
     }
 
-    fn process_object(
-        &self,
-        obj: &Map<String, JsonValue>,
-        encrypt: bool,
-        mask_encrypted: bool,
-    ) -> Result<JsonValue, ApplicationError> {
-        let mut new_obj = Map::new();
-        for (k, v) in obj {
-            new_obj.insert(
-                k.clone(),
-                self.process_value(v, encrypt, mask_encrypted)?,
-            );
-        }
-        Ok(JsonValue::Object(new_obj))
-    }
-
-    fn process_array(
-        &self,
-        arr: &[JsonValue],
-        encrypt: bool,
-        mask_encrypted: bool,
-    ) -> Result<JsonValue, ApplicationError> {
-        let new_arr: Result<Vec<JsonValue>, _> = arr
-            .iter()
-            .map(|v| self.process_settings(v, encrypt, mask_encrypted))
-            .collect();
-        Ok(JsonValue::Array(new_arr?))
-    }
-
     fn handle_encryption(
         &self,
         value: &JsonValue,
@@ -568,3 +509,92 @@ impl UserProfileDbHandler {
         }
     }
 }
+
+impl UserProfileDbHandler {
+    pub async fn register_encryption_key(
+        &self,
+        name: &str,
+        file_path: &PathBuf,
+        key_type: &str,
+    ) -> Result<(), ApplicationError> {
+        let hash = EncryptionHandler::get_private_key_hash(file_path)?;
+        let mut db = self.db.lock().await;
+        db.process_queue_with_result(|tx| {
+            tx.execute(
+                "INSERT INTO encryption_keys (name, file_path, sha256_hash, \
+                 key_type) VALUES (?, ?, ?, ?)",
+                params![name, file_path.to_str().unwrap(), hash, key_type],
+            )
+            .map_err(|e| DatabaseOperationError::SqliteError(e))?;
+            Ok(())
+        })
+        .map_err(ApplicationError::from)
+    }
+
+    pub async fn get_encryption_key(
+        &self,
+        name: &str,
+    ) -> Result<(String, String, String), ApplicationError> {
+        let mut db = self.db.lock().await;
+        db.process_queue_with_result(|tx| {
+            tx.query_row(
+                "SELECT file_path, sha256_hash, key_type FROM encryption_keys \
+                 WHERE name = ?",
+                params![name],
+                |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)),
+            )
+            .map_err(|e| DatabaseOperationError::SqliteError(e))
+        })
+        .map_err(ApplicationError::from)
+    }
+
+    pub async fn remove_encryption_key(
+        &self,
+        name: &str,
+    ) -> Result<(), ApplicationError> {
+        let mut db = self.db.lock().await;
+        db.process_queue_with_result(|tx| {
+            tx.execute(
+                "DELETE FROM encryption_keys WHERE name = ?",
+                params![name],
+            )
+            .map_err(|e| DatabaseOperationError::SqliteError(e))?;
+            Ok(())
+        })
+        .map_err(ApplicationError::from)
+    }
+
+    pub async fn list_encryption_keys(
+        &self,
+        key_type: Option<&str>,
+    ) -> Result<Vec<String>, ApplicationError> {
+        let mut db = self.db.lock().await;
+        db.process_queue_with_result(|tx| {
+            let query = match key_type {
+                Some(_) => {
+                    "SELECT name FROM encryption_keys WHERE key_type = ?"
+                }
+                None => "SELECT name FROM encryption_keys",
+            };
+
+            let mut stmt = tx
+                .prepare(query)
+                .map_err(|e| DatabaseOperationError::SqliteError(e))?;
+
+            let row_mapper = |row: &rusqlite::Row| row.get(0);
+
+            let rows = match key_type {
+                Some(ktype) => stmt.query_map(params![ktype], row_mapper),
+                None => stmt.query_map([], row_mapper),
+            }
+            .map_err(|e| DatabaseOperationError::SqliteError(e))?;
+
+            let keys = rows
+                .collect::<Result<Vec<String>, _>>()
+                .map_err(|e| DatabaseOperationError::SqliteError(e))?;
+
+            Ok(keys)
+        })
+        .map_err(ApplicationError::from)
+    }
+}

lumni/src/apps/builtin/llm/prompt/src/cli/mod.rs

@@ -7,7 +7,9 @@ pub use subcommands::db::handle_db_subcommand;
 use subcommands::profile::create_profile_subcommand;
 pub use subcommands::profile::handle_profile_subcommand;
 
-use super::chat::db::{ConversationDatabase, UserProfileDbHandler};
+use super::chat::db::{
+    ConversationDatabase, EncryptionHandler, UserProfileDbHandler,
+};
 use crate::external as lumni;
 
 pub fn parse_cli_arguments(spec: ApplicationSpec) -> Command {

lumni/src/apps/builtin/llm/prompt/src/cli/subcommands/mod.rs

@@ -1,4 +1,5 @@
 pub mod db;
 pub mod profile;
+pub mod profile_helper;
 
-use super::{ConversationDatabase, UserProfileDbHandler};
+use super::{ConversationDatabase, EncryptionHandler, UserProfileDbHandler};