From f4b355c299a5072fee1c2ec79d4411e476cbf21e Mon Sep 17 00:00:00 2001 From: Yuxuan Lu Date: Mon, 5 Jun 2023 20:03:52 +0800 Subject: [PATCH] Automatically login & reauth (#1691) * Automatically login & reauth * Revert automatical format by vscode * Only ask user to choose username once * Use dialoguer::Confirm; fix logic --- Cargo.lock | 12 +++++ Cargo.toml | 1 + tools/cli/Cargo.toml | 1 + tools/cli/src/cli/common.rs | 101 ++++++++++++++++++++++++++++++------ tools/cli/src/opt/kanidm.rs | 17 +++--- 5 files changed, 108 insertions(+), 24 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8f92785d9..3bc195037 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -316,6 +316,17 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "async-recursion" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e97ce7de6cf12de5d7226c73f5ba9811622f4db3a5b91b55c53e987e5f91cba" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.17", +] + [[package]] name = "async-session" version = "2.0.1" @@ -2388,6 +2399,7 @@ dependencies = [ name = "kanidm_tools" version = "1.1.0-beta.13-dev" dependencies = [ + "async-recursion", "clap", "clap_complete", "compact_jwt", diff --git a/Cargo.toml b/Cargo.toml index 6c7228b86..ebf8b0e7b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -38,6 +38,7 @@ homepage = "https://github.com/kanidm/kanidm/" repository = "https://github.com/kanidm/kanidm/" [workspace.dependencies] +async-recursion = "1.0.4" async-trait = "^0.1.68" base32 = "^0.4.0" base64 = "^0.21.0" diff --git a/tools/cli/Cargo.toml b/tools/cli/Cargo.toml index 140dbe98e..51b6fd9d0 100644 --- a/tools/cli/Cargo.toml +++ b/tools/cli/Cargo.toml @@ -30,6 +30,7 @@ name = "kanidm_ssh_authorizedkeys_direct" path = "src/ssh_authorizedkeys.rs" [dependencies] +async-recursion.workplace = true clap = { workspace = true, features = ["derive", "env"] } compact_jwt = { workspace = true, features = ["openssl"] } dialoguer.workspace = true diff --git a/tools/cli/src/cli/common.rs b/tools/cli/src/cli/common.rs index 1693993b5..a81249c94 100644 --- a/tools/cli/src/cli/common.rs +++ b/tools/cli/src/cli/common.rs @@ -1,20 +1,30 @@ +use std::env; use std::str::FromStr; +use async_recursion::async_recursion; use compact_jwt::{Jws, JwsUnverified}; use dialoguer::theme::ColorfulTheme; -use dialoguer::Select; +use dialoguer::{Confirm, Select}; use kanidm_client::{KanidmClient, KanidmClientBuilder}; use kanidm_proto::constants::{DEFAULT_CLIENT_CONFIG_PATH, DEFAULT_CLIENT_CONFIG_PATH_HOME}; use kanidm_proto::v1::UserAuthToken; use crate::session::read_tokens; -use crate::CommonOpt; +use crate::{CommonOpt, LoginOpt, ReauthOpt}; +#[derive(Clone)] pub enum OpType { Read, Write, } +#[derive(Debug)] +pub enum ToClientError { + NeedLogin(String), + NeedReauth(String), + Other, +} + impl CommonOpt { pub fn to_unauth_client(&self) -> KanidmClient { let config_path: String = shellexpand::tilde(DEFAULT_CLIENT_CONFIG_PATH_HOME).into_owned(); @@ -73,14 +83,14 @@ impl CommonOpt { }) } - pub async fn to_client(&self, optype: OpType) -> KanidmClient { + async fn try_to_client(&self, optype: OpType) -> Result { let client = self.to_unauth_client(); // Read the token file. let tokens = match read_tokens() { Ok(t) => t, Err(_e) => { error!("Error retrieving authentication token store"); - std::process::exit(1); + return Err(ToClientError::Other); } }; @@ -88,18 +98,22 @@ impl CommonOpt { error!( "No valid authentication tokens found. Please login with the 'login' subcommand." ); - std::process::exit(1); + 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) => { + Some(_username) => { + username = _username.clone(); // Is it in the store? - match tokens.get(username) { + match tokens.get(&username) { Some(t) => t.clone(), None => { error!("No valid authentication tokens found for {}.", username); - std::process::exit(1); + return Err(ToClientError::NeedLogin(username)); } } } @@ -109,12 +123,16 @@ 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() } 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_token() { - Ok(value) => value, + match prompt_for_username_get_values() { + Ok((f_uname, f_token)) => { + username = f_uname; + f_token + } Err(msg) => { error!("{}", msg); std::process::exit(1); @@ -128,7 +146,7 @@ impl CommonOpt { Ok(jwtu) => jwtu, Err(e) => { error!("Unable to parse token - {:?}", e); - std::process::exit(1); + return Err(ToClientError::Other); } }; @@ -145,7 +163,7 @@ impl CommonOpt { "Session has expired for {} - you may need to login again.", uat.spn ); - std::process::exit(1); + return Err(ToClientError::NeedLogin(username)); } } @@ -158,7 +176,7 @@ impl CommonOpt { "Privileges have expired for {} - you need to re-authenticate again.", uat.spn ); - std::process::exit(1); + return Err(ToClientError::NeedReauth(username)); } } } @@ -166,14 +184,67 @@ impl CommonOpt { Err(e) => { error!("Unable to read token for requested user - you may need to login again."); debug!(?e, "JWT Error"); - std::process::exit(1); + return Err(ToClientError::NeedLogin(username)); } }; // Set it into the client client.set_token(token).await; - client + Ok(client) + } + + #[async_recursion] + pub async fn to_client(&self, optype: OpType) -> KanidmClient { + match self.try_to_client(optype.clone()).await { + Ok(c) => c, + Err(e) => { + match e { + ToClientError::NeedLogin(username) => { + if !Confirm::new() + .with_prompt("Would you like to login again?") + .interact() + .expect("Failed to interact with interactive session") + { + std::process::exit(1); + } + let mut copt = self.clone(); + copt.username = Some(username); + let login_opt = LoginOpt { + copt, + password: env::var("KANIDM_PASSWORD").ok(), + }; + login_opt.exec().await; + // we still use `to_client` instead of `try_to_client` because we may need to prompt user to re-auth again. + // since reauth_opt will call `to_client`, this function is recursive anyway. + // we use copt since it's username is updated. + return login_opt.copt.to_client(optype).await; + } + ToClientError::NeedReauth(username) => { + if !Confirm::new() + .with_prompt("Would you like to re-authenticate?") + .interact() + .expect("Failed to interact with interactive session") + { + std::process::exit(1); + } + let mut copt = self.clone(); + copt.username = Some(username); + let reauth_opt = ReauthOpt { copt }; + // calls `to_client` recursively + // but should not goes into `NeedLogin` branch again + reauth_opt.exec().await; + if let Ok(c) = reauth_opt.copt.try_to_client(optype).await { + return c; + } + } + ToClientError::Other => { + std::process::exit(1); + } + } + std::process::exit(1); + } + } } } diff --git a/tools/cli/src/opt/kanidm.rs b/tools/cli/src/opt/kanidm.rs index 0b32c9fc1..f6c212198 100644 --- a/tools/cli/src/opt/kanidm.rs +++ b/tools/cli/src/opt/kanidm.rs @@ -14,7 +14,7 @@ pub struct DebugOpt { pub debug: bool, } -#[derive(Debug, Args)] +#[derive(Debug, Args, Clone)] pub struct CommonOpt { /// Enable debbuging of the kanidm tool #[clap(short, long, env = "KANIDM_DEBUG")] @@ -29,14 +29,14 @@ pub struct CommonOpt { #[clap(parse(from_os_str), short = 'C', long = "ca", env = "KANIDM_CA_PATH")] pub ca_path: Option, /// Log format (still in very early development) - #[clap(short, long = "output", env = "KANIDM_OUTPUT", default_value="text")] + #[clap(short, long = "output", env = "KANIDM_OUTPUT", default_value = "text")] output_mode: String, } #[derive(Debug, Args)] pub struct GroupNamedMembers { name: String, - #[clap(required=true,min_values=1)] + #[clap(required = true, min_values = 1)] members: Vec, #[clap(flatten)] copt: CommonOpt, @@ -503,7 +503,7 @@ pub enum RecycleOpt { pub struct LoginOpt { #[clap(flatten)] copt: CommonOpt, - #[clap(short, long, env="KANIDM_PASSWORD", hide=true)] + #[clap(short, long, env = "KANIDM_PASSWORD", hide = true)] /// Supply a password to the login option password: Option, } @@ -711,7 +711,6 @@ pub struct OptSetDomainDisplayName { new_display_name: String, } - #[derive(Debug, Subcommand)] pub enum PwBadlistOpt { #[clap[name = "show"]] @@ -724,7 +723,7 @@ pub enum PwBadlistOpt { Upload { #[clap(flatten)] copt: CommonOpt, - #[clap(parse(from_os_str),required=true,min_values=1)] + #[clap(parse(from_os_str), required = true, min_values = 1)] paths: Vec, /// Perform a dry run and display the list that would have been uploaded instead. #[clap(short = 'n', long)] @@ -736,9 +735,9 @@ pub enum PwBadlistOpt { Remove { #[clap(flatten)] copt: CommonOpt, - #[clap(parse(from_os_str), required=true, min_values=1)] + #[clap(parse(from_os_str), required = true, min_values = 1)] paths: Vec, - } + }, } #[derive(Debug, Subcommand)] @@ -863,7 +862,7 @@ pub enum SystemOpt { Synch { #[clap(subcommand)] commands: SynchOpt, - } + }, } #[derive(Debug, Subcommand)]