jwt-verify 0.1.0

use jwt_verify::{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) => {
            println!("✅ ID Token verified successfully!");
            println!("  Subject: {}", claims.get_sub());
            println!("  Issuer: {}", claims.get_iss());
            println!("  Audience: {}", claims.get_aud());

            if let Some(email) = claims.get_email() {
                println!("  Email: {}", email);
                println!("  Email verified: {}", claims.is_email_verified());
            }

            if let Some(name) = claims.get_name() {
                println!("  Name: {}", name);
            }
        }
        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) => {
            println!("✅ Access Token verified successfully!");
            println!("  Subject: {}", claims.get_sub());
            println!("  Issuer: {}", claims.get_iss());
            println!("  Scopes: {:?}", claims.get_scopes());

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

            if let Some(client_id) = claims.get_client_id() {
                println!("  Client ID: {}", client_id);
            }
        }
        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(())
}