jwt-hack 2.5.0

use anyhow::Result;
use colored::Colorize;
use serde_json::Value;
use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;

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

/// Options for encoding operations
#[allow(dead_code)]
pub struct EncodeOptions {
    pub secret: Option<String>,
    pub private_key_path: Option<PathBuf>,
    pub algorithm: String,
    pub no_signature: bool,
    pub headers: Vec<(String, String)>,
    pub compress: bool,
    pub jwe: bool,
}

/// Encodes JSON data into a JWT token with various algorithm and signing options
#[allow(clippy::too_many_arguments)]
pub fn execute(
    json_str: &str,
    secret: Option<&str>,
    private_key_path: Option<&PathBuf>,
    algorithm: &str,
    no_signature: bool,
    headers: &[(String, String)],
    compress: bool,
    jwe: bool,
) {
    if jwe {
        if let Err(e) = encode_jwe(json_str, secret) {
            utils::log_error(format!("JWE Encode Error: {e}"));
            utils::log_error("e.g jwt-hack encode {JSON} --jwe --secret={YOUR_SECRET}");
        }
    } else if let Err(e) = encode_json(
        json_str,
        secret,
        private_key_path,
        algorithm,
        no_signature,
        headers,
        compress,
    ) {
        utils::log_error(format!("JSON Encode Error: {e}"));
        utils::log_error("e.g jwt-hack encode {JSON} --secret={YOUR_SECRET}");
        utils::log_error(
            "or with RSA: jwt-hack encode {JSON} --private-key=private.pem --algorithm=RS256",
        );
    }
}

/// Helper function to convert header vector to optional hashmap
fn create_header_map(headers: &[(String, String)]) -> Option<HashMap<&str, &str>> {
    if headers.is_empty() {
        None
    } else {
        let mut map = HashMap::new();
        for (key, value) in headers {
            map.insert(key.as_str(), value.as_str());
        }
        Some(map)
    }
}

/// Helper function to display encoding success information
fn display_encoding_result(
    token: &str,
    algorithm: &str,
    key_info: &str,
    headers: &[(String, String)],
) {
    println!("  {:<14}{}", "Algorithm".bold(), algorithm.cyan());
    println!("  {:<14}{}", "Key".bold(), key_info);

    if !headers.is_empty() {
        println!("\n  {}", "Headers".bold());
        for (key, value) in headers {
            println!("  {:<14}{}", key.to_string().dimmed(), value);
        }
    }

    println!("\n  {}", "Token".bold());
    println!("  {}", utils::format_jwt_token(token));
}

fn encode_json(
    json_str: &str,
    secret: Option<&str>,
    private_key_path: Option<&PathBuf>,
    algorithm: &str,
    no_signature: bool,
    headers: &[(String, String)],
    compress: bool,
) -> Result<()> {
    // Parse the input JSON into a Value object
    let claims: Value = serde_json::from_str(json_str)?;

    // Convert custom header key-value pairs into a hashmap for JWT encoding
    let header_map = create_header_map(headers);

    // Build JWT encoding options based on provided parameters
    // Private key option is handled separately due to Rust lifetime requirements
    let options = if no_signature {
        // Use 'none' algorithm (creates unsigned JWT token)
        jwt::EncodeOptions {
            algorithm: "none",
            key_data: jwt::KeyData::None,
            header_params: header_map,
            compress_payload: compress,
        }
    } else if let Some(path) = private_key_path {
        // Read RSA/EC private key from file for asymmetric algorithms
        let key_content = fs::read_to_string(path)?;

        // Create encoding options with the private key content (keeping ownership in this scope)
        let options = jwt::EncodeOptions {
            algorithm,
            key_data: jwt::KeyData::PrivateKeyPem(&key_content),
            header_params: header_map,
            compress_payload: compress,
        };

        // Encode JWT immediately while private key content is in scope
        let token = jwt::encode_with_options(&claims, &options)?;

        let key_info = format!("{} ({})", path.display(), "Private Key".dimmed());
        display_encoding_result(&token, algorithm, &key_info, headers);

        return Ok(());
    } else {
        // Default case: use HMAC with provided secret (or empty string)
        jwt::EncodeOptions {
            algorithm,
            key_data: jwt::KeyData::Secret(secret.unwrap_or("")),
            header_params: header_map,
            compress_payload: compress,
        }
    };

    // Encode the JWT token using the configured options
    let token = jwt::encode_with_options(&claims, &options)?;

    // Determine key information display based on signature type
    let key_info = if no_signature || secret.unwrap_or("").is_empty() {
        "None (unsigned)".dimmed().to_string()
    } else {
        "****".to_string()
    };

    display_encoding_result(&token, algorithm, &key_info, headers);

    Ok(())
}

fn encode_jwe(json_str: &str, secret: Option<&str>) -> Result<()> {
    let _claims: Value = serde_json::from_str(json_str)?;
    let key = secret.unwrap_or("default_jwe_key");
    let token = jwt::encode_jwe_demo(json_str, key)?;

    println!("  {:<14}{}", "Key Mgmt".bold(), "dir".cyan());
    println!("  {:<14}{}", "Encryption".bold(), "A256GCM".cyan());

    println!("\n  {}", "Token".bold());
    println!("  {}", token);

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;

    #[test]
    fn test_create_header_map() {
        // Test empty headers
        let headers = Vec::new();
        assert_eq!(create_header_map(&headers), None);

        // Test single header
        let headers = vec![("key1".to_string(), "value1".to_string())];
        let map = create_header_map(&headers).expect("Should return Some map");
        assert_eq!(map.len(), 1);
        assert_eq!(map.get("key1"), Some(&"value1"));

        // Test multiple headers
        let headers = vec![
            ("key1".to_string(), "value1".to_string()),
            ("key2".to_string(), "value2".to_string()),
        ];
        let map = create_header_map(&headers).expect("Should return Some map");
        assert_eq!(map.len(), 2);
        assert_eq!(map.get("key1"), Some(&"value1"));
        assert_eq!(map.get("key2"), Some(&"value2"));
    }

    #[test]
    fn test_execute_with_secret() {
        // Create a simple JSON payload
        let json_str = r#"{"sub":"1234567890","name":"John Doe"}"#;
        let secret = Some("test_secret");
        let private_key_path = None;
        let algorithm = "HS256";
        let no_signature = false;
        let headers = Vec::new();

        // Execute should not panic
        let result = std::panic::catch_unwind(|| {
            execute(
                json_str,
                secret,
                private_key_path,
                algorithm,
                no_signature,
                &headers,
                false, // compress
                false, // jwe
            );
        });

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

    #[test]
    fn test_execute_with_no_signature() {
        // Create a simple JSON payload
        let json_str = r#"{"sub":"1234567890","name":"John Doe"}"#;
        let secret = None;
        let private_key_path = None;
        let algorithm = "none";
        let no_signature = true;
        let headers = Vec::new();

        // Execute should not panic
        let result = std::panic::catch_unwind(|| {
            execute(
                json_str,
                secret,
                private_key_path,
                algorithm,
                no_signature,
                &headers,
                false, // compress
                false, // jwe
            );
        });

        assert!(result.is_ok(), "execute() panicked with no signature");
    }

    #[test]
    fn test_execute_with_custom_headers() {
        // Create a simple JSON payload
        let json_str = r#"{"sub":"1234567890","name":"John Doe"}"#;
        let secret = Some("test_secret");
        let private_key_path = None;
        let algorithm = "HS256";
        let no_signature = false;
        let headers = vec![
            ("kid".to_string(), "1234".to_string()),
            ("typ".to_string(), "JWT+AT".to_string()),
        ];

        // Execute should not panic
        let result = std::panic::catch_unwind(|| {
            execute(
                json_str,
                secret,
                private_key_path,
                algorithm,
                no_signature,
                &headers,
                false, // compress
                false, // jwe
            );
        });

        assert!(result.is_ok(), "execute() panicked with custom headers");
    }

    #[test]
    fn test_execute_with_invalid_json() {
        // Create an invalid JSON payload
        let json_str = r#"{"sub":"1234567890","name":"John Doe"#; // Missing closing brace
        let secret = Some("test_secret");
        let private_key_path = None;
        let algorithm = "HS256";
        let no_signature = false;
        let headers = Vec::new();

        // Execute should handle the error and not panic
        let result = std::panic::catch_unwind(|| {
            execute(
                json_str,
                secret,
                private_key_path,
                algorithm,
                no_signature,
                &headers,
                false, // compress
                false, // jwe
            );
        });

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

    #[test]
    fn test_encode_json_with_rsa_key() {
        // This test requires creating a temporary RSA key file
        let temp_dir = tempdir().expect("Failed to create temp directory");
        let key_path = temp_dir.path().join("test_key.pem");

        // Write sample RSA private key (this is just a placeholder for testing)
        let sample_key = "-----BEGIN RSA PRIVATE KEY-----\nMIIEogIBAAKCAQEAnzyis1ZjfNB0bBgKFMSvvkTtwlvBsaJq7S5wA+kzeVOVpVWw\nkWdVha4s38XM/pa/yr47av7+z3VTmvDRyAHcaT92whREFpLv9cj5lTeJSibyr/Mr\nm/YtjCZVWgaOYIhwrXwKLqPr/11inWsAkfIytvHWTxZYEcXLgAXFuUuaS3uF9gEi\nNQwzGTU1v0FqkqTBr4B8nW3HCN47XUu0t8Y0e+lf4s4OxQawWD79J9/5d3Ry0vbV\n3Am1FtGJiJvOwRsIfVChDpYStTcHTCMqtvWbV6L11BWkpzGXSW4Hv43qa+GSYOD2\nQU68Mb59oSk2OB+BtOLpJofmbGEGgvmwyCI9MwIDAQAB\n-----END RSA PRIVATE KEY-----";
        std::fs::write(&key_path, sample_key).expect("Failed to write test key file");

        // Create a simple JSON payload
        let json_str = r#"{"sub":"1234567890","name":"John Doe"}"#;
        let secret = None;
        let private_key_path = Some(&key_path);
        let algorithm = "RS256";
        let no_signature = false;
        let headers = Vec::new();

        // Execute with RSA key shouldn't panic (even if the key is invalid for actual signing)
        let result = std::panic::catch_unwind(|| {
            encode_json(
                json_str,
                secret,
                private_key_path,
                algorithm,
                no_signature,
                &headers,
                false, // compress
            )
        });

        assert!(
            result.is_err() || result.is_ok(),
            "Properly handled RSA key attempt"
        );

        // Clean up
        temp_dir.close().expect("Failed to clean up temp directory");
    }

    #[test]
    fn test_execute_with_jwe_flag() {
        // Test JWE encoding execution
        let json_str = r#"{"sub":"test","name":"JWE User"}"#;
        let secret = Some("test_secret");
        let private_key_path = None;
        let algorithm = "HS256";
        let no_signature = false;
        let headers = Vec::new();
        let compress = false;
        let jwe = true;

        // Execute with JWE flag shouldn't panic
        let result = std::panic::catch_unwind(|| {
            execute(
                json_str,
                secret,
                private_key_path,
                algorithm,
                no_signature,
                &headers,
                compress,
                jwe,
            );
        });

        assert!(result.is_ok(), "execute() with JWE flag should not panic");
    }

    #[test]
    fn test_encode_jwe_function() {
        // Test the encode_jwe function directly
        let json_str = r#"{"sub":"test","name":"JWE User"}"#;
        let secret = Some("test_secret");

        let result = encode_jwe(json_str, secret);
        assert!(result.is_ok(), "encode_jwe should succeed with valid JSON");
    }
}