PXP-638: [CLI] Implement new config

cli/Cargo.toml

@@ -62,6 +62,8 @@ semver = "1.0.20"
 # Telemetry
 wukong-telemetry-macro = { path = "../telemetry-macro" }
 wukong-telemetry = { path = "../telemetry" }
+time = { version = "0.3.31", features = ["formatting", "parsing"] }
+anyhow = "1.0.79"
 
 [dev-dependencies]
 httpmock = "0.6.7"

cli/src/auth/google_cloud.rs

@@ -1,10 +1,12 @@
-use once_cell::sync::Lazy;
+use crate::config::Config;
 use serde::{Deserialize, Serialize};
 use std::{future::Future, pin::Pin};
+use time::{format_description, OffsetDateTime};
+use tonic::async_trait;
 use yup_oauth2::{
     authenticator_delegate::{DefaultInstalledFlowDelegate, InstalledFlowDelegate},
     hyper, hyper_rustls,
-    storage::TokenInfo,
+    storage::{TokenInfo, TokenStorage},
     ApplicationSecret, InstalledFlowAuthenticator, InstalledFlowReturnMethod,
 };
 
@@ -14,25 +16,6 @@ struct JSONToken {
     token: TokenInfo,
 }
 
-pub static CONFIG_PATH: Lazy<Option<String>> = Lazy::new(|| {
-    #[cfg(feature = "prod")]
-    return dirs::home_dir().map(|mut path| {
-        path.extend([".config", "wukong"]);
-        path.to_str().unwrap().to_string()
-    });
-
-    #[cfg(not(feature = "prod"))]
-    {
-        match std::env::var("WUKONG_DEV_GCLOUD_FILE") {
-            Ok(config) => Some(config),
-            Err(_) => dirs::home_dir().map(|mut path| {
-                path.extend([".config", "wukong", "dev"]);
-                path.to_str().unwrap().to_string()
-            }),
-        }
-    }
-});
-
 /// async function to be pinned by the `present_user_url` method of the trait
 /// we use the existing `DefaultInstalledFlowDelegate::present_user_url` method as a fallback for
 /// when the browser did not open for example, the user still see's the URL.
@@ -74,7 +57,64 @@ const AUTH_URI: &str = "https://accounts.google.com/o/oauth2/auth";
 const REDIRECT_URI: &str = "http://127.0.0.1/8855";
 const AUTH_PROVIDER_X509_CERT_URL: &str = "https://www.googleapis.com/oauth2/v1/certs";
 
-pub async fn get_token_or_login() -> String {
+#[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Clone)]
+pub struct GoogleCloudConfig {
+    /// used when authorizing calls to oauth2 enabled services.
+    pub access_token: String,
+    /// used to refresh an expired access_token.
+    pub refresh_token: String,
+    /// The time when the token expires.
+    pub expiry_time: String,
+    /// Optionally included by the OAuth2 server and may contain information to verify the identity
+    /// used to obtain the access token.
+    /// Specifically Google API:s include this if the additional scopes "email" and/or "profile"
+    /// are used. In that case the content is an JWT token.
+    pub id_token: Option<String>,
+}
+
+struct ConfigTokenStore {
+    config: Config,
+}
+
+#[async_trait]
+impl TokenStorage for ConfigTokenStore {
+    async fn set(&self, _scopes: &[&str], token: TokenInfo) -> anyhow::Result<()> {
+        let mut config = self.config.clone();
+
+        config.auth.google_cloud = Some(GoogleCloudConfig {
+            access_token: token.access_token.expect("Invalid access token"),
+            refresh_token: token.refresh_token.expect("Invalid refresh token"),
+            expiry_time: token
+                .expires_at
+                .expect("Invalid expiry time")
+                .format(&format_description::well_known::Rfc3339)?,
+            id_token: token.id_token,
+        });
+
+        config.save_to_default_path()?;
+
+        Ok(())
+    }
+
+    async fn get(&self, _target_scopes: &[&str]) -> Option<TokenInfo> {
+        let google_cloud = self.config.auth.google_cloud.clone()?;
+
+        Some(TokenInfo {
+            access_token: Some(google_cloud.access_token),
+            refresh_token: Some(google_cloud.refresh_token),
+            expires_at: Some(
+                OffsetDateTime::parse(
+                    &google_cloud.expiry_time,
+                    &format_description::well_known::Rfc3339,
+                )
+                .expect("Invalid expiry time"),
+            ),
+            id_token: google_cloud.id_token,
+        })
+    }
+}
+
+pub async fn get_token_or_login(config: Option<Config>) -> String {
     let secret = ApplicationSecret {
         client_id: GOOGLE_CLIENT_ID.to_string(),
         client_secret: GOOGLE_CLIENT_SECRET.to_string(),
@@ -96,17 +136,17 @@ pub async fn get_token_or_login() -> String {
             .build(),
     );
 
+    let config = match config {
+        Some(config) => config,
+        None => Config::load_from_default_path().expect("Unable to load config"),
+    };
+
     let authenticator = InstalledFlowAuthenticator::with_client(
         secret,
         InstalledFlowReturnMethod::HTTPPortRedirect(8855),
         client,
     )
-    .persist_tokens_to_disk(format!(
-        "{}/gcloud_logging",
-        CONFIG_PATH
-            .as_ref()
-            .expect("Unable to identify user's home directory"),
-    ))
+    .with_storage(Box::new(ConfigTokenStore { config }))
     .flow_delegate(Box::new(InstalledFlowBrowserDelegate))
     .build()
     .await
@@ -122,39 +162,15 @@ pub async fn get_token_or_login() -> String {
 }
 
 pub async fn get_access_token() -> Option<String> {
-    let contents = tokio::fs::read(format!(
-        "{}/gcloud_logging",
-        CONFIG_PATH
-            .as_ref()
-            .expect("Unable to identify user's home directory")
-    ))
-    .await;
-
-    let tokens = contents
-        .map(|contents| {
-            serde_json::from_slice::<Vec<JSONToken>>(&contents)
-                .map_err(|_| {
-                    eprintln!("Failed to parse token file.");
-                })
-                .ok()
-        })
-        .unwrap_or(None);
-
-    let json_token = tokens.and_then(|tokens| {
-        tokens
-            .iter()
-            .find(|token| {
-                token
-                    .scopes
-                    .contains(&"https://www.googleapis.com/auth/logging.read".to_string())
-            })
-            .map(|token| token.token.access_token.clone())
-    });
+    let config = match Config::load_from_default_path() {
+        Ok(config) => config,
+        Err(_) => return None,
+    };
 
     // Sometimes access token exist but is expired, so call get_token_or_login() to refresh it
     // before returning it.
-    if json_token.is_some() {
-        Some(get_token_or_login().await)
+    if config.auth.google_cloud.is_some() {
+        Some(get_token_or_login(None).await)
     } else {
         None
     }

cli/src/auth/okta.rs

@@ -1,5 +1,5 @@
 use crate::{
-    config::{AuthConfig, Config},
+    config::{Config, OktaConfig},
     error::{AuthError, WKCliError},
     utils::compare_with_current_time,
 };
@@ -53,7 +53,7 @@ pub struct OktaAuth {
     pub expiry_time: String,
 }
 
-impl From<OktaAuth> for AuthConfig {
+impl From<OktaAuth> for OktaConfig {
     fn from(value: OktaAuth) -> Self {
         Self {
             account: value.account,
@@ -67,9 +67,13 @@ impl From<OktaAuth> for AuthConfig {
 }
 
 pub fn need_tokens_refresh(config: &Config) -> Result<bool, WKCliError> {
-    let auth_config = config.auth.as_ref().ok_or(WKCliError::UnAuthenticated)?;
+    let okta_config = config
+        .auth
+        .okta
+        .as_ref()
+        .ok_or(WKCliError::UnAuthenticated)?;
 
-    let remaining_duration = compare_with_current_time(&auth_config.expiry_time);
+    let remaining_duration = compare_with_current_time(&okta_config.expiry_time);
     Ok(remaining_duration < EXPIRY_REMAINING_TIME_IN_MINS.minutes())
 }
 
@@ -245,7 +249,11 @@ pub async fn login(config: &Config) -> Result<OktaAuth, WKCliError> {
 }
 
 pub async fn refresh_tokens(config: &Config) -> Result<OktaAuth, WKCliError> {
-    let auth_config = config.auth.as_ref().ok_or(WKCliError::UnAuthenticated)?;
+    let okta_config = config
+        .auth
+        .okta
+        .as_ref()
+        .ok_or(WKCliError::UnAuthenticated)?;
     let okta_client_id = ClientId::new(config.core.okta_client_id.clone());
 
     let issuer_url =
@@ -260,7 +268,7 @@ pub async fn refresh_tokens(config: &Config) -> Result<OktaAuth, WKCliError> {
     let client = CoreClient::from_provider_metadata(provider_metadata, okta_client_id, None);
 
     let token_exchange_result = client
-        .exchange_refresh_token(&RefreshToken::new(auth_config.refresh_token.clone()))
+        .exchange_refresh_token(&RefreshToken::new(okta_config.refresh_token.clone()))
         .request_async(async_http_client)
         .await;
 
@@ -276,7 +284,7 @@ pub async fn refresh_tokens(config: &Config) -> Result<OktaAuth, WKCliError> {
                         && error_description.contains("refresh token")
                         && error_description.contains("expired")
                     {
-                        introspect_token(config, &auth_config.refresh_token).await?;
+                        introspect_token(config, &okta_config.refresh_token).await?;
 
                         AuthError::OktaRefreshTokenExpired {
                             message: error_description,
@@ -318,8 +326,8 @@ pub async fn refresh_tokens(config: &Config) -> Result<OktaAuth, WKCliError> {
         .to_rfc3339();
 
     Ok(OktaAuth {
-        account: auth_config.account.clone(),
-        subject: auth_config.subject.clone(),
+        account: okta_config.account.clone(),
+        subject: okta_config.subject.clone(),
         id_token: id_token.to_string(),
         access_token: access_token.secret().to_owned(),
         expiry_time: expiry,

cli/src/auth/vault.rs

@@ -48,10 +48,14 @@ pub static BASE_URL: Lazy<String> = Lazy::new(|| {
 });
 
 pub async fn get_token_or_login(config: &mut Config) -> Result<String, WKCliError> {
-    let auth_config = config.auth.as_ref().ok_or(WKCliError::UnAuthenticated)?;
+    let okta_config = config
+        .auth
+        .okta
+        .as_ref()
+        .ok_or(WKCliError::UnAuthenticated)?;
     let client = reqwest::Client::new();
 
-    match &config.vault {
+    match &config.auth.vault {
         Some(vault_config) => match verify_token(&client, &BASE_URL, &vault_config.api_token).await
         {
             Ok(_) => {
@@ -71,7 +75,7 @@ pub async fn get_token_or_login(config: &mut Config) -> Result<String, WKCliErro
                     let new_expiry_time =
                         calculate_expiry_time(renew_token_resp.auth.lease_duration);
 
-                    config.vault = Some(VaultConfig {
+                    config.auth.vault = Some(VaultConfig {
                         api_token: token.clone(),
                         expiry_time: new_expiry_time,
                     });
@@ -84,7 +88,7 @@ pub async fn get_token_or_login(config: &mut Config) -> Result<String, WKCliErro
             }
             Err(WKCliError::AuthError(AuthError::VaultPermissionDenied)) => {
                 // login
-                colored_println!("Login Vault with okta account {}", auth_config.account);
+                colored_println!("Login Vault with okta account {}", okta_config.account);
 
                 let password = dialoguer::Password::with_theme(&ColorfulTheme::default())
                     .with_prompt("Enter okta password")
@@ -93,25 +97,25 @@ pub async fn get_token_or_login(config: &mut Config) -> Result<String, WKCliErro
                 let loader = new_spinner();
                 loader.set_message("Authenticating the user... You may need to check your device for an MFA notification.");
 
-                let email = &auth_config.account;
+                let email = &okta_config.account;
                 let login_resp = login(&client, &BASE_URL, email, &password).await?;
 
                 loader.finish_and_clear();
                 let expiry_time = calculate_expiry_time(login_resp.auth.lease_duration);
 
-                config.vault = Some(VaultConfig {
+                config.auth.vault = Some(VaultConfig {
                     api_token: login_resp.auth.client_token.clone(),
                     expiry_time,
                 });
                 config.save_to_default_path()?;
 
-                colored_println!("You are now logged in as {}.\n", email);
+                colored_println!("You are now logged in as: {}.\n", email);
                 Ok(login_resp.auth.client_token)
             }
             Err(err) => Err(err),
         },
         None => {
-            colored_println!("Login Vault with okta account {}", auth_config.account);
+            colored_println!("Login Vault with okta account {}", okta_config.account);
 
             let password = dialoguer::Password::with_theme(&ColorfulTheme::default())
                 .with_prompt("Enter okta password")
@@ -122,19 +126,19 @@ pub async fn get_token_or_login(config: &mut Config) -> Result<String, WKCliErro
             "Authenticating the user... You may need to check your device for an MFA notification.",
         );
 
-            let email = &auth_config.account;
+            let email = &okta_config.account;
             let login_resp = login(&client, &BASE_URL, email, &password).await?;
 
             loader.finish_and_clear();
             let expiry_time = calculate_expiry_time(login_resp.auth.lease_duration);
 
-            config.vault = Some(VaultConfig {
+            config.auth.vault = Some(VaultConfig {
                 api_token: login_resp.auth.client_token.clone(),
                 expiry_time,
             });
             config.save_to_default_path()?;
 
-            colored_println!("You are now logged in as {}.\n", email);
+            colored_println!("You are now logged in as: {}.\n", email);
             Ok(login_resp.auth.client_token)
         }
     }