cipherstash-client 0.12.5

pub mod user_token;

mod auth0;
mod okta;

use std::{
    path::Path,
    time::{Duration, SystemTime, UNIX_EPOCH},
};

use async_mutex::Mutex as AsyncMutex;
use async_trait::async_trait;
use auth0::Auth0UserCredentials;
use log::debug;
use miette::Diagnostic;
use okta::OktaUserCredentials;
use serde::Deserialize;
use thiserror::Error;
use url::Url;

use crate::{
    config::idp_provider::IdpProvider,
    credentials::{
        token_store::TokenStore, AutoRefreshable, ClearTokenError, Credentials, GetTokenError,
        TokenExpiry,
    },
};

pub use user_token::UserToken;

// The offline_access scope is used for requesting a refresh token
// The cipherstash:admin scope is used by self hosted CTS to allow access to
// management endpoints
pub const DEFAULT_REQUESTED_SCOPES: &str = "offline_access cipherstash:admin";

#[derive(Deserialize)]
pub(crate) struct PollingInfo {
    pub user_code: String,
    pub device_code: String,
    pub verification_uri_complete: String,
}

#[derive(Deserialize)]
pub(crate) struct AccessTokenResponse {
    pub refresh_token: String,
    pub access_token: String,
    pub expires_in: u64,
}

impl From<AccessTokenResponse> for UserToken {
    fn from(value: AccessTokenResponse) -> Self {
        Self {
            access_token: value.access_token,
            refresh_token: value.refresh_token,
            expiry: value.expires_in + now_secs(),
        }
    }
}

pub fn now_secs() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .expect("Expected system time to be greater than UNIX_EPOCH")
        .as_secs()
}

/// Show a prompt in the terminal and open a browser (if available) with an authentication link for generating
/// an access token.
pub(crate) fn prompt_user(polling_info: &PollingInfo) {
    if open::that(&polling_info.verification_uri_complete).is_err() {
        println!(
            "Failed to open web browser. Please manually click the link in the following message."
        )
    }

    let user_code = &polling_info.user_code;
    let code_len = user_code.len();

    println!();
    println!("### ACTION REQUIRED ###");
    println!();
    println!(
        "Visit {} to complete authentication by following the below steps:",
        polling_info.verification_uri_complete
    );
    println!();
    println!("1. Verify that this code matches the code in your browser");
    println!();
    println!("             +------{}------+", "-".repeat(code_len));
    println!("             |      {}      |", " ".repeat(code_len));
    println!("             |      {user_code}      |");
    println!("             |      {}      |", " ".repeat(code_len));
    println!("             +------{}------+", "-".repeat(code_len));
    println!();
    println!("2. If the codes match, click on the confirm button in the browser");
    println!();
    println!("Waiting for authentication...");
}

#[derive(Diagnostic, Error, Debug)]
pub enum RefreshTokenError {
    #[error("Failed to redeem refresh token: {0}")]
    RequestFailed(reqwest::Error),

    #[error("Failed to parse json response: {0}")]
    BadResponse(reqwest::Error),
}

#[derive(Diagnostic, Error, Debug)]
pub enum NewTokenError {
    #[error("Failed to parse Url: {0}")]
    UrlParse(#[from] url::ParseError),
    #[error("Failed to get device code: {0}")]
    DeviceCodeRequestFailed(reqwest::Error),

    #[error("Failed to parse polling info json response: {0}")]
    DeviceCodeBadResponse(reqwest::Error),

    #[error("Failed to poll for new token: {0}")]
    PollTokenRequestFailed(reqwest::Error),

    #[error("Failed to parse access token response: {0}")]
    PollTokenBadResponse(reqwest::Error),

    #[error("Failed to parse pending auth response: {0}")]
    PollTokenBadPendingResponse(reqwest::Error),

    #[error("Device code authentication failed: {0}")]
    PollTokenAuthFailed(String),

    #[error("Unexpected error code in response body: {0}")]
    PollTokenUnexpected(String),
}

pub struct UserCredentials {
    token_store: AsyncMutex<TokenStore<UserToken>>,
    provider: UserCredentialsProvider,
}

enum UserCredentialsProvider {
    Auth0(Auth0UserCredentials),
    Okta(OktaUserCredentials),
}

impl UserCredentials {
    pub fn new(
        idp_token_path: &Path,
        idp_base_url: &Url,
        idp_audience: &str,
        idp_client_id: &str,
        idp_provider: IdpProvider,
    ) -> Self {
        let provider = match idp_provider {
            IdpProvider::Auth0 => UserCredentialsProvider::Auth0(Auth0UserCredentials::new(
                idp_base_url,
                idp_audience,
                idp_client_id,
            )),

            IdpProvider::Okta => {
                UserCredentialsProvider::Okta(OktaUserCredentials::new(idp_base_url, idp_client_id))
            }
        };

        Self {
            token_store: AsyncMutex::new(TokenStore::new(idp_token_path)),
            provider,
        }
    }
}

impl UserCredentialsProvider {
    async fn refresh_access_token(
        &self,
        cached_token: &UserToken,
    ) -> Result<Option<UserToken>, RefreshTokenError> {
        match self {
            Self::Auth0(creds) => creds.refresh_access_token(cached_token).await,
            Self::Okta(creds) => creds.refresh_access_token(cached_token).await,
        }
    }

    async fn acquire_new_token(&self) -> Result<UserToken, NewTokenError> {
        match self {
            Self::Auth0(creds) => creds.acquire_new_token().await,
            Self::Okta(creds) => creds.acquire_new_token().await,
        }
    }
}

#[async_trait]
impl Credentials for UserCredentials {
    type Token = UserToken;

    async fn get_token(&self) -> Result<Self::Token, GetTokenError> {
        let mut token_store = self.token_store.lock().await;

        // Check to see if we can make a new token from the cache
        if let Some(cached_token) = &token_store.get() {
            // If the token hasn't expired yet just return with it immediately
            if !cached_token.is_expired() {
                return Ok(cached_token.clone());
            }

            // The cached token has expired so try and get a new one from the refresh token.
            // If auto refresh is enabled, this should not have to happen
            if let Some(new_token) = self
                .provider
                .refresh_access_token(cached_token)
                .await
                .map_err(|e| GetTokenError::RefreshTokenFailed(Box::new(e)))?
            {
                token_store
                    .set(&new_token)
                    .map_err(|e| GetTokenError::PersistTokenError(Box::new(e)))?;
                return Ok(new_token);
            }
        }

        // We need to ask the user to allow us to get a new token via device code
        let new_token = self
            .provider
            .acquire_new_token()
            .await
            .map_err(|err| GetTokenError::AcquireNewTokenFailed(Box::new(err)))?;

        // Saves the token to disk
        token_store
            .set(&new_token)
            .map_err(|e| GetTokenError::PersistTokenError(Box::new(e)))?;
        Ok(new_token)
    }

    async fn clear_token(&self) -> Result<(), ClearTokenError> {
        let mut token_store = self.token_store.lock().await;
        token_store
            .clear()
            .map_err(|e| ClearTokenError(Box::new(e)))
    }
}

#[async_trait]
impl AutoRefreshable for UserCredentials {
    async fn refresh(&self) -> Duration {
        let token = {
            // Drop the guard early to allow other tasks get the current cached token if still valid
            let mut token_store = self.token_store.lock().await;
            token_store.get()
        };

        // Check to see if we have a token from the cache or disk
        if let Some(cached_token) = token {
            debug!(target: "console_credentials", "Found token on disk");

            // If the token is still new, we do an early return
            if !cached_token.should_refresh() {
                debug!(target: "console_credentials", "Access token is still new");
                return cached_token.refresh_interval();
            }

            // The cached token is close to expiry, so try and get a new one from the refresh token
            debug!(target: "console_credentials", "Access token close to expiry, refreshing");
            if let Ok(Some(new_token)) = self.provider.refresh_access_token(&cached_token).await {
                let mut token_store = self.token_store.lock().await;
                if token_store.set(&new_token).is_ok() {
                    debug!(target: "console_credentials", "Access token refreshed and saved to disk");
                    return new_token.refresh_interval();
                }
            }
        }

        Self::Token::min_refresh_interval()
    }
}