prvng_platform/crates/control-center/src/services/jwt.rs

use std::sync::Arc;

use anyhow::Context;
use jsonwebtoken::{
    decode, encode, Algorithm, DecodingKey, EncodingKey, Header, TokenData, Validation,
};
// RSA key generation using rsa crate
use serde::{Deserialize, Serialize};
use tracing::info;
use uuid::Uuid;

use crate::error::{auth, ControlCenterError, Result};
use crate::models::{JwtClaims, TokenResponse};
// Use the configuration from simple_config
use crate::simple_config::JwtConfig;

/// RSA key pair for JWT signing
pub struct RsaKeys {
    pub private_key_pem: String,
    pub public_key_pem: String,
}

/// JWT service for handling token operations with RS256
#[derive(Clone)]
pub struct JwtService {
    config: JwtConfig,
    encoding_key: Arc<EncodingKey>,
    decoding_key: Arc<DecodingKey>,
    validation: Validation,
}

impl JwtService {
    /// Create a new JWT service with RS256 algorithm
    pub fn new(config: JwtConfig) -> Result<Self> {
        let encoding_key = EncodingKey::from_rsa_pem(config.private_key_pem.as_bytes())
            .context("Failed to create encoding key from private key PEM")?;

        let decoding_key = DecodingKey::from_rsa_pem(config.public_key_pem.as_bytes())
            .context("Failed to create decoding key from public key PEM")?;

        let mut validation = Validation::new(Algorithm::RS256);
        validation.set_issuer(&[&config.issuer]);
        validation.set_audience(&[&config.audience]);
        validation.validate_exp = true;
        validation.validate_nbf = false;

        info!("JWT service initialized with RS256 algorithm");

        Ok(Self {
            config,
            encoding_key: Arc::new(encoding_key),
            decoding_key: Arc::new(decoding_key),
            validation,
        })
    }

    /// Generate access token
    pub fn generate_access_token(
        &self,
        user_id: Uuid,
        session_id: Uuid,
        roles: Vec<String>,
    ) -> Result<String> {
        let claims = JwtClaims::new(
            user_id,
            session_id,
            roles,
            self.config.access_token_expiration_hours,
            self.config.issuer.clone(),
            self.config.audience.clone(),
        );

        let header = Header::new(Algorithm::RS256);

        encode(&header, &claims, &self.encoding_key)
            .context("Failed to encode access token")
            .map_err(ControlCenterError::from)
    }

    /// Generate refresh token (longer-lived, simpler claims)
    pub fn generate_refresh_token(&self, user_id: Uuid, session_id: Uuid) -> Result<String> {
        let claims = RefreshTokenClaims::new(
            user_id,
            session_id,
            self.config.refresh_token_expiration_hours,
            self.config.issuer.clone(),
            self.config.audience.clone(),
        );

        let header = Header::new(Algorithm::RS256);

        encode(&header, &claims, &self.encoding_key)
            .context("Failed to encode refresh token")
            .map_err(ControlCenterError::from)
    }

    /// Generate both access and refresh tokens
    pub fn generate_token_pair(
        &self,
        user_id: Uuid,
        session_id: Uuid,
        roles: Vec<String>,
    ) -> Result<TokenResponse> {
        let access_token = self.generate_access_token(user_id, session_id, roles)?;
        let refresh_token = self.generate_refresh_token(user_id, session_id)?;

        Ok(TokenResponse::new(
            access_token,
            refresh_token,
            self.config.access_token_expiration_hours * 3600, // Convert to seconds
            session_id.to_string(),
        ))
    }

    /// Verify and decode access token
    pub fn verify_access_token(&self, token: &str) -> Result<TokenData<JwtClaims>> {
        decode::<JwtClaims>(token, &self.decoding_key, &self.validation)
            .context("Failed to verify access token")
            .map_err(ControlCenterError::from)
    }

    /// Verify and decode refresh token
    pub fn verify_refresh_token(&self, token: &str) -> Result<TokenData<RefreshTokenClaims>> {
        decode::<RefreshTokenClaims>(token, &self.decoding_key, &self.validation)
            .context("Failed to verify refresh token")
            .map_err(ControlCenterError::from)
    }

    /// Extract user ID from token without full verification (for expired
    /// tokens)
    pub fn extract_user_id_unsafe(&self, token: &str) -> Result<Uuid> {
        let mut validation = self.validation.clone();
        validation.validate_exp = false; // Don't validate expiration

        let token_data = decode::<JwtClaims>(token, &self.decoding_key, &validation)
            .context("Failed to extract user ID from token")?;

        Uuid::parse_str(&token_data.claims.sub)
            .context("Invalid user ID in token")
            .map_err(ControlCenterError::from)
    }

    /// Get token expiration in seconds
    pub fn get_access_token_expiration(&self) -> i64 {
        self.config.access_token_expiration_hours * 3600
    }

    /// Get refresh token expiration in seconds
    pub fn get_refresh_token_expiration(&self) -> i64 {
        self.config.refresh_token_expiration_hours * 3600
    }
}

/// Refresh token claims (simpler than access token)
#[derive(Debug, Serialize, Deserialize)]
pub struct RefreshTokenClaims {
    pub sub: String,        // Subject (user_id)
    pub exp: i64,           // Expiration timestamp
    pub iat: i64,           // Issued at timestamp
    pub iss: String,        // Issuer
    pub aud: String,        // Audience
    pub session_id: String, // Session ID
    pub token_type: String, // "refresh"
}

impl RefreshTokenClaims {
    /// Create new refresh token claims
    pub fn new(
        user_id: Uuid,
        session_id: Uuid,
        expiration_hours: i64,
        issuer: String,
        audience: String,
    ) -> Self {
        let now = chrono::Utc::now();
        Self {
            sub: user_id.to_string(),
            exp: (now + chrono::Duration::hours(expiration_hours)).timestamp(),
            iat: now.timestamp(),
            iss: issuer,
            aud: audience,
            session_id: session_id.to_string(),
            token_type: "refresh".to_string(),
        }
    }
}

/// Generate RSA key pair for JWT signing (RS256)
pub fn generate_rsa_key_pair() -> Result<RsaKeys> {
    use rsa::rand_core::OsRng;
    use rsa::{RsaPrivateKey, RsaPublicKey};
    use rsa::pkcs8::{EncodePrivateKey, EncodePublicKey, LineEnding};

    // Generate 2048-bit RSA key pair with OS randomness for cryptographic security
    let private_key =
        RsaPrivateKey::new(&mut OsRng, 2048).context("Failed to generate RSA private key")?;

    let public_key = RsaPublicKey::from(&private_key);

    // Convert to PKCS#8 PEM format (required by jsonwebtoken and signature crates)
    let private_key_pem = private_key
        .to_pkcs8_pem(LineEnding::LF)
        .context("Failed to encode private key as PKCS#8 PEM")?
        .to_string();

    let public_key_pem = public_key
        .to_public_key_pem(LineEnding::LF)
        .context("Failed to encode public key as PEM")?;

    Ok(RsaKeys {
        private_key_pem,
        public_key_pem,
    })
}

/// Load RSA key pair from files
pub fn load_rsa_keys_from_files(private_key_path: &str, public_key_path: &str) -> Result<RsaKeys> {
    let private_key_pem = std::fs::read_to_string(private_key_path)
        .with_context(|| format!("Failed to read private key from {}", private_key_path))?;

    let public_key_pem = std::fs::read_to_string(public_key_path)
        .with_context(|| format!("Failed to read public key from {}", public_key_path))?;

    // Validate keys by creating encoding/decoding keys
    EncodingKey::from_rsa_pem(private_key_pem.as_bytes())
        .context("Invalid private key PEM format")?;

    DecodingKey::from_rsa_pem(public_key_pem.as_bytes())
        .context("Invalid public key PEM format")?;

    Ok(RsaKeys {
        private_key_pem,
        public_key_pem,
    })
}

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

    fn create_test_jwt_config() -> JwtConfig {
        let keys = generate_rsa_key_pair().expect("Failed to generate test keys");
        JwtConfig {
            issuer: "test-issuer".to_string(),
            audience: "test-audience".to_string(),
            access_token_expiration_hours: 1,
            refresh_token_expiration_hours: 24,
            private_key_pem: keys.private_key_pem,
            public_key_pem: keys.public_key_pem,
        }
    }

    #[tokio::test]
    async fn test_jwt_token_generation_and_verification() {
        let config = create_test_jwt_config();
        let jwt_service = JwtService::new(config)
            .expect("Failed to create JWT service for test");

        let user_id = Uuid::new_v4();
        let session_id = Uuid::new_v4();
        let roles = vec!["user".to_string(), "admin".to_string()];

        // Generate tokens
        let token_response = jwt_service
            .generate_token_pair(user_id, session_id, roles.clone())
            .expect("Failed to generate token pair for test");

        // Verify access token
        let access_claims = jwt_service
            .verify_access_token(&token_response.access_token)
            .expect("Failed to verify access token for test");
        assert_eq!(access_claims.claims.sub, user_id.to_string());
        assert_eq!(access_claims.claims.session_id, session_id.to_string());
        assert_eq!(access_claims.claims.roles, roles);

        // Verify refresh token
        let refresh_claims = jwt_service
            .verify_refresh_token(&token_response.refresh_token)
            .unwrap();
        assert_eq!(refresh_claims.claims.sub, user_id.to_string());
        assert_eq!(refresh_claims.claims.session_id, session_id.to_string());
        assert_eq!(refresh_claims.claims.token_type, "refresh");
    }

    #[tokio::test]
    async fn test_key_generation() {
        let keys = generate_rsa_key_pair().unwrap();

        // Keys should be valid PEM format
        assert!(keys.private_key_pem.contains("-----BEGIN PRIVATE KEY-----"));
        assert!(keys.private_key_pem.contains("-----END PRIVATE KEY-----"));
        assert!(keys.public_key_pem.contains("-----BEGIN PUBLIC KEY-----"));
        assert!(keys.public_key_pem.contains("-----END PUBLIC KEY-----"));

        // Should be able to create JWT service with generated keys
        let config = JwtConfig {
            private_key_pem: keys.private_key_pem,
            public_key_pem: keys.public_key_pem,
            ..JwtConfig::default()
        };

        let _jwt_service = JwtService::new(config).unwrap();
    }
}