Improve cli to support multi-domain handling. (#1786)

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 {
+        let (spn, token) = match &self.username {
-            Some(_username) => {
+            Some(filter_username) => {
-                username = _username.clone();
+                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) {
+                match possible_token {
-                    Some(t) => t.clone(),
+                    Some(t) => t,
                     None => {
-                        error!("No valid authentication tokens found for {}.", username);
+                        error!(
-                        return Err(ToClientError::NeedLogin(username));
+                            "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_uname.clone(), f_token.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)) => {
+                        Ok(tuple) => tuple,
-                            username = f_uname;
-                            f_token
-                        }
                         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));
             }
         };
 

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);
         }
     }
 }