From 28c8e9df69375ed368fcbfe296a81c5e33137e68 Mon Sep 17 00:00:00 2001 From: Firstyear Date: Tue, 27 Jun 2023 16:45:45 +1000 Subject: [PATCH] Improve cli to support multi-domain handling. (#1786) --- tools/cli/src/cli/common.rs | 57 ++++++++++++++++++++++++------------ tools/cli/src/cli/session.rs | 20 ++++++------- 2 files changed, 48 insertions(+), 29 deletions(-) diff --git a/tools/cli/src/cli/common.rs b/tools/cli/src/cli/common.rs index 4e035f204..4227b544f 100644 --- a/tools/cli/src/cli/common.rs +++ b/tools/cli/src/cli/common.rs @@ -101,19 +101,42 @@ impl CommonOpt { return Err(ToClientError::Other); } - // we need to store guessed username for login and reauth. - let username; - // If we have a username, use that to select tokens - let token = match &self.username { - Some(_username) => { - username = _username.clone(); + let (spn, token) = match &self.username { + Some(filter_username) => { + let possible_token = if filter_username.contains('@') { + // If there is an @, it's an spn so just get the token directly. + tokens + .get(filter_username) + .map(|t| (filter_username.clone(), t.clone())) + } else { + let filter_username = format!("{}@", filter_username); + // First, filter for tokens that match. + let mut token_refs: Vec<_> = tokens + .iter() + .filter(|(t, _)| t.starts_with(&filter_username)) + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + + match token_refs.len() { + 0 => None, + 1 => token_refs.pop(), + _ => { + error!("Multiple authentication tokens found for {}. Please specify the full spn to proceed", filter_username); + return Err(ToClientError::Other); + } + } + }; + // Is it in the store? - match tokens.get(&username) { - Some(t) => t.clone(), + match possible_token { + Some(t) => t, None => { - error!("No valid authentication tokens found for {}.", username); - return Err(ToClientError::NeedLogin(username)); + error!( + "No valid authentication tokens found for {}.", + filter_username + ); + return Err(ToClientError::NeedLogin(filter_username.clone())); } } } @@ -123,16 +146,12 @@ impl CommonOpt { let (f_uname, f_token) = tokens.iter().next().expect("Memory Corruption"); // else pick the first token debug!("Using cached token for name {}", f_uname); - username = f_uname.clone(); - f_token.clone() + (f_uname.clone(), f_token.clone()) } else { // Unable to automatically select the user because multiple tokens exist // so we'll prompt the user to select one match prompt_for_username_get_values() { - Ok((f_uname, f_token)) => { - username = f_uname; - f_token - } + Ok(tuple) => tuple, Err(msg) => { error!("{}", msg); std::process::exit(1); @@ -163,7 +182,7 @@ impl CommonOpt { "Session has expired for {} - you may need to login again.", uat.spn ); - return Err(ToClientError::NeedLogin(username)); + return Err(ToClientError::NeedLogin(spn)); } } @@ -176,7 +195,7 @@ impl CommonOpt { "Privileges have expired for {} - you need to re-authenticate again.", uat.spn ); - return Err(ToClientError::NeedReauth(username)); + return Err(ToClientError::NeedReauth(spn)); } } } @@ -184,7 +203,7 @@ impl CommonOpt { Err(e) => { error!("Unable to read token for requested user - you may need to login again."); debug!(?e, "JWT Error"); - return Err(ToClientError::NeedLogin(username)); + return Err(ToClientError::NeedLogin(spn)); } }; diff --git a/tools/cli/src/cli/session.rs b/tools/cli/src/cli/session.rs index db1f8a5e2..fc8acdd07 100644 --- a/tools/cli/src/cli/session.rs +++ b/tools/cli/src/cli/session.rs @@ -304,7 +304,7 @@ async fn process_auth_state( }); // Add our new one - let (username, tonk) = match client.get_token().await { + let (spn, tonk) = match client.get_token().await { Some(t) => { let tonk = match JwsUnverified::from_str(&t).and_then(|jwtu| { jwtu.validate_embeded() @@ -317,9 +317,9 @@ async fn process_auth_state( } }; - let username = tonk.name().to_string(); + let spn = tonk.spn; // Return the un-parsed token - (username, t) + (spn, t) } None => { error!("Error retrieving client session"); @@ -327,7 +327,7 @@ async fn process_auth_state( } }; - tokens.insert(username.clone(), tonk); + tokens.insert(spn.clone(), tonk); // write them out. if write_tokens(&tokens).is_err() { @@ -336,7 +336,7 @@ async fn process_auth_state( }; // Success! - println!("Login Success for {}", username); + println!("Login Success for {}", spn); } impl LoginOpt { @@ -428,7 +428,7 @@ impl LogoutOpt { } pub async fn exec(&self) { - let username: String = if self.local_only { + let spn: String = if self.local_only { // For now we just remove this from the token store. let mut _tmp_username = String::new(); match &self.copt.username { @@ -474,7 +474,7 @@ impl LogoutOpt { } }; - uat.name().to_string() + uat.spn }; let mut tokens = read_tokens().unwrap_or_else(|_| { @@ -483,15 +483,15 @@ impl LogoutOpt { }); // Remove our old one - if tokens.remove(&username).is_some() { + if tokens.remove(&spn).is_some() { // write them out. if let Err(_e) = write_tokens(&tokens) { error!("Error persisting authentication token store"); std::process::exit(1); }; - println!("Removed session for {}", username); + println!("Removed session for {}", spn); } else { - println!("No sessions for {}", username); + println!("No sessions for {}", spn); } } }