Add Token Encryption

Cargo.toml

@@ -23,6 +23,12 @@ libsqlite3-sys = { version = "0.27.0", features = ["bundled"] }
 beam-lib = { git = "https://github.com/samply/beam", branch = "develop", features = ["http-util"] }
 async-sse = "5.1.0"
 futures-util = { version = "0.3", features = ["io"] }
+#encrypt 
+rust-crypto = "^0.2"
+aes = "0.8.4"
+ctr = "0.9.2"
+cipher = "0.4.4"
+base64 = "0.22.0"
 
 # Logging
 tracing = { version = "0.1" }

src/db.rs

@@ -13,6 +13,7 @@ use crate::schema::tokens;
 use crate::enums::{OpalTokenStatus, OpalProjectStatus};
 use crate::handlers::{check_project_status_request, fetch_project_tables_names_request, check_token_status_request};
 use crate::models::{NewToken, TokenManager, TokenParams, TokenStatus, TokensQueryParams, ProjectQueryParams};
+use crate::utils::{decrypt_data, generate_r_script};
 
 const MIGRATIONS: EmbeddedMigrations = embed_migrations!();
 
@@ -62,28 +63,35 @@ impl Db {
     }
 
     pub fn update_token_db(&mut self, token_update: NewToken) {
-
-
-        let target = tokens.filter(
-            user_id.eq(&token_update.user_id)
+        // Attempt to find the `id` of the last record matching the criteria
+        let maybe_last_id = tokens
+            .filter(
+                user_id.eq(&token_update.user_id)
                 .and(project_id.eq(&token_update.project_id))
                 .and(bk.eq(&token_update.bk))
-        );
-
-        match diesel::update(target)
-            .set((
-                token.eq(token_update.token),
-                token_status.eq("UPDATED"),
-                token_created_at.eq(&token_update.token_created_at),
-            ))
-            .execute(&mut self.0)
-        {
-            Ok(_) => {
-                info!("Token updated in DB");
-            }
-            Err(error) => {
-                warn!("Error updating token: {}", error);
+            )
+            .select(id)
+            .order(id.desc())
+            .first::<i32>(&mut self.0) // Assuming `id` is of type `i32`
+            .optional(); // Use `.optional()` to handle cases where no records are found
+
+        if let Ok(Some(last_id)) = maybe_last_id {
+            // If a last record is found, perform the update
+            let target = tokens.filter(id.eq(last_id));
+            match diesel::update(target)
+                .set((
+                    token.eq(&token_update.token),
+                    token_status.eq("UPDATED"),
+                    token_created_at.eq(&token_update.token_created_at),
+                ))
+                .execute(&mut self.0)
+            {
+                Ok(_) => info!("Token updated in DB"),
+                Err(error) => warn!("Error updating token: {}", error),
             }
+        } else if let Err(error) = maybe_last_id {
+            // Handle potential errors from the `first` query
+            warn!("Error finding last token record: {}", error);
         }
     }
 
@@ -153,15 +161,17 @@ impl Db {
         tokens
             .filter(user_id.eq(user))
             .filter(project_id.eq(project))
+            .order(id.desc())
             .select(token_name)
             .first::<String>(&mut self.0)
-            .optional() 
+            .optional()
     }
 
     pub fn get_token_value(&mut self, user: String, project: String) -> Result<Option<String>, Error> {
         tokens
             .filter(user_id.eq(user))
             .filter(project_id.eq(project))
+            .order(id.desc())
             .select(token)
             .first::<String>(&mut self.0)
             .optional() 
@@ -274,11 +284,12 @@ impl Db {
 
                     match records_result {
                         Ok(record) => {
+                            let token_decrypt = decrypt_data(record.token.clone(), &record.token_name.clone().as_bytes()[..16]); 
                                 for table in &tables {
                                     let site_name = record.bk.split('.').nth(1).expect("Valid app id");
                                     script_lines.push(format!(
                                         "builder$append(server='{}', url='https://{}/opal/', token='{}', table='{}', driver='OpalDriver')",
-                                        site_name, record.bk, record.token, table
+                                        site_name, record.bk, token_decrypt, table
                                     ));
                                 }
                             }
@@ -302,31 +313,4 @@ impl Db {
             },
         }
     }
-}    
-
-fn generate_r_script(script_lines: Vec<String>) -> String {
-    let mut builder_script = String::from(
-        r#"library(DSI)
-library(DSOpal)
-library(dsBaseClient)
-set_config(use_proxy(url="http://beam-connect", port=8062))
-set_config( config( ssl_verifyhost = 0L, ssl_verifypeer = 0L ) )
-
-builder <- DSI::newDSLoginBuilder(.silent = FALSE)
-"#,
-    );
-
-    // Append each line to the script.
-    for line in script_lines {
-        builder_script.push_str(&line);
-        builder_script.push('\n');
-    }
-
-    // Finish the script with the login and assignment commands.
-    builder_script.push_str(
-        "logindata <- builder$build()
-connections <- DSI::datashield.login(logins = logindata, assign = TRUE, symbol = 'D')\n",
-    );
-
-    builder_script
-}
+}    

src/handlers.rs

@@ -6,6 +6,8 @@ use crate::config::CONFIG;
 use crate::db::Db;
 use crate::enums::{OpalResponse, OpalProjectTablesResponse, OpalProjectStatus, OpalProjectStatusResponse, OpalTokenStatus, OpalRequestType};
 use crate::models::{NewToken, OpalRequest, TokenParams, TokensQueryParams, ProjectQueryParams};
+use crate::utils::{encrypt_data, decrypt_data};
+use base64::{engine::general_purpose::STANDARD, Engine};
 use anyhow::Result;
 use async_sse::Event;
 use axum::http::StatusCode;
@@ -119,7 +121,7 @@ pub async fn refresh_token_request(mut db: Db, token_params: TokenParams) -> Res
     };
 
     let token_value = match db.get_token_value(token_params.user_id.clone(), token_params.project_id.clone()) {
-        Ok(Some(value)) => value,
+        Ok(Some(value)) => decrypt_data(value, &token_name.clone().as_bytes()[..16]),
         Ok(None) => {
             return Err(anyhow::Error::msg("Token value not found")) 
         },
@@ -136,8 +138,6 @@ pub async fn refresh_token_request(mut db: Db, token_params: TokenParams) -> Res
         Some(token_value.clone())
     ).await?;
 
-    info!("Refresh token task  {task:#?}");
-
     tokio::task::spawn(update_tokens_from_beam(db, task, token_params, token_name.clone()));
     Ok(OpalResponse::Ok { token: "OK".to_string() })
 }
@@ -290,10 +290,13 @@ async fn save_tokens_from_beam(mut db: Db, task: TaskRequest<OpalRequest>, token
                 last_error = Some(format!("Error: {error}"));
             },
             OpalResponse::Ok { token } => {
-                let site_name = result.from.as_ref();//.split('.').nth(1).expect("Valid app id");
+                let encryp_token = encrypt_data(&token.clone().as_bytes(), &token_name.clone().as_bytes()[..16]);
+                let token_encoded =  STANDARD.encode(encryp_token); 
+                let site_name = result.from.as_ref();
+
                 let new_token = NewToken {
                     token_name: &token_name,
-                    token: &token,
+                    token: &token_encoded,
                     project_id: &token_params.project_id,
                     bk: &site_name,
                     token_status: OpalTokenStatus::CREATED.as_str(),
@@ -345,10 +348,13 @@ async fn update_tokens_from_beam(mut db: Db, task: TaskRequest<OpalRequest>, tok
                 last_error = Some(format!("Error: {error}"));
             },
             OpalResponse::Ok { token } => {
-                let site_name = result.from.as_ref(); //.split('.').nth(1).expect("Valid app id");
+                let encryp_token = encrypt_data(&token.clone().as_bytes(), &token_name.clone().as_bytes()[..16]);
+                let token_encoded =  STANDARD.encode(encryp_token); 
+                let site_name = result.from.as_ref(); 
+
                 let new_token = NewToken {
                     token_name: &token_name,
-                    token: &token,
+                    token: &token_encoded,
                     project_id: &token_params.project_id,
                     bk: &site_name,
                     token_status: OpalTokenStatus::CREATED.as_str(),

src/main.rs

@@ -5,6 +5,7 @@ mod models;
 mod routes;
 mod schema;
 mod enums;
+mod utils;
 
 use crate::config::CONFIG;
 use axum::Router;

src/utils.rs

@@ -0,0 +1,52 @@
+use aes::Aes256; 
+use ctr::Ctr128BE; 
+use cipher::{KeyIvInit, StreamCipher};
+use base64::{engine::general_purpose::STANDARD, Engine};
+
+
+pub fn encrypt_data(data: &[u8], nonce: &[u8]) -> Vec<u8> {
+
+    let key: &[u8; 32] = b"0123456789abcdef0123456789ABCDEF";     
+    let mut cipher = Ctr128BE::<Aes256>::new_from_slices(key, nonce).unwrap();
+    let mut encrypted = data.to_vec();
+    cipher.apply_keystream(&mut encrypted);
+
+    encrypted
+}
+
+pub fn decrypt_data(data: String, nonce: &[u8]) -> String {
+    let toke_decode = STANDARD.decode(data).unwrap();
+    let decrypted_token =  encrypt_data(&toke_decode, nonce);
+    let token_str = String::from_utf8(decrypted_token).unwrap();
+    token_str 
+}
+
+pub fn generate_r_script(script_lines: Vec<String>) -> String {
+    let mut builder_script = String::from(
+        r#"library(DSI)
+library(DSOpal)
+library(dsBaseClient)
+set_config(use_proxy(url="http://beam-connect", port=8062))
+set_config( config( ssl_verifyhost = 0L, ssl_verifypeer = 0L ) )
+
+builder <- DSI::newDSLoginBuilder(.silent = FALSE)
+"#,
+    );
+
+    // Append each line to the script.
+    for line in script_lines {
+        builder_script.push_str(&line);
+        builder_script.push('\n');
+    }
+
+    // Finish the script with the login and assignment commands.
+    builder_script.push_str(
+        "logindata <- builder$build()
+connections <- DSI::datashield.login(logins = logindata, assign = TRUE, symbol = 'D')\n",
+    );
+
+    builder_script
+}
+
+
+