Adapt user script generation (#7)

samply · Feb 13, 2025 · 156db72 · 156db72
1 parent cd958a8
commit 156db72
Show file tree

Hide file tree

Showing 6 changed files with 133 additions and 116 deletions.
diff --git a/.env b/.env
@@ -2,3 +2,4 @@ BEAM_URL=http://localhost:8081
 BEAM_ID=app1.proxy1.broker
 BEAM_SECRET=App1Secret
 TOKEN_MANAGER_DB_PATH=./file.db  
+AUTH_SCRIPT_TEMPLATE_PATH=./auth-script-template.R
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -0,0 +1,10 @@
+# Change Log
+All notable changes to this project will be documented in this file.
+
+The format is based on [Keep a Changelog](http://keepachangelog.com/)
+and this project adheres to [Semantic Versioning](http://semver.org/).
+
+## [1.0.0 - 2025-02-11]
+### Changed
+- Generate user script
+
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "token-manager"
-version = "0.1.0"
+version = "1.0.0"
 edition = "2021"
 license = "Apache-2.0"
 

diff --git a/src/config.rs b/src/config.rs
@@ -32,6 +32,9 @@ pub struct Config {
     #[clap(env, default_value = "0123456789abcdef0123456789ABCDEF")]
     pub token_encrypt_key: String,
 
+    #[clap(long, env)]
+    pub auth_script_template_path: String,
+
     #[clap(long, env, default_value = "info")]
     pub rust_log: String,
 }

diff --git a/src/db.rs b/src/db.rs
@@ -10,7 +10,7 @@ use diesel::r2d2::{ConnectionManager, Pool, PooledConnection};
 use diesel::result::Error;
 use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness};
 use serde_json::json;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use tracing::{error, info, warn};
 
 use crate::config::CONFIG;
@@ -23,7 +23,7 @@ use crate::models::{
 };
 use crate::schema::tokens;
 use crate::schema::tokens::dsl::*;
-use crate::utils::{decrypt_data, generate_r_script};
+use crate::utils::{decrypt_data, fetch_tables_prefix, generate_r_script};
 
 const MIGRATIONS: EmbeddedMigrations = embed_migrations!();
 
@@ -65,7 +65,10 @@ impl Db {
             .execute(&mut self.0)
         {
             Ok(_) => {
-                info!("Token Saved in DB for user: {} in BK: {}", new_token.user_id, new_token.bk);
+                info!(
+                    "Token Saved in DB for user: {} in BK: {}",
+                    new_token.user_id, new_token.bk
+                );
             }
             Err(error) => {
                 warn!("Error connecting to {}", error);
@@ -84,7 +87,7 @@ impl Db {
             .select(id)
             .order(id.desc())
             .first::<i32>(&mut self.0)
-            .optional(); 
+            .optional();
 
         if let Ok(Some(last_id)) = maybe_last_id {
             let target = tokens.filter(id.eq(last_id));
@@ -96,7 +99,10 @@ impl Db {
                 ))
                 .execute(&mut self.0)
             {
-                Ok(_) => info!("Token Updated in DB for user: {} in BK: {}", token_update.user_id, token_update.bk),
+                Ok(_) => info!(
+                    "Token Updated in DB for user: {} in BK: {}",
+                    token_update.user_id, token_update.bk
+                ),
                 Err(error) => warn!("Error updating token: {}", error),
             }
         } else if let Err(error) = maybe_last_id {
@@ -117,20 +123,26 @@ impl Db {
             .execute(&mut self.0)
         {
             Ok(_) => {
-                info!("Token status updated in DB for user: {} in BK: {}", token_update.user_id, token_update.bk);
+                info!(
+                    "Token status updated in DB for user: {} in BK: {}",
+                    token_update.user_id, token_update.bk
+                );
             }
             Err(error) => {
                 warn!("Error updating token status: {}", error);
             }
         }
     }
 
-    pub fn delete_project_db(&mut self, project_params: &ProjectQueryParams,) {
+    pub fn delete_project_db(&mut self, project_params: &ProjectQueryParams) {
         let target = tokens.filter(project_id.eq(&project_params.project_id));
 
         match diesel::delete(target).execute(&mut self.0) {
             Ok(_) => {
-                info!("Project and Token deleted from DB for project: {} in BK: {}", &project_params.project_id, &project_params.bk);
+                info!(
+                    "Project and Token deleted from DB for project: {} in BK: {}",
+                    &project_params.project_id, &project_params.bk
+                );
             }
             Err(error) => {
                 warn!("Error deleting token: {}", error);
@@ -143,7 +155,10 @@ impl Db {
 
         match diesel::delete(target).execute(&mut self.0) {
             Ok(_) => {
-                info!("Token deleted from DB for user: {} in BK: {}", token_params.user_id, token_params.bk);
+                info!(
+                    "Token deleted from DB for user: {} in BK: {}",
+                    token_params.user_id, token_params.bk
+                );
             }
             Err(error) => {
                 warn!("Error deleting token: {}", error);
@@ -206,11 +221,11 @@ impl Db {
             Ok(true) => {
                 info!("Token available for user: {}", params.user_id);
                 Ok("true".to_string())
-            },
+            }
             Ok(false) => {
                 info!("No Token available for user: {}", params.user_id);
                 Ok("false".to_string())
-            },
+            }
             Err(_) => Err((
                 StatusCode::INTERNAL_SERVER_ERROR,
                 "Error checking token availability.".to_string(),
@@ -235,7 +250,7 @@ impl Db {
             bk: params.bk.clone(),
             project_id: params.project_id.clone(),
         })
-        .await
+            .await
         {
             token_status_json["project_status"] = json_response.0["project_status"].clone();
         } else {
@@ -244,15 +259,13 @@ impl Db {
 
         let token_name_response = match self.get_token_name(&params) {
             Ok(Some(name)) => name,
-            Ok(None) =>{ 
+            Ok(None) => {
                 info!(
                     "Received status response for token. User ID: {}, BK: {}, Response: {}",
-                    params.user_id,
-                    params.bk,
-                    token_status_json["token_status"]
+                    params.user_id, params.bk, token_status_json["token_status"]
                 );
-                return Ok(Json(token_status_json))
-            },
+                return Ok(Json(token_status_json));
+            }
             Err(e) => return Err((StatusCode::INTERNAL_SERVER_ERROR, e.to_string())),
         };
 
@@ -267,9 +280,7 @@ impl Db {
             Ok(_) => {
                 info!(
                     "Received status response for token. User ID: {}, BK: {}, Response: {}",
-                    params.user_id,
-                    params.bk,
-                    token_status_json["token_status"]
+                    params.user_id, params.bk, token_status_json["token_status"]
                 );
                 return Ok(Json(token_status_json));
             }
@@ -290,7 +301,7 @@ impl Db {
             token_name_response.clone(),
             token_value.clone(),
         )
-        .await
+            .await
         {
             token_status_json["token_status"] = json_response.0["token_status"].clone();
 
@@ -307,87 +318,67 @@ impl Db {
 
         info!(
             "Received status response for token. User ID: {}, BK: {}, Response: {}",
-            params.user_id,
-            params.bk,
-            token_status_json["token_status"]
+            params.user_id, params.bk, token_status_json["token_status"]
         );
 
         Ok(Json(token_status_json))
     }
 
     pub async fn generate_user_script(&mut self, query: TokenParams) -> Result<String, String> {
         let tables_per_bridgehead_result = fetch_project_tables_names_request(query.clone()).await;
-
         match tables_per_bridgehead_result {
-            Ok(tables_per_bridgehead) => {
-                let mut script_lines = Vec::new();
-                let all_tables = tables_per_bridgehead
-                    .values()
-                    .flat_map(|tables| tables.iter())
-                    .collect::<HashSet<_>>();
-
-                for bridgehead in &query.bridgehead_ids {
-                    let records_result = tokens
-                        .filter(project_id.eq(&query.project_id))
-                        .filter(user_id.eq(&query.user_id))
-                        .filter(bk.eq(bridgehead))
-                        .order(id.desc())
-                        .select(TokenManager::as_select())
-                        .first::<TokenManager>(&mut self.0);
-
-                    match records_result {
-                        Ok(record) => {
-                            let token_decrypt = decrypt_data(
-                                record.token.clone(),
-                                &record.token_name.clone().as_bytes()[..16],
-                            );
-                            if let Some(tables) = tables_per_bridgehead.get(bridgehead) {
-                                let tables_set: HashSet<_> = tables.iter().collect();
-                                let missing_tables: HashSet<_> =
-                                    all_tables.difference(&tables_set).collect();
-                                if !missing_tables.is_empty() {
-                                    info!(
-                                        "Bridgehead {} is missing tables: {:?}",
-                                        bridgehead, missing_tables
-                                    );
-                                    script_lines.push(format!(
-                                        "\n # Tables not available for bridgehead '{}': {:?}",
-                                        bridgehead, missing_tables
-                                    ));
-                                }
-
-                                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, token_decrypt, table.clone()
-                                    ));
-                                }
-                                script_lines.push("".to_string());
-                            }
-                        }
-                        Err(_) => {
-                            info!("Token not available for Bridgehead {}", bridgehead);
-                            script_lines.push(format!(
-                                "\n # Token not available for bridgehead '{}'",
-                                bridgehead
-                            ));
-                        }
-                    }
-                }
-
-                if !script_lines.is_empty() {
-                    let script = generate_r_script(script_lines);
-                    Ok(script)
-                } else {
-                    Ok("No records found for the given project and user.".into())
-                }
-            }
+            Ok(tables_per_bridgehead) => 
+                self.generate_user_script_using_tables(query, tables_per_bridgehead).await,
             Err(e) => {
                 error!("Error in fetch_project_tables_names_request: {:?}", e);
                 Err("Error obtaining table names.".into())
             }
         }
     }
-}
+
+    async fn generate_user_script_using_tables(&mut self, query: TokenParams, bridgehead_tables: HashMap<String, HashSet<String>>) -> Result<String, String> {
+        let mut script_lines = Vec::new();
+        script_lines.push("\"SiteName\",\"URL\",\"ProjectName\",\"Token\"".to_string());
+        for bridgehead in &query.bridgehead_ids {
+            let records_result = tokens
+                .filter(project_id.eq(&query.project_id))
+                .filter(user_id.eq(&query.user_id))
+                .filter(bk.eq(bridgehead))
+                .order(id.desc())
+                .select(TokenManager::as_select())
+                .first::<TokenManager>(&mut self.0);
+
+            match records_result {
+                Ok(record) => {
+                    let token_decrypt = decrypt_data(
+                        record.token.clone(),
+                        &record.token_name.clone().as_bytes()[..16],
+                    );
+                    let site_name = record.bk.split('.').nth(1).expect("Valid app id");
+                    let tables_prefix = fetch_tables_prefix(&bridgehead_tables, bridgehead, &query.project_id);
+                    // TODO: Maybe in the future, it makes sense to pass record.bk instead of site_name as URL
+                    // e.g. "https://token-manager.dktk-test.broker.ccp-it.dktk.dkfz.de/opal/" instead of "https://dktk-test/opal/"
+                    // It looks more like an absolute beam path. In more complex beam contexts, it would avoid ambiguity.
+                    script_lines.push(format!("\"{}\",\"https://{}/opal/\",\"{}\",\"{}\"",
+                                              site_name, site_name, tables_prefix, token_decrypt));
+                }
+                Err(_) => {
+                    info!("Token not available for Bridgehead {}", bridgehead);
+                    script_lines.push(format!(
+                        "\n # Token not available for bridgehead '{}'",
+                        bridgehead
+                    ));
+                }
+            }
+        }
+        if !script_lines.is_empty() {
+            match generate_r_script(script_lines.join("\n")) {
+                Ok(script) => Ok(script), // Return the script if successful
+                Err(_) => Err("# Failed to generate user script due to template error.".to_string()),
+            }
+        } else {
+            Ok("No records found for the given project and user.".into())
+        }
+    }
+
+}
diff --git a/src/utils.rs b/src/utils.rs
@@ -1,3 +1,5 @@
+use std::{fs, io};
+use std::collections::{HashMap, HashSet};
 use crate::config::CONFIG;
 use aes::Aes256;
 use base64::{engine::general_purpose::STANDARD, Engine};
@@ -29,29 +31,39 @@ pub fn decrypt_data(data: String, nonce: &[u8]) -> String {
     String::from_utf8(decrypted_token).unwrap()
 }
 
-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');
-    }
+pub fn generate_r_script(script_config: String) -> Result<String, io::Error> {
+    // Read the auth script template from the file
+    let template_path = &CONFIG.auth_script_template_path;
+    let template_content = fs::read_to_string(template_path)?;
+
+    // Replace the placeholder with the credentials CSV
+    let processed_content = template_content.replace("${CSV_CREDENTIALS_CONFIG}", script_config.as_str());
+
+    // Return the processed content
+    Ok(processed_content)
+}
 
-    // 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",
-    );
+pub fn fetch_tables_prefix(
+    bridgehead_tables: &HashMap<String, HashSet<String>>, bridgehead: &str, default: &str) -> String {
+    if let Some(set) = bridgehead_tables.get(bridgehead) {
+        let mut prefix: Option<String> = None;
 
-    builder_script
+        for value in set {
+            if let Some((current_prefix, _)) = value.split_once('.') {
+                match &prefix {
+                    Some(existing_prefix) if existing_prefix != current_prefix => {
+                        return default.to_string();
+                    }
+                    None => {
+                        prefix = Some(current_prefix.to_string());
+                    }
+                    _ => {}
+                }
+            }
+        }
+        prefix.unwrap_or_else(|| default.to_string())
+    } else {
+        default.to_string()
+    }
 }
+