jwt-hack 2.6.0

use anyhow::Result;
use colored::Colorize;
use serde_json::Value;
use std::time::SystemTime;

use crate::jwt;
use crate::utils;

/// Helper function to format Unix timestamp to human-readable format
fn format_unix_timestamp(seconds: u64) -> String {
    let time = SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(seconds);
    chrono::DateTime::<chrono::Utc>::from(time)
        .format("%Y-%m-%d %H:%M:%S UTC")
        .to_string()
}

/// Annotate issued-at claim with human-readable timestamp in the JSON output
fn process_issued_at_claim(claims: &Value, claims_map: &mut Value) {
    if let Some(iat) = claims.get("iat") {
        if let Some(iat_val) = iat.as_f64() {
            let iat_seconds = iat_val as u64;
            let formatted_time = format_unix_timestamp(iat_seconds);

            if let Some(obj) = claims_map.as_object_mut() {
                obj.insert("iat_time".to_string(), Value::String(formatted_time));
            }
        }
    }
}

/// Annotate expiration claim with human-readable timestamp and status in the JSON output
fn process_expiration_claim(claims: &Value, claims_map: &mut Value) {
    if let Some(exp) = claims.get("exp") {
        if let Some(exp_val) = exp.as_f64() {
            let exp_seconds = exp_val as u64;
            let exp_time = SystemTime::UNIX_EPOCH + std::time::Duration::from_secs(exp_seconds);
            let formatted_time = format_unix_timestamp(exp_seconds);

            let now = SystemTime::now();
            let is_expired = now > exp_time;

            if let Some(obj) = claims_map.as_object_mut() {
                obj.insert("exp_time".to_string(), Value::String(formatted_time));
                obj.insert(
                    "exp_status".to_string(),
                    Value::String(if is_expired { "EXPIRED" } else { "VALID" }.to_string()),
                );
            }
        }
    }
}

/// Decodes and displays JWT token components with formatted output
pub fn execute(token: &str) {
    if let Err(e) = decode_token(token) {
        utils::log_error(format!("JWT Decode Error: {e}"));
        utils::log_error("e.g jwt-hack decode {JWT_CODE}");
    }
}

pub fn execute_json(token: &str) -> Result<Value> {
    decode_token_json(token)
}

fn decode_token(token: &str) -> Result<()> {
    // Detect token type first
    let token_type = jwt::detect_token_type(token);

    match token_type {
        jwt::TokenType::Jwt => decode_jwt_token(token),
        jwt::TokenType::Jwe => decode_jwe_token(token),
        jwt::TokenType::Unknown => {
            let parts: Vec<&str> = token.split('.').collect();
            Err(anyhow::anyhow!(
                "Unknown token format: expected 3 parts (JWT) or 5 parts (JWE), got {} parts",
                parts.len()
            ))
        }
    }
}

fn decode_token_json(token: &str) -> Result<Value> {
    let token_type = jwt::detect_token_type(token);

    match token_type {
        jwt::TokenType::Jwt => decode_jwt_token_json(token),
        jwt::TokenType::Jwe => decode_jwe_token_json(token),
        jwt::TokenType::Unknown => {
            let parts: Vec<&str> = token.split('.').collect();
            Err(anyhow::anyhow!(
                "Unknown token format: expected 3 parts (JWT) or 5 parts (JWE), got {} parts",
                parts.len()
            ))
        }
    }
}

fn decode_jwt_token(token: &str) -> Result<()> {
    let decoded = jwt::decode(token)?;

    let alg_str = format!("{:?}", decoded.algorithm);
    let typ = decoded
        .header
        .get("typ")
        .and_then(|v| v.as_str())
        .unwrap_or("JWT");

    println!("  {:<14}{}", "Algorithm".bold(), alg_str.cyan());
    println!("  {:<14}{}", "Type".bold(), typ);

    println!("\n  {}", "Header".bold());
    let header_json = serde_json::to_string_pretty(&decoded.header)?;
    println!("  {}", header_json.replace('\n', "\n  "));

    let mut claims_map: Value = decoded.claims.clone();
    process_issued_at_claim(&decoded.claims, &mut claims_map);
    process_expiration_claim(&decoded.claims, &mut claims_map);

    println!("\n  {}", "Payload".bold());
    let payload_json = serde_json::to_string_pretty(&claims_map)?;
    println!("  {}", payload_json.replace('\n', "\n  "));

    Ok(())
}

fn decode_jwt_token_json(token: &str) -> Result<Value> {
    let decoded = jwt::decode(token)?;

    let alg_str = format!("{:?}", decoded.algorithm);
    let typ = decoded
        .header
        .get("typ")
        .and_then(|v| v.as_str())
        .unwrap_or("JWT");

    let mut claims_map: Value = decoded.claims.clone();
    process_issued_at_claim(&decoded.claims, &mut claims_map);
    process_expiration_claim(&decoded.claims, &mut claims_map);

    Ok(serde_json::json!({
        "success": true,
        "token_type": "jwt",
        "algorithm": alg_str,
        "typ": typ,
        "header": decoded.header,
        "payload": claims_map
    }))
}

fn decode_jwe_token(token: &str) -> Result<()> {
    let decoded = jwt::decode_jwe(token)?;

    println!("  {:<14}{}", "Key Mgmt".bold(), decoded.algorithm.cyan());
    println!("  {:<14}{}", "Encryption".bold(), decoded.encryption.cyan());

    println!("\n  {}", "Header".bold());
    let header_json = serde_json::to_string_pretty(&decoded.header)?;
    println!("  {}", header_json.replace('\n', "\n  "));

    println!(
        "\n  {:<18}{}",
        "Encrypted Key".bold(),
        if decoded.encrypted_key.is_empty() {
            "(empty)".dimmed().to_string()
        } else {
            utils::format_base64_preview(&decoded.encrypted_key)
        }
    );
    println!(
        "  {:<18}{}",
        "IV".bold(),
        if decoded.iv.is_empty() {
            "(empty)".dimmed().to_string()
        } else {
            utils::format_base64_preview(&decoded.iv)
        }
    );
    println!(
        "  {:<18}{}",
        "Ciphertext".bold(),
        utils::format_base64_preview(&decoded.ciphertext)
    );
    println!(
        "  {:<18}{}",
        "Auth Tag".bold(),
        if decoded.tag.is_empty() {
            "(empty)".dimmed().to_string()
        } else {
            utils::format_base64_preview(&decoded.tag)
        }
    );

    eprintln!(
        "\n  {}",
        "JWE payload is encrypted and cannot be decoded without the appropriate key".dimmed()
    );

    // Check for misconfigurations
    let misconfigs = jwt::detect_jwe_misconfigurations(&decoded);
    if !misconfigs.is_empty() {
        eprintln!("\n  {}", "Security Issues Detected:".yellow().bold());
        for issue in misconfigs {
            eprintln!("  {}", issue);
        }
    }

    Ok(())
}

fn decode_jwe_token_json(token: &str) -> Result<Value> {
    let decoded = jwt::decode_jwe(token)?;

    let misconfigs = jwt::detect_jwe_misconfigurations(&decoded);

    Ok(serde_json::json!({
        "success": true,
        "token_type": "jwe",
        "key_mgmt": decoded.algorithm,
        "encryption": decoded.encryption,
        "header": decoded.header,
        "encrypted_key": decoded.encrypted_key,
        "iv": decoded.iv,
        "ciphertext": decoded.ciphertext,
        "auth_tag": decoded.tag,
        "security_issues": misconfigs,
        "note": "JWE payload is encrypted and cannot be decoded without the appropriate key"
    }))
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::Utc;
    use serde_json::json;
    use std::time::{Duration, UNIX_EPOCH};

    #[test]
    fn test_execute_valid_token() {
        // Create a valid token for testing
        let now = Utc::now();
        let claims = json!({
            "sub": "test_user",
            "name": "Test User",
            "iat": now.timestamp(),
            "exp": (now + chrono::Duration::days(1)).timestamp()
        });

        let token = jwt::encode(&claims, "", "HS256").expect("Failed to create test token");

        // Execute should not panic with a valid token
        let result = std::panic::catch_unwind(|| {
            execute(&token);
        });

        assert!(result.is_ok(), "execute() panicked with valid token");
    }

    #[test]
    fn test_execute_invalid_token() {
        // Create an invalid token for testing
        let invalid_token = "invalid.token.format";

        // Execute should handle the error and not panic
        let result = std::panic::catch_unwind(|| {
            execute(invalid_token);
        });

        assert!(result.is_ok(), "execute() panicked with invalid token");
    }

    #[test]
    fn test_decode_token_with_timestamps() {
        // Create a token with timestamp fields
        let now = SystemTime::now();
        let since_epoch = now.duration_since(UNIX_EPOCH).unwrap();
        let iat = since_epoch.as_secs() as i64;
        let exp = (since_epoch + Duration::from_secs(86400)).as_secs() as i64;

        let claims = json!({
            "sub": "test_user",
            "iat": iat,
            "exp": exp
        });

        let token = jwt::encode(&claims, "", "HS256").expect("Failed to create test token");

        // Test that decode_token processes timestamps correctly
        let result = decode_token(&token);
        assert!(
            result.is_ok(),
            "decode_token failed for valid token with timestamps"
        );
    }

    #[test]
    fn test_decode_token_without_timestamps() {
        // Create a token without timestamp fields
        let claims = json!({
            "sub": "test_user",
            "name": "Test User"
        });

        let token = jwt::encode(&claims, "", "HS256").expect("Failed to create test token");

        // Test that decode_token handles tokens without timestamps
        let result = decode_token(&token);
        assert!(
            result.is_ok(),
            "decode_token failed for valid token without timestamps"
        );
    }

    #[test]
    fn test_decode_jwe_token() {
        // Test JWE token decoding
        let jwe_token = "eyJhbGciOiJkaXIiLCJlbmMiOiJBMjU2R0NNIn0..ZHVtbXlfaXZfMTIzNDU2.eyJzdWIiOiJ0ZXN0In0.ZHVtbXlfdGFn";

        let result = decode_token(jwe_token);
        assert!(
            result.is_ok(),
            "decode_token should succeed for valid JWE token"
        );
    }

    #[test]
    fn test_decode_unknown_token_format() {
        // Test token with invalid number of parts
        let invalid_token = "invalid.token.with.too.many.parts.here";

        let result = decode_token(invalid_token);
        assert!(
            result.is_err(),
            "decode_token should fail for invalid token format"
        );
    }

    #[test]
    fn test_execute_json_success() {
        let claims = json!({ "sub": "u", "iat": 0, "exp": 1 });
        let token = jwt::encode(&claims, "", "HS256").expect("token");
        let value = execute_json(&token).expect("json decode");
        assert_eq!(value.get("success").and_then(|v| v.as_bool()), Some(true));
        assert_eq!(
            value.get("token_type").and_then(|v| v.as_str()),
            Some("jwt")
        );
        assert!(value.get("payload").is_some());
    }
}