nushell-plugins/nu_plugin_auth/src/auth.rs

//! JWT authentication and verification logic.
//!
//! This module provides RS256 JWT token verification, claims extraction,
//! and token management utilities.

use base64::{engine::general_purpose, Engine as _};
use jsonwebtoken::{decode, decode_header, Algorithm, DecodingKey, Validation};
use serde::{Deserialize, Serialize};

use crate::error::{AuthError, AuthErrorKind};

/// Default Control Center URL for authentication.
pub const DEFAULT_CONTROL_CENTER_URL: &str = "http://localhost:8081";

/// JWT Claims structure for provisioning platform tokens.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Claims {
    /// Subject (user ID)
    pub sub: String,
    /// Username
    pub username: String,
    /// User email
    pub email: String,
    /// User roles
    pub roles: Vec<String>,
    /// Expiration time (Unix timestamp)
    pub exp: i64,
    /// Issued at time (Unix timestamp)
    pub iat: i64,
    /// Not before time (Unix timestamp)
    #[serde(default)]
    pub nbf: i64,
    /// JWT ID (unique identifier)
    #[serde(default)]
    pub jti: String,
    /// Issuer
    #[serde(default)]
    pub iss: String,
    /// Audience
    #[serde(default)]
    pub aud: String,
}

/// Result of token verification.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VerificationResult {
    /// Whether the token is valid
    pub valid: bool,
    /// Extracted claims if valid
    pub claims: Option<Claims>,
    /// Error message if invalid
    pub error: Option<String>,
    /// Time remaining until expiration (seconds)
    pub expires_in: Option<i64>,
}

impl VerificationResult {
    /// Creates a successful verification result.
    pub fn success(claims: Claims) -> Self {
        let now = chrono::Utc::now().timestamp();
        let expires_in = claims.exp - now;
        Self {
            valid: true,
            claims: Some(claims),
            error: None,
            expires_in: Some(expires_in),
        }
    }

    /// Creates a failed verification result.
    pub fn failure(error: impl Into<String>) -> Self {
        Self {
            valid: false,
            claims: None,
            error: Some(error.into()),
            expires_in: None,
        }
    }
}

/// Decodes and extracts claims from a JWT without verification.
///
/// This is useful for inspecting token contents before verification
/// or when verification is not possible (e.g., no public key available).
///
/// # Arguments
///
/// * `token` - The JWT token string
///
/// # Returns
///
/// Returns the decoded claims or an error if the token format is invalid.
pub fn decode_claims_unverified(token: &str) -> Result<Claims, AuthError> {
    let parts: Vec<&str> = token.split('.').collect();
    if parts.len() != 3 {
        return Err(AuthError::invalid_token(
            "Token must have 3 parts separated by '.'",
        ));
    }

    let payload = parts[1];
    let decoded = general_purpose::URL_SAFE_NO_PAD
        .decode(payload)
        .map_err(|e| AuthError::invalid_token(format!("Failed to decode payload: {}", e)))?;

    let claims: Claims = serde_json::from_slice(&decoded)
        .map_err(|e| AuthError::invalid_token(format!("Failed to parse claims: {}", e)))?;

    Ok(claims)
}

/// Verifies a JWT token using RS256 algorithm with the provided public key.
///
/// # Arguments
///
/// * `token` - The JWT token string
/// * `public_key_pem` - The RSA public key in PEM format
///
/// # Returns
///
/// Returns a VerificationResult indicating whether the token is valid
/// and containing the claims if verification succeeded.
pub fn verify_token_rs256(
    token: &str,
    public_key_pem: &str,
) -> Result<VerificationResult, AuthError> {
    // Verify the token header uses RS256
    let header = decode_header(token)
        .map_err(|e| AuthError::invalid_token(format!("Failed to decode header: {}", e)))?;

    if header.alg != Algorithm::RS256 {
        return Err(AuthError::new(
            AuthErrorKind::SignatureVerificationFailed,
            format!("Expected RS256 algorithm, got {:?}", header.alg),
        ));
    }

    // Create decoding key from PEM
    let decoding_key = DecodingKey::from_rsa_pem(public_key_pem.as_bytes())
        .map_err(|e| AuthError::configuration_error(format!("Invalid public key: {}", e)))?;

    // Set up validation
    let mut validation = Validation::new(Algorithm::RS256);
    validation.validate_exp = true;
    validation.validate_nbf = true;

    // Decode and verify
    match decode::<Claims>(token, &decoding_key, &validation) {
        Ok(token_data) => Ok(VerificationResult::success(token_data.claims)),
        Err(e) => {
            let error_msg = match e.kind() {
                jsonwebtoken::errors::ErrorKind::ExpiredSignature => "Token has expired",
                jsonwebtoken::errors::ErrorKind::InvalidSignature => "Invalid signature",
                jsonwebtoken::errors::ErrorKind::InvalidToken => "Invalid token format",
                jsonwebtoken::errors::ErrorKind::InvalidIssuer => "Invalid issuer",
                jsonwebtoken::errors::ErrorKind::InvalidAudience => "Invalid audience",
                _ => "Token verification failed",
            };
            Ok(VerificationResult::failure(error_msg))
        }
    }
}

/// Verifies a JWT token locally by checking format and expiration.
///
/// This performs basic validation without cryptographic verification:
/// - Token format (3 parts)
/// - Expiration time
/// - Not-before time
///
/// Use this for quick local checks when the public key is not available.
pub fn verify_token_local(token: &str) -> Result<VerificationResult, AuthError> {
    let claims = decode_claims_unverified(token)?;
    let now = chrono::Utc::now().timestamp();

    // Check expiration
    if claims.exp < now {
        return Ok(VerificationResult::failure(format!(
            "Token expired {} seconds ago",
            now - claims.exp
        )));
    }

    // Check not-before (if set)
    if claims.nbf > 0 && claims.nbf > now {
        return Ok(VerificationResult::failure(format!(
            "Token not valid for {} more seconds",
            claims.nbf - now
        )));
    }

    Ok(VerificationResult::success(claims))
}

/// Checks if a token is expired.
pub fn is_token_expired(token: &str) -> Result<bool, AuthError> {
    let claims = decode_claims_unverified(token)?;
    let now = chrono::Utc::now().timestamp();
    Ok(claims.exp < now)
}

/// Gets the time remaining until token expiration in seconds.
pub fn get_token_expiry_seconds(token: &str) -> Result<i64, AuthError> {
    let claims = decode_claims_unverified(token)?;
    let now = chrono::Utc::now().timestamp();
    Ok(claims.exp - now)
}

/// Extracts the user ID from a token.
pub fn get_user_id(token: &str) -> Result<String, AuthError> {
    let claims = decode_claims_unverified(token)?;
    Ok(claims.sub)
}

/// Extracts the username from a token.
pub fn get_username(token: &str) -> Result<String, AuthError> {
    let claims = decode_claims_unverified(token)?;
    Ok(claims.username)
}

/// Extracts the roles from a token.
pub fn get_roles(token: &str) -> Result<Vec<String>, AuthError> {
    let claims = decode_claims_unverified(token)?;
    Ok(claims.roles)
}

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

    // Test token (expired, but valid format)
    const TEST_TOKEN: &str = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJ1c2VyLTEyMyIsInVzZXJuYW1lIjoiYWRtaW4iLCJlbWFpbCI6ImFkbWluQGV4YW1wbGUuY29tIiwicm9sZXMiOlsiYWRtaW4iLCJ1c2VyIl0sImV4cCI6MTcwMDAwMDAwMCwiaWF0IjoxNjk5OTk2NDAwLCJuYmYiOjAsImp0aSI6Imp0aS0xMjMiLCJpc3MiOiJwcm92aXNpb25pbmciLCJhdWQiOiJwcm92aXNpb25pbmctY2xpIn0.signature";

    #[test]
    fn test_decode_claims_unverified_invalid_format() {
        let result = decode_claims_unverified("not.a.valid.token.format");
        assert!(result.is_err());
    }

    #[test]
    fn test_decode_claims_unverified_missing_parts() {
        let result = decode_claims_unverified("only.two");
        assert!(result.is_err());
        let error = result.unwrap_err();
        assert_eq!(error.kind, AuthErrorKind::InvalidToken);
    }

    #[test]
    fn test_verification_result_success() {
        let claims = Claims {
            sub: "user-123".to_string(),
            username: "admin".to_string(),
            email: "admin@example.com".to_string(),
            roles: vec!["admin".to_string()],
            exp: chrono::Utc::now().timestamp() + 3600,
            iat: chrono::Utc::now().timestamp(),
            nbf: 0,
            jti: "jti-123".to_string(),
            iss: "provisioning".to_string(),
            aud: "cli".to_string(),
        };
        let result = VerificationResult::success(claims);
        assert!(result.valid);
        assert!(result.claims.is_some());
        assert!(result.expires_in.is_some());
    }

    #[test]
    fn test_verification_result_failure() {
        let result = VerificationResult::failure("test error");
        assert!(!result.valid);
        assert!(result.claims.is_none());
        assert_eq!(result.error, Some("test error".to_string()));
    }

    #[test]
    fn test_is_token_expired_handles_invalid() {
        let result = is_token_expired("invalid");
        assert!(result.is_err());
    }
}