jwt-verify 0.1.2

use jwt_verify::{
    CognitoAccessTokenClaims, CognitoIdTokenClaims, CognitoJwtVerifier, JwtError, JwtVerifier,
    VerifierConfig,
};
use std::{env, time::Duration};

/// Load environment variables from .env files
fn load_env() {
    dotenv::from_path("examples/.env").ok();
    dotenv::dotenv().ok();
}

/// Get environment variable with a default fallback
fn get_env_or_default(key: &str, default: &str) -> String {
    env::var(key).unwrap_or_else(|_| default.to_string())
}

/// Example 1: Basic single user pool verification
async fn example_single_pool() -> Result<(), JwtError> {
    println!("Example 1: Basic Cognito JWT verification with a single user pool");

    let region = get_env_or_default("AWS_REGION", "us-east-1");
    let user_pool_id = get_env_or_default("COGNITO_USER_POOL_ID", "us-east-1_example");
    let client_id = get_env_or_default("COGNITO_CLIENT_ID", "client1");

    println!(
        "Using Cognito User Pool: {} in region {} - client id {}",
        user_pool_id, region, client_id
    );

    let verifier = CognitoJwtVerifier::new_single_pool(&region, &user_pool_id, &[client_id])?;

    let id_token = env::var("COGNITO_ID_TOKEN").unwrap_or_else(|_| {
        println!("Warning: COGNITO_ID_TOKEN not set, using placeholder");
        "your_jwt_token_here".to_string()
    });

    let access_token = env::var("COGNITO_ACCESS_TOKEN").unwrap_or_else(|_| {
        println!("Warning: COGNITO_ACCESS_TOKEN not set, using placeholder");
        "your_access_token_here".to_string()
    });

    verify_id_token(&verifier, &id_token).await;
    verify_access_token(&verifier, &access_token).await;

    Ok(())
}

/// Example 2: Multiple user pools verification
async fn example_multiple_pools() -> Result<(), JwtError> {
    println!("\nExample 2: Cognito JWT verification with multiple user pools");

    let region = get_env_or_default("AWS_REGION", "us-east-1");
    let user_pool_id = get_env_or_default("COGNITO_USER_POOL_ID", "us-east-1_example");
    let client_id = get_env_or_default("COGNITO_CLIENT_ID", "client1");
    let client_id3 = get_env_or_default("COGNITO_CLIENT_ID_3", "client3");

    let region2 = get_env_or_default("AWS_REGION_2", "us-west-2");
    let user_pool_id2 = get_env_or_default("COGNITO_USER_POOL_ID_2", "us-west-2_example2");
    let client_id2 = get_env_or_default("COGNITO_CLIENT_ID_2", "client2");

    println!("Using multiple Cognito User Pools:");
    println!("  1. {} in region {}", user_pool_id, region);
    println!("  2. {} in region {}", user_pool_id2, region2);

    let config1 = VerifierConfig::new(&region, &user_pool_id, &[client_id, client_id3], None)?;
    let config2 = VerifierConfig::new(&region2, &user_pool_id2, &[client_id2], None)?;

    let multi_pool_verifier = CognitoJwtVerifier::new(vec![config1, config2])?;

    let access_token = env::var("COGNITO_ACCESS_TOKEN").unwrap_or_else(|_| {
        println!("Warning: COGNITO_ACCESS_TOKEN not set, using placeholder");
        "your_access_token_here".to_string()
    });

    verify_access_token(&multi_pool_verifier, &access_token).await;

    Ok(())
}

/// Example 3: Single pool with multiple client IDs
async fn example_multiple_clients() -> Result<(), JwtError> {
    println!("\nExample 3: Single user pool with multiple client IDs");

    let region = get_env_or_default("AWS_REGION", "us-east-1");
    let user_pool_id = get_env_or_default("COGNITO_USER_POOL_ID", "us-east-1_example");
    let client_id = get_env_or_default("COGNITO_CLIENT_ID", "client1");
    let client_id2 = get_env_or_default("COGNITO_CLIENT_ID_2", "client2");

    println!("Using single user pool with multiple client IDs:");
    println!("  User Pool: {} in region {}", user_pool_id, region);
    println!("  Client ID 1: {}", client_id);
    println!("  Client ID 2: {}", client_id2);

    let multi_client_config = VerifierConfig::new(
        &region,
        &user_pool_id,
        &[client_id.clone(), client_id2.clone()],
        None,
    )?
    .with_clock_skew(Duration::from_secs(120))
    .with_cache_duration(Duration::from_secs(3600 * 12));

    println!("Configuration:");
    println!("  - Clock skew: 120 seconds");
    println!("  - Cache duration: 12 hours");
    println!(
        "  - Allowed client IDs: {:?}",
        multi_client_config.client_ids
    );

    let multi_client_verifier = CognitoJwtVerifier::new(vec![multi_client_config])?;

    println!("\nTesting tokens with multi-client verifier:");

    let id_token = env::var("COGNITO_ID_TOKEN").unwrap_or_else(|_| {
        println!("Warning: COGNITO_ID_TOKEN not set, using placeholder");
        "your_jwt_token_here".to_string()
    });

    let access_token = env::var("COGNITO_ACCESS_TOKEN").unwrap_or_else(|_| {
        println!("Warning: COGNITO_ACCESS_TOKEN not set, using placeholder");
        "your_access_token_here".to_string()
    });

    verify_id_token(&multi_client_verifier, &id_token).await;
    verify_access_token(&multi_client_verifier, &access_token).await;

    prefetch_jwks(&multi_client_verifier).await;

    Ok(())
}

/// Verify an ID token and print claims
async fn verify_id_token(verifier: &CognitoJwtVerifier, token: &str) {
    match verifier.verify_id_token(token).await {
        Ok(claims) => {
            // Downcast immediately to access all Cognito-specific fields
            if let Some(cognito_claims) = claims.downcast_ref::<CognitoIdTokenClaims>() {
                println!("✅ ID Token verified successfully!");
                println!("  Subject: {}", cognito_claims.base.sub);
                println!("  Issuer: {}", cognito_claims.base.iss);
                println!("  Audience: {}", cognito_claims.aud);
                println!("  Token use: {}", cognito_claims.base.token_use);

                if let Some(email) = &cognito_claims.email {
                    println!("  Email: {}", email);
                    println!(
                        "  Email verified: {}",
                        cognito_claims.email_verified.unwrap_or(false)
                    );
                }

                if let Some(name) = &cognito_claims.name {
                    println!("  Name: {}", name);
                }

                if let Some(username) = &cognito_claims.cognito_username {
                    println!("  Cognito username: {}", username);
                }

                if let Some(groups) = &cognito_claims.cognito_groups {
                    println!("  Cognito groups: {:?}", groups);
                }
            } else {
                // Fallback if downcast fails (shouldn't happen with CognitoJwtVerifier)
                println!("✅ ID Token verified successfully!");
                println!("  Subject: {}", claims.get_sub());
                println!("  Issuer: {}", claims.get_iss());
                println!("  Audience: {}", claims.get_aud());
            }
        }
        Err(e) => {
            println!("❌ ID Token verification failed: {}", e);
        }
    }
}

/// Verify an access token and print claims
async fn verify_access_token(verifier: &CognitoJwtVerifier, token: &str) {
    match verifier.verify_access_token(token).await {
        Ok(claims) => {
            // Downcast immediately to access all Cognito-specific fields
            if let Some(cognito_claims) = claims.downcast_ref::<CognitoAccessTokenClaims>() {
                println!("✅ Access Token verified successfully!");
                println!("  Subject: {}", cognito_claims.base.sub);
                println!("  Issuer: {}", cognito_claims.base.iss);
                println!("  Token use: {}", cognito_claims.base.token_use);
                println!("  Scopes: {:?}", cognito_claims.get_scopes());

                let has_read_scope = cognito_claims.has_scope("read");
                println!("  Has 'read' scope: {}", has_read_scope);

                println!("  Client ID: {}", cognito_claims.base.client_id);

                if let Some(username) = &cognito_claims.base.username {
                    println!("  Username: {}", username);
                }

                if let Some(version) = cognito_claims.version {
                    println!("  Token version: {}", version);
                }
            } else {
                // Fallback if downcast fails (shouldn't happen with CognitoJwtVerifier)
                println!("✅ Access Token verified successfully!");
                println!("  Subject: {}", claims.get_sub());
                println!("  Issuer: {}", claims.get_iss());
                println!("  Scopes: {:?}", claims.get_scopes());
            }
        }
        Err(e) => {
            println!("❌ Access Token verification failed: {}", e);
        }
    }
}

/// Example 4: Negative test cases - demonstrating common errors
async fn example_negative_cases() -> Result<(), JwtError> {
    println!("\nExample 4: Negative test cases - error handling");

    let region = get_env_or_default("AWS_REGION", "us-east-1");
    let user_pool_id = get_env_or_default("COGNITO_USER_POOL_ID", "us-east-1_example");
    let client_id = get_env_or_default("COGNITO_CLIENT_ID", "client1");

    // Create verifier with only client1 allowed
    let verifier = CognitoJwtVerifier::new_single_pool(&region, &user_pool_id, &[client_id])?;

    let id_token = env::var("COGNITO_ID_TOKEN").unwrap_or_else(|_| {
        println!("Warning: COGNITO_ID_TOKEN not set, using placeholder");
        "your_jwt_token_here".to_string()
    });

    let access_token = env::var("COGNITO_ACCESS_TOKEN").unwrap_or_else(|_| {
        println!("Warning: COGNITO_ACCESS_TOKEN not set, using placeholder");
        "your_access_token_here".to_string()
    });

    // Test 1: Using access token with verify_id_token (wrong token type)
    println!("\n[Test 1] Verifying access token as ID token (should fail):");
    match verifier.verify_id_token(&access_token).await {
        Ok(_) => println!("  ❌ Unexpected success - should have failed!"),
        Err(e) => println!("  ✅ Expected error: {}", e),
    }

    // Test 2: Using ID token with verify_access_token (wrong token type)
    println!("\n[Test 2] Verifying ID token as access token (should fail):");
    match verifier.verify_access_token(&id_token).await {
        Ok(_) => println!("  ❌ Unexpected success - should have failed!"),
        Err(e) => println!("  ✅ Expected error: {}", e),
    }

    // Test 3: Token from disallowed client ID (same user pool)
    let disallowed_client_id = get_env_or_default("COGNITO_CLIENT_ID_2", "client2");
    println!(
        "\n[Test 3] Token from disallowed client ID (allowed: {}, token from: {}):",
        get_env_or_default("COGNITO_CLIENT_ID", "client1"),
        disallowed_client_id
    );

    if let Ok(other_client_token) = env::var("COGNITO_ID_TOKEN_CLIENT2") {
        match verifier.verify_id_token(&other_client_token).await {
            Ok(_) => println!("  ❌ Unexpected success - should have rejected disallowed client!"),
            Err(e) => println!("  ✅ Expected error: {}", e),
        }
    } else {
        println!("  ⊘ Skipped - COGNITO_ID_TOKEN_CLIENT2 not set");
    }

    // Test 4: Token from different user pool (CLIENT_ID_3 is in a different pool)
    println!("\n[Test 4] Token from different user pool (should fail):");

    if let Ok(different_pool_id_token) = env::var("COGNITO_ID_TOKEN_CLIENT3") {
        match verifier.verify_id_token(&different_pool_id_token).await {
            Ok(_) => {
                println!(
                    "  ❌ Unexpected success - should have rejected token from different pool!"
                )
            }
            Err(e) => println!("  ✅ Expected error (ID token): {}", e),
        }
    } else {
        println!("  ⊘ Skipped (ID token) - COGNITO_ID_TOKEN_CLIENT3 not set");
    }

    if let Ok(different_pool_access_token) = env::var("COGNITO_ACESS_TOKEN_CLIENT3") {
        match verifier
            .verify_access_token(&different_pool_access_token)
            .await
        {
            Ok(_) => {
                println!(
                    "  ❌ Unexpected success - should have rejected token from different pool!"
                )
            }
            Err(e) => println!("  ✅ Expected error (access token): {}", e),
        }
    } else {
        println!("  ⊘ Skipped (access token) - COGNITO_ACESS_TOKEN_CLIENT3 not set");
    }

    // Test 5: Malformed token
    println!("\n[Test 5] Malformed token (should fail):");
    let malformed_token = "not.a.valid.jwt.token";
    match verifier.verify_id_token(malformed_token).await {
        Ok(_) => println!("  ❌ Unexpected success - should have failed!"),
        Err(e) => println!("  ✅ Expected error: {}", e),
    }

    // Test 6: Empty token
    println!("\n[Test 6] Empty token (should fail):");
    match verifier.verify_id_token("").await {
        Ok(_) => println!("  ❌ Unexpected success - should have failed!"),
        Err(e) => println!("  ✅ Expected error: {}", e),
    }

    // Test 7: Expired ID token (if available)
    println!("\n[Test 7] Expired ID token (should fail):");
    if let Ok(expired_token) = env::var("EXPIRED_ID_TOKEN") {
        match verifier.verify_id_token(&expired_token).await {
            Ok(_) => println!("  ❌ Unexpected success - should have rejected expired token!"),
            Err(e) => println!("  ✅ Expected error: {}", e),
        }
    } else {
        println!("  ⊘ Skipped - EXPIRED_ID_TOKEN not set");
    }

    // Test 8: Expired access token (if available)
    println!("\n[Test 8] Expired access token (should fail):");
    if let Ok(expired_token) = env::var("EXPIRED_ACCESS_TOKEN") {
        match verifier.verify_access_token(&expired_token).await {
            Ok(_) => println!("  ❌ Unexpected success - should have rejected expired token!"),
            Err(e) => println!("  ✅ Expected error: {}", e),
        }
    } else {
        println!("  ⊘ Skipped - EXPIRED_ACCESS_TOKEN not set");
    }

    Ok(())
}

/// Prefetch JWKs to warm up the cache
async fn prefetch_jwks(verifier: &CognitoJwtVerifier) {
    println!("\nPrefetching JWKs:");
    let hydration_results = verifier.hydrate().await;
    for (pool_id, result) in hydration_results {
        match result {
            Ok(_) => println!("✅ Successfully prefetched JWKs for pool {}", pool_id),
            Err(e) => println!("❌ Failed to prefetch JWKs for pool {}: {}", pool_id, e),
        }
    }
}

#[tokio::main]
async fn main() -> Result<(), JwtError> {
    load_env();

    example_single_pool().await?;
    example_multiple_pools().await?;
    example_multiple_clients().await?;
    example_negative_cases().await?;

    Ok(())
}