jwt-verify 0.1.0

use jwt_verify::{JwtError, JwtVerifier, OidcJwtVerifier, OidcProviderConfig};
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 provider verification
async fn example_single_provider() -> Result<(), JwtError> {
    println!("Example 1: Basic OIDC JWT verification with a single provider");

    let issuer = get_env_or_default("OIDC_ISSUER", "https://accounts.example.com");
    let jwks_url = env::var("OIDC_JWKS_URL")
        .unwrap_or_else(|_| format!("{}/.well-known/jwks.json", issuer.trim_end_matches('/')));
    let client_id = get_env_or_default("OIDC_CLIENT_ID", "client1");

    println!("Using OIDC Provider:");
    println!("  Issuer: {}", issuer);
    println!("  JWKS URL: {}", jwks_url);
    println!("  Client ID: {}", client_id);

    let config = OidcProviderConfig::new(&issuer, Some(&jwks_url), &[client_id], None)?;
    let verifier = OidcJwtVerifier::new(vec![config])?;

    let id_token = env::var("OIDC_ID_TOKEN").unwrap_or_else(|_| {
        println!("Warning: OIDC_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 providers verification
async fn example_multiple_providers() -> Result<(), JwtError> {
    println!("\nExample 2: OIDC JWT verification with multiple providers");

    let issuer = get_env_or_default("OIDC_ISSUER", "https://accounts.example.com");
    let jwks_url = env::var("OIDC_JWKS_URL")
        .unwrap_or_else(|_| format!("{}/.well-known/jwks.json", issuer.trim_end_matches('/')));
    let client_id = get_env_or_default("OIDC_CLIENT_ID", "client1");
    let client_id3 = get_env_or_default("OIDC_CLIENT_ID_3", "client3");

    let issuer2 = get_env_or_default("OIDC_ISSUER_2", "https://accounts.example2.com");
    let jwks_url2 = env::var("OIDC_JWKS_URL_2")
        .unwrap_or_else(|_| format!("{}/.well-known/jwks.json", issuer2.trim_end_matches('/')));
    let client_id2 = get_env_or_default("OIDC_CLIENT_ID_2", "client2");

    println!("Using multiple OIDC Providers:");
    println!("  1. Issuer: {}", issuer);
    println!("  2. Issuer: {}", issuer2);

    let provider1 =
        OidcProviderConfig::new(&issuer, Some(&jwks_url), &[client_id, client_id3], None)?;
    let provider2 = OidcProviderConfig::new(&issuer2, Some(&jwks_url2), &[client_id2], None)?;

    let multi_provider_verifier = OidcJwtVerifier::new(vec![provider1, provider2])?;

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

    verify_access_token(&multi_provider_verifier, &access_token).await;

    Ok(())
}

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

    let issuer = get_env_or_default("OIDC_ISSUER", "https://accounts.example.com");
    let jwks_url = env::var("OIDC_JWKS_URL")
        .unwrap_or_else(|_| format!("{}/.well-known/jwks.json", issuer.trim_end_matches('/')));
    let client_id = get_env_or_default("OIDC_CLIENT_ID", "client1");
    let client_id2 = get_env_or_default("OIDC_CLIENT_ID_2", "client2");

    println!("Using single provider with multiple client IDs:");
    println!("  Issuer: {}", issuer);
    println!("  Client ID 1: {}", client_id);
    println!("  Client ID 2: {}", client_id2);

    let multi_client_config = OidcProviderConfig::new(
        &issuer,
        Some(&jwks_url),
        &[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 = OidcJwtVerifier::new(vec![multi_client_config])?;

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

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

    let access_token = env::var("OIDC_ACCESS_TOKEN").unwrap_or_else(|_| {
        println!("Warning: OIDC_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: &OidcJwtVerifier, 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: &OidcJwtVerifier, 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 issuer = get_env_or_default("OIDC_ISSUER", "https://accounts.example.com");
    let jwks_url = env::var("OIDC_JWKS_URL")
        .unwrap_or_else(|_| format!("{}/.well-known/jwks.json", issuer.trim_end_matches('/')));
    let client_id = get_env_or_default("OIDC_CLIENT_ID", "client1");

    let config = OidcProviderConfig::new(&issuer, Some(&jwks_url), &[client_id], None)?;
    let verifier = OidcJwtVerifier::new(vec![config])?;

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

    let access_token = env::var("OIDC_ACCESS_TOKEN").unwrap_or_else(|_| {
        println!("Warning: OIDC_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 provider)
    let disallowed_client_id = get_env_or_default("OIDC_CLIENT_ID_2", "client2");
    println!(
        "\n[Test 3] Token from disallowed client ID (allowed: {}, token from: {}):",
        get_env_or_default("OIDC_CLIENT_ID", "client1"),
        disallowed_client_id
    );

    if let Ok(other_client_token) = env::var("OIDC_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 - OIDC_ID_TOKEN_CLIENT2 not set");
    }

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

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

    if let Ok(different_provider_access_token) = env::var("OIDC_ACCESS_TOKEN_CLIENT3") {
        match verifier
            .verify_access_token(&different_provider_access_token)
            .await
        {
            Ok(_) => {
                println!(
                    "  ❌ Unexpected success - should have rejected token from different provider!"
                )
            }
            Err(e) => println!("  ✅ Expected error (access token): {}", e),
        }
    } else {
        println!("  ⊘ Skipped (access token) - OIDC_ACCESS_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("OIDC_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 - OIDC_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("OIDC_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 - OIDC_EXPIRED_ACCESS_TOKEN not set");
    }

    Ok(())
}

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

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

    example_single_provider().await?;
    example_multiple_providers().await?;
    example_multiple_clients().await?;
    example_negative_cases().await?;

    Ok(())
}