diff --git a/Cargo.toml b/Cargo.toml index 6f41e8a..b9df850 100644 --- a/Cargo.toml +++ b/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" } diff --git a/src/db.rs b/src/db.rs index 3945761..8091462 100644 --- a/src/db.rs +++ b/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::(&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::(&mut self.0) - .optional() + .optional() } pub fn get_token_value(&mut self, user: String, project: String) -> Result, Error> { tokens .filter(user_id.eq(user)) .filter(project_id.eq(project)) + .order(id.desc()) .select(token) .first::(&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 { - 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 -} +} \ No newline at end of file diff --git a/src/handlers.rs b/src/handlers.rs index fde092f..81e1cb5 100644 --- a/src/handlers.rs +++ b/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, 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, 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(), diff --git a/src/main.rs b/src/main.rs index b63c367..6992d79 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,6 +5,7 @@ mod models; mod routes; mod schema; mod enums; +mod utils; use crate::config::CONFIG; use axum::Router; diff --git a/src/utils.rs b/src/utils.rs new file mode 100644 index 0000000..5daf94c --- /dev/null +++ b/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 { + + let key: &[u8; 32] = b"0123456789abcdef0123456789ABCDEF"; + let mut cipher = Ctr128BE::::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 { + 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 +} + + +