profile encryption fixes

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

@@ -169,10 +169,6 @@ pub async fn run_cli(
         env.get_config_dir().expect("Config directory not defined");
     let sqlite_file = config_dir.join("chat.db");
 
-    // TODO: add support for custom encryption keys in conversations
-    //let encryption_handler =
-    //    EncryptionHandler::new_from_path(custom_key_path)?.map(Arc::new);
-
     let db_conn = Arc::new(ConversationDatabase::new(&sqlite_file, None)?);
 
     let mut profile_handler = db_conn.get_profile_handler(None);

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

@@ -22,12 +22,14 @@ use crate::external as lumni;
 pub struct EncryptionHandler {
     public_key: RsaPublicKey,
     private_key: RsaPrivateKey,
+    key_path: PathBuf,
 }
 
 impl EncryptionHandler {
     pub fn new(
         public_key_pem: &str,
         private_key_pem: &str,
+        key_path: PathBuf,
     ) -> Result<Self, ApplicationError> {
         let public_key = RsaPublicKey::from_public_key_pem(public_key_pem)
             .map_err(EncryptionError::from)?;
@@ -36,6 +38,7 @@ impl EncryptionHandler {
         Ok(Self {
             public_key,
             private_key,
+            key_path,
         })
     }
 
@@ -71,7 +74,11 @@ impl EncryptionHandler {
             .to_pkcs8_pem(LineEnding::LF)
             .map_err(|e| EncryptionError::Pkcs8Error(e))?;
 
-        Ok(Some(Self::new(&public_key_pem, &private_key_pem)?))
+        Ok(Some(Self::new(
+            &public_key_pem,
+            &private_key_pem,
+            private_key_path.clone(),
+        )?))
     }
 
     pub fn get_private_key_pem(&self) -> Result<String, ApplicationError> {
@@ -85,6 +92,10 @@ impl EncryptionHandler {
             })
     }
 
+    pub fn get_key_path(&self) -> &PathBuf {
+        &self.key_path
+    }
+
     fn is_encrypted_key(key_pem: &str) -> bool {
         key_pem.contains("ENCRYPTED")
     }
@@ -255,17 +266,6 @@ impl EncryptionHandler {
             .map_err(|e| ApplicationError::from(EncryptionError::RsaError(e)))
     }
 
-    pub fn get_ssh_private_key(&self) -> Result<Vec<u8>, EncryptionError> {
-        // Convert the RSA private key to PKCS#8 PEM format with LF line endings
-        let pem = self
-            .private_key
-            .to_pkcs8_pem(LineEnding::LF)
-            .map_err(EncryptionError::from)?;
-
-        // Convert the PEM string to bytes
-        Ok(pem.as_bytes().to_vec())
-    }
-
     pub fn encrypt_string(
         &self,
         data: &str,
@@ -374,10 +374,20 @@ impl EncryptionHandler {
     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 handler = EncryptionHandler::new_from_path(private_key_path)?
+            .ok_or_else(|| {
+                ApplicationError::EncryptionError(EncryptionError::InvalidKey(
+                    "No encryption handler available".to_string(),
+                ))
+            })?;
+
+        let key_data = handler.get_private_key_pem().map_err(|e| {
+            ApplicationError::EncryptionError(EncryptionError::InvalidKey(
+                e.to_string(),
+            ))
+        })?;
         let mut hasher = Sha256::new();
-        hasher.update(&file_content);
+        hasher.update(key_data.as_bytes());
         Ok(format!("{:x}", hasher.finalize()))
     }
 }
@@ -437,6 +447,7 @@ impl EncryptionHandler {
         Ok(Self {
             public_key,
             private_key,
+            key_path: key_path.clone(),
         })
     }
 }

lumni/src/apps/builtin/llm/prompt/src/chat/db/user_profile/database_operations.rs

@@ -1,6 +1,6 @@
 use std::path::PathBuf;
 
-use lumni::api::error::{ApplicationError, EncryptionError};
+use lumni::api::error::ApplicationError;
 use rusqlite::{params, OptionalExtension};
 
 use super::{DatabaseOperationError, EncryptionHandler, UserProfileDbHandler};
@@ -97,47 +97,18 @@ impl UserProfileDbHandler {
         .map_err(ApplicationError::from)
     }
 
-    pub async fn load_encryption_handler(
-        &self,
-        encryption_key_id: i64,
-    ) -> Result<EncryptionHandler, ApplicationError> {
-        let mut db = self.db.lock().await;
-        let key_path: String = db.process_queue_with_result(|tx| {
-            tx.query_row(
-                "SELECT file_path FROM encryption_keys WHERE id = ?",
-                params![encryption_key_id],
-                |row| row.get(0),
-            )
-            .map_err(DatabaseOperationError::SqliteError)
-        })?;
-
-        EncryptionHandler::new_from_path(&PathBuf::from(key_path))
-            .map_err(|e| {
-                ApplicationError::EncryptionError(EncryptionError::InvalidKey(
-                    e.to_string(),
-                ))
-            })?
-            .ok_or_else(|| {
-                ApplicationError::EncryptionError(EncryptionError::InvalidKey(
-                    "Failed to create encryption handler from key file"
-                        .to_string(),
-                ))
-            })
-    }
-
     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],
+                "INSERT INTO encryption_keys (name, file_path, sha256_hash) \
+                 VALUES (?, ?, ?)",
+                params![name, file_path.to_str().unwrap(), hash],
             )
             .map_err(|e| DatabaseOperationError::SqliteError(e))?;
             Ok(())
@@ -148,14 +119,14 @@ impl UserProfileDbHandler {
     pub async fn get_encryption_key(
         &self,
         name: &str,
-    ) -> Result<(String, String, String), ApplicationError> {
+    ) -> Result<(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 = ?",
+                "SELECT file_path, sha256_hash FROM encryption_keys WHERE \
+                 name = ?",
                 params![name],
-                |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)),
+                |row| Ok((row.get(0)?, row.get(1)?)),
             )
             .map_err(|e| DatabaseOperationError::SqliteError(e))
         })

lumni/src/apps/builtin/llm/prompt/src/chat/db/user_profile/encryption_operations.rs

@@ -21,12 +21,12 @@ impl UserProfileDbHandler {
                         EncryptionError::EncryptionFailed(e.to_string()),
                     )
                 })?;
-
             Ok(json!({
                 "content": encrypted_content,
                 "encryption_key": encryption_key,
             }))
         } else {
+            eprintln!("No encryption handler available");
             Ok(JsonValue::String(content.to_string()))
         }
     }
@@ -119,6 +119,7 @@ impl UserProfileDbHandler {
             ))
         }
     }
+
     pub fn generate_encryption_key(
         profile_name: &str,
     ) -> Result<(EncryptionHandler, PathBuf, String), ApplicationError> {
@@ -141,7 +142,6 @@ impl UserProfileDbHandler {
                     e.to_string(),
                 ))
             })?;
-
         Ok((encryption_handler, key_path, key_hash))
     }
 }

lumni/src/apps/builtin/llm/prompt/src/chat/db/user_profile/profile_operations.rs

@@ -1,8 +1,8 @@
 use std::path::PathBuf;
 use std::sync::Arc;
 
-use lumni::api::error::{ApplicationError, EncryptionError};
-use rusqlite::{params, OptionalExtension};
+use lumni::api::error::ApplicationError;
+use rusqlite::{params, OptionalExtension, Transaction};
 use serde_json::Value as JsonValue;
 
 use super::{
@@ -17,7 +17,9 @@ impl UserProfileDbHandler {
         profile_name: &str,
         new_settings: &JsonValue,
     ) -> Result<(), ApplicationError> {
-        let (encryption_key_id, merged_settings, new_encryption_handler) = {
+        let has_encryption_handler = self.encryption_handler.is_some();
+        let mut created_encryption_handler: Option<EncryptionHandler> = None;
+        let (encryption_key_id, merged_settings) = {
             let mut db = self.db.lock().await;
             db.process_queue_with_result(|tx| {
                 let existing_profile: Option<(i64, String, i64)> = tx
@@ -36,26 +38,60 @@ impl UserProfileDbHandler {
                             Some(existing_options),
                             new_settings,
                         )?;
-                        Ok((existing_key_id, merged, None))
-                    }
-                    None => {
-                        let (new_encryption_handler, key_path, key_hash) =
-                            Self::generate_encryption_key(profile_name)?;
 
-                        let key_id: i64 = tx
-                            .query_row(
-                                "INSERT INTO encryption_keys (file_path, \
-                                 sha256_hash) VALUES (?, ?) RETURNING id",
-                                params![key_path.to_str().unwrap(), key_hash],
-                                |row| row.get(0),
-                            )
-                            .map_err(DatabaseOperationError::SqliteError)?;
+                        // if encryption handler is not available, create a new one from the key path
+                        if !has_encryption_handler {
+                            let key_path: String = tx
+                                .query_row(
+                                    "SELECT file_path FROM encryption_keys \
+                                     WHERE id = ?",
+                                    params![existing_key_id],
+                                    |row| row.get(0),
+                                )
+                                .map_err(DatabaseOperationError::SqliteError)?;
 
-                        Ok((
-                            key_id,
-                            new_settings.clone(),
-                            Some(new_encryption_handler),
-                        ))
+                            let encryption_handler =
+                                EncryptionHandler::new_from_path(
+                                    &PathBuf::from(&key_path),
+                                )?
+                                .ok_or_else(
+                                    || {
+                                        ApplicationError::InvalidInput(
+                                            "Failed to load encryption handler"
+                                                .to_string(),
+                                        )
+                                    },
+                                )?;
+                            created_encryption_handler =
+                                Some(encryption_handler);
+                        }
+                        Ok((existing_key_id, merged))
+                    }
+                    None => {
+                        if !has_encryption_handler {
+                            let (new_encryption_handler, key_path, key_hash) =
+                                Self::generate_encryption_key(profile_name)?;
+                            let key_id = self.get_or_insert_encryption_key(
+                                tx, &key_path, &key_hash,
+                            )?;
+                            created_encryption_handler =
+                                Some(new_encryption_handler);
+                            Ok((key_id, new_settings.clone()))
+                        } else {
+                            let key_path = self
+                                .encryption_handler
+                                .as_ref()
+                                .unwrap()
+                                .get_key_path();
+                            let key_hash =
+                                EncryptionHandler::get_private_key_hash(
+                                    &key_path,
+                                )?;
+                            let key_id = self.get_or_insert_encryption_key(
+                                tx, &key_path, &key_hash,
+                            )?;
+                            Ok((key_id, new_settings.clone()))
+                        }
                     }
                 }
             })
@@ -67,21 +103,16 @@ impl UserProfileDbHandler {
             })?
         };
 
-        if let Some(encryption_handler) = new_encryption_handler {
-            self.encryption_handler = Some(Arc::new(encryption_handler));
-        } else if self.encryption_handler.is_none() {
-            let encryption_handler =
-                self.load_encryption_handler(encryption_key_id).await?;
-            self.encryption_handler = Some(Arc::new(encryption_handler));
-        }
-
         if self.encryption_handler.is_none() {
-            return Err(ApplicationError::EncryptionError(
-                EncryptionError::Other(Box::new(std::io::Error::new(
-                    std::io::ErrorKind::Other,
-                    "Failed to set up encryption handler",
-                ))),
-            ));
+            // Set new encryption handler outside the closure
+            if let Some(new_encryption_handler) = created_encryption_handler {
+                self.encryption_handler =
+                    Some(Arc::new(new_encryption_handler));
+            } else {
+                return Err(ApplicationError::InvalidInput(
+                    "Failed to create encryption handler".to_string(),
+                ));
+            }
         }
 
         let processed_settings = self.process_settings(
@@ -120,6 +151,37 @@ impl UserProfileDbHandler {
         Ok(())
     }
 
+    fn get_or_insert_encryption_key<'a>(
+        &self,
+        tx: &Transaction<'a>,
+        key_path: &PathBuf,
+        key_hash: &str,
+    ) -> Result<i64, DatabaseOperationError> {
+        // First, try to find an existing key with the same hash
+        let existing_key_id: Option<i64> = tx
+            .query_row(
+                "SELECT id FROM encryption_keys WHERE sha256_hash = ?",
+                params![key_hash],
+                |row| row.get(0),
+            )
+            .optional()
+            .map_err(DatabaseOperationError::SqliteError)?;
+
+        match existing_key_id {
+            Some(id) => Ok(id),
+            None => {
+                // If no existing key found, insert a new one
+                tx.query_row(
+                    "INSERT INTO encryption_keys (file_path, sha256_hash) \
+                     VALUES (?, ?) RETURNING id",
+                    params![key_path.to_str().unwrap(), key_hash],
+                    |row| row.get(0),
+                )
+                .map_err(DatabaseOperationError::SqliteError)
+            }
+        }
+    }
+
     pub async fn get_profile_settings(
         &mut self,
         profile_name: &str,
@@ -141,7 +203,6 @@ impl UserProfileDbHandler {
                 .map_err(DatabaseOperationError::SqliteError)
             })?
         };
-
         if self.encryption_handler.is_none() {
             let encryption_handler =
                 EncryptionHandler::new_from_path(&PathBuf::from(&key_path))?