ones-oidc 0.3.7

use super::device::{get_jwks, make_device_jwt, make_device_jwt_ciba, request_device_access_token, device_access_token};
use super::oidc_types::JwtPayload;
use super::oidc_types::{
    AuthenticationMethod, AuthenticationResult, CibaResponse, CibaStatusResponse,
    ClientCredentialsIntrospection, SubjectIdentity,
};
use super::{AuthenticatedEntityKind, LoginHint};
use super::http_client::default_client;
use super::config::OnesOidcConfig;
use crate::errors::{DeviceError, OidcError, OidcRequirementsError};
use crate::oidc_backend::QrStatusRequest;
use crate::oidc_types::{LoginHintKind, OidcErrorResponse, QrAuthSessionIdp};
use jsonwebtoken::jwk::JwkSet;
use jsonwebtoken::EncodingKey;
use log::debug;
use openidconnect::core::{CoreProviderMetadata, CoreTokenType};
use openidconnect::{
    AccessToken, ClientId, EmptyExtraTokenFields, IssuerUrl, OAuth2TokenResponse, StandardTokenResponse
};
use serde::de::DeserializeOwned;
use std::str::FromStr;
use uuid::Uuid;

/// Configuration for CIBA request parameters
#[derive(Debug, Clone)]
pub struct CibaRequestConfig {
    pub login_hint: LoginHint,
    pub scope: String,
    pub binding_message: String,
    pub resource: Option<String>,
    pub qr_session_id: Option<String>,
}

/**
 * Handle OIDC response and parse it to the specified type or return an appropriate error
 */
async fn handle_oidc_response<T>(response: Result<reqwest::Response, reqwest::Error>) -> Result<T, OidcError> 
where 
    T: DeserializeOwned,
{
    match response {
        Ok(res) => {
            let status = res.status();
            let status_code = status.as_u16();
            
            if status.is_success() {
                match res.json::<T>().await {
                    Ok(parsed_response) => Ok(parsed_response),
                    Err(e) => Err(OidcError::RequestErrorWithDetails {
                        status_code,
                        error: "request_error".to_string(),
                        error_description: format!("Failed to parse response: {}", e),
                    }),
                }
            } else {
                // Get the response body as text
                match res.text().await {
                    Ok(body) => {
                        // Try to parse as JSON
                        match serde_json::from_str::<OidcErrorResponse>(&body) {
                            Ok(error_response) => {
                                if error_response.error == "authorization_pending" {
                                    Err(OidcError::CibaAuthenticationPending)
                                } else {
                                    Err(OidcError::RequestErrorWithDetails {
                                        status_code,
                                        error: error_response.error,
                                        error_description: error_response.error_description,
                                    })
                                }
                            },
                            Err(_) => {
                                Err(OidcError::RequestErrorWithDetails {
                                    status_code,
                                    error: "request_error".to_string(),
                                    error_description: body,
                                })
                            }
                        }
                    },
                    Err(e) => {
                        Err(OidcError::RequestErrorWithDetails {
                            status_code,
                            error: "request_error".to_string(),
                            error_description: format!("Failed to read response body: {}", e),
                        })
                    }
                }
            }
        },
        Err(e) => {
            Err(OidcError::RequestErrorWithDetails {
                status_code: 0, // No status code available for connection errors
                error: "request_error".to_string(),
                error_description: e.to_string(),
            })
        }
    }
}

/**
 * Check if provided token is in JWT format
 * @param token: &String
 */
pub fn is_jwt_token(token: &str) -> bool {
    token.contains(".")
}

/**
 * IDP token introspection
 * for non-JWT tokens; uses introspection endpoint
 */
async fn token_introspection(
    base_url: &str,
    access_token: &AccessToken,
    device_jwt: &String,
    device_client_id: &String,
) -> Result<ClientCredentialsIntrospection, OidcError> {
    let client = default_client()?;
    let params = [
        ("token", access_token.secret()),
        ("client_id", &device_client_id.to_string()),
        (
            "client_assertion_type",
            &"urn:ietf:params:oauth:client-assertion-type:jwt-bearer".to_string(),
        ),
        ("client_assertion", device_jwt),
    ];
    let response = client
        .post(format!("{}/token/introspection", base_url))
        .form(&params)
        .header("content-type", "application/x-www-form-urlencoded")
        .send()
        .await?
        .error_for_status()?
        .json::<ClientCredentialsIntrospection>()
        .await?;

    Ok(response)
}

/**
 * IDP subject identification
 * for non-JWT tokens; uses identify endpoint
 */
async fn identify_subject(
    base_url: &str,
    access_token: &String,
) -> Result<SubjectIdentity, OidcError> {
    let client = default_client()?;
    let response = client
        .get(format!("{}/identify", base_url))
        .header("Authorization", format!("Bearer {}", access_token))
        .send()
        .await?
        .error_for_status()?
        .json::<SubjectIdentity>()
        .await?;

    Ok(response)
}

/**
 * Sub: Make a CIBA request to IDP
 */
async fn ciba_request(
    base_url: &str,
    device_client_id: &ClientId,
    device_jwt: &str,
    device_jwt_ciba: &str,
) -> Result<CibaResponse, OidcError> {
    let client = default_client()?;

    let params = [
        ("client_id", device_client_id.as_str()),
        (
            "client_assertion_type",
            "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
        ),
        ("client_assertion", device_jwt),
        ("request", device_jwt_ciba),
    ];

    let response = client
        .post(format!("{}/backchannel", base_url))
        .form(&params)
        .header("content-type", "application/x-www-form-urlencoded")
        .send()
        .await;
        
    handle_oidc_response::<CibaResponse>(response).await
}

/**
 * Make a CIBA request to IDP
 */
async fn make_ciba_request(
    device_client_id: &ClientId,
    provider_metadata: &CoreProviderMetadata,
    private_key: &EncodingKey,
    device_jwt: &str,
    config: &CibaRequestConfig,
) -> Result<CibaResponse, OidcError> {
    let signed_request = make_device_jwt_ciba(
        device_client_id,
        provider_metadata,
        &config.login_hint,
        &config.scope,
        &config.binding_message,
        config.resource.clone(),
        private_key,
        config.qr_session_id.clone(),
    )?;

    let ciba_response = ciba_request(
        provider_metadata.issuer().as_str(),
        device_client_id,
        device_jwt,
        &signed_request,
    )
    .await?;

    Ok(ciba_response)
}

/**
 * Check CIBA status with IDP
 */
async fn check_ciba_status(
    provider_metadata: &CoreProviderMetadata,
    auth_request_id: &str,
    device_jwt: &str,
    device_client_id: &ClientId,
) -> Result<CibaStatusResponse, OidcError> {
    let client = default_client()?;
    let token_endpoint = provider_metadata
        .token_endpoint()
        .ok_or(OidcRequirementsError::MissingTokenEndpoint)?;

    let params = [
        ("grant_type", "urn:openid:params:grant-type:ciba"),
        ("auth_req_id", auth_request_id),
        ("client_id", device_client_id.as_str()),
        (
            "client_assertion_type",
            "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
        ),
        ("client_assertion", device_jwt),
    ];

    let response = client
        .post(token_endpoint.as_str())
        .form(&params)
        .header("content-type", "application/x-www-form-urlencoded")
        .send()
        .await;
        
    handle_oidc_response::<CibaStatusResponse>(response).await
}

async fn make_qr_auth_session_idp(
    base_url: &str,
    device_access_token: &str,
) -> Result<QrAuthSessionIdp, OidcError> {
    let client = default_client()?;

    let response = client
        .post(format!("{}/qr/make", base_url))
        .bearer_auth(device_access_token)
        .send()
        .await?
        .error_for_status()?
        .json::<QrAuthSessionIdp>()
        .await?;

    Ok(response)
}


async fn verify_qr_auth_session_idp(
    base_url: &str,
    session: QrStatusRequest,
    device_access_token: &str,
) -> Result<Option<QrAuthSessionIdp>, OidcError> {
    let client = default_client()?;
    
    let response = client
        .post(format!("{}/qr/status", base_url))
        .bearer_auth(device_access_token)
        .json(&session)
        .send()
        .await?
        .error_for_status()?
        .json::<Option<QrAuthSessionIdp>>()
        .await?;
    
    Ok(response)
}

/**
 * Validate a JWT token received from IDP
 */
pub fn validate_jwt(
    jwt: &str,
    issuer_jwks: &JwkSet,
    _client_id: &ClientId,
) -> Result<JwtPayload, OidcError> {
    if issuer_jwks.keys.is_empty() {
        return Err(OidcError::InvalidJwkSet);
    }
    
    let decoding_key = jsonwebtoken::DecodingKey::from_jwk(&issuer_jwks.keys[0].clone())
        .map_err(|e| OidcError::InvalidJwk(e.to_string()))?;
    let mut validation = jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::RS256);
    // debug!("Setting audience to Client ID: {}", client_id.to_string());
    // TODO: How'to handle this in tests, with dynamic apps?
    // validation.set_audience(&vec![client_id.to_string()]);
    validation.validate_aud = false;
    debug!("Validating JWT: {}", jwt);
    let token_data = jsonwebtoken::decode::<JwtPayload>(jwt, &decoding_key, &validation)?;

    Ok(token_data.claims)
}

#[derive(Clone)]
pub struct OpenIdconnectClient {
    pub client_id: ClientId,
    pub issuer_url: IssuerUrl,
    pub issuer_jwks: Option<JwkSet>,
    pub provider_metadata: CoreProviderMetadata,
    pub private_key: EncodingKey,
    pub config: OnesOidcConfig,
}

impl OpenIdconnectClient {
    pub fn new(
        client_id: ClientId,
        issuer_url: IssuerUrl,
        provider_metadata: CoreProviderMetadata,
        private_key: EncodingKey,
    ) -> Self {
        OpenIdconnectClient {
            client_id,
            issuer_url,
            issuer_jwks: None,
            provider_metadata,
            private_key,
            config: OnesOidcConfig::default(),
        }
    }

    pub fn with_config(
        client_id: ClientId,
        issuer_url: IssuerUrl,
        provider_metadata: CoreProviderMetadata,
        private_key: EncodingKey,
        config: OnesOidcConfig,
    ) -> Self {
        OpenIdconnectClient {
            client_id,
            issuer_url,
            issuer_jwks: None,
            provider_metadata,
            private_key,
            config,
        }
    }

    pub fn client_id(&self) -> &ClientId {
        &self.client_id
    }

    pub fn issuer_url(&self) -> &IssuerUrl {
        &self.issuer_url
    }

    pub fn provider_metadata(&self) -> &CoreProviderMetadata {
        &self.provider_metadata
    }

    pub fn private_key(&self) -> &EncodingKey {
        &self.private_key
    }

    pub fn device_jwt(&self) -> Result<String, DeviceError> {
        make_device_jwt(&self.client_id, &self.provider_metadata, &self.private_key)
    }

    pub async fn make_ciba_request(
        &self,
        login_hint: &LoginHint,
        scope: &str,
        binding_message: &str,
        resource: Option<String>,
        qr_session_id: Option<String>,
    ) -> Result<CibaResponse, OidcError> {
        let device_jwt = self.device_jwt()?;
        let config = CibaRequestConfig {
            login_hint: login_hint.clone(),
            scope: scope.to_string(),
            binding_message: binding_message.to_string(),
            resource,
            qr_session_id,
        };
        make_ciba_request(
            &self.client_id,
            &self.provider_metadata,
            &self.private_key,
            &device_jwt,
            &config,
        )
        .await
    }

    pub async fn check_ciba_status(
        &self,
        auth_request_id: &str,
    ) -> Result<CibaStatusResponse, OidcError> {
        let device_jwt = self.device_jwt()?;
        check_ciba_status(
            &self.provider_metadata,
            auth_request_id,
            &device_jwt,
            &self.client_id,
        )
        .await
    }

    pub async fn make_qr_auth_session_idp(
        &self,
    ) -> Result<QrAuthSessionIdp, OidcError> {
        let device_access_token_res = &self.device_access_token().await?;
        let device_access_token = device_access_token_res.access_token().secret();
        make_qr_auth_session_idp(&self.issuer_url, device_access_token).await
    }

    // Ones authRequestId is filled, switch to polling CIBA
    pub async fn verify_qr_auth_session_idp(
        &self,
        session: QrStatusRequest,
        // CIBA properties
        scope: &str,
        binding_message: &str,
        resource: Option<String>,
    ) -> Result<Option<QrAuthSessionIdp>, OidcError> {
        let device_access_token_res = &self.device_access_token().await?;
        let device_access_token = device_access_token_res.access_token().secret();
        let result =
            verify_qr_auth_session_idp(&self.issuer_url, session, device_access_token).await?;
        if result.is_none() {
            return Ok(None);
        }

        let result = result.ok_or(OidcError::InvalidQrSession)?;
        if let (Some(login_hint_token), None) = (&result.login_hint_token, &result.auth_request_id) {
            let login_hint = LoginHint {
                kind: LoginHintKind::LoginHintToken,
                value: login_hint_token.clone(),
            };
            let _ = self
                .make_ciba_request(
                    &login_hint,
                    scope,
                    binding_message,
                    resource,
                    Some(result.session_id.clone()),
                )
                .await?;
        }

        Ok(Some(result))
    }

    pub async fn token_introspection(
        &self,
        access_token: &AccessToken,
    ) -> Result<ClientCredentialsIntrospection, OidcError> {
        let device_jwt = self.device_jwt()?;
        token_introspection(
            &self.issuer_url,
            access_token,
            &device_jwt,
            &self.client_id,
        )
        .await
    }

    pub async fn get_jwks(&mut self) -> Result<JwkSet, reqwest::Error> {
        if let Some(jwks) = &self.issuer_jwks {
            return Ok(jwks.clone());
        }

        let jwks = get_jwks(&self.provider_metadata).await?;
        self.issuer_jwks = Some(jwks.clone());
        Ok(jwks)
    }

    pub async fn request_device_access_token(
        &self,
    ) -> Result<
        StandardTokenResponse<EmptyExtraTokenFields, CoreTokenType>,
        DeviceError,
    > {
        let device_jwt = self.device_jwt()?;
        request_device_access_token(&self.provider_metadata, &self.client_id, &device_jwt)
            .await
    }

    pub async fn device_access_token(
        &self,
    ) -> Result<StandardTokenResponse<EmptyExtraTokenFields, CoreTokenType>, DeviceError> {
        let device_jwt = self.device_jwt()?;
        device_access_token(&self.provider_metadata, &self.client_id, &device_jwt).await
    }

    pub fn make_device_jwt(&self) -> Result<String, DeviceError> {
        make_device_jwt(&self.client_id, &self.provider_metadata, &self.private_key)
    }

    pub async fn validate_jwt(&mut self, jwt: &str) -> Result<JwtPayload, OidcError> {
        self.get_jwks().await?;
        validate_jwt(jwt, &self.get_jwks().await?, &self.client_id)
    }

    pub async fn validate_token(
        &mut self,
        token: &String,
        permitted_idp_scopes: Option<Vec<String>>,
    ) -> Result<AuthenticationResult, OidcError> {
        if is_jwt_token(token) {
            let claims = self.validate_jwt(token).await?;

            let aud = if let Some(aud_claim) = &claims.aud {
                Some(Uuid::from_str(aud_claim.as_str())?)
            } else {
                None
            };

            let scope = claims.scope.clone();
            let resource = claims.resource.clone();

            if let (Some(scope_val), Some(resource_val), Some(permitted_scopes)) = (&scope, &resource, &permitted_idp_scopes) {
                if resource_val == "idp-server" && permitted_scopes.contains(scope_val) {
                    return Ok(AuthenticationResult {
                        entity: AuthenticatedEntityKind::Device,
                        iss: claims.iss,
                        sub: claims.sub,
                        aud,
                        scope: Some(scope_val.clone()),
                        username: claims.username,
                        client_id: claims.client_id,
                        method: AuthenticationMethod::IdpJwt,
                        idp_role: claims.idp_role,
                    });
                }
            }

            Ok(AuthenticationResult {
                entity: AuthenticatedEntityKind::User,
                iss: claims.iss,
                sub: claims.sub,
                aud,
                scope,
                username: claims.username,
                client_id: claims.client_id,
                method: AuthenticationMethod::UserJwt,
                idp_role: claims.idp_role,
            })
        } else {
            // This may be a device, mobile, or user token
            let access_token = AccessToken::new(token.to_string());
            let introspection = self.token_introspection(&access_token).await
                .map_err(|e| OidcError::TokenIntrospectionFailed(e.to_string()))?;
                
            if !introspection.active {
                return Err(OidcError::TokenNotActive);
            }

            let iss = introspection.iss.ok_or(OidcError::TokenIntrospectionFailed(
                "Missing issuer in token".to_string()
            ))?;

            let identify = identify_subject(self.issuer_url.as_str(), token).await
                .map_err(|e| OidcError::TokenIdentificationFailed(e.to_string()))?;
            let mut method = AuthenticationMethod::UserDevice;

            if introspection.sub.is_some() && identify.subject_type == AuthenticatedEntityKind::User
            {
                method = AuthenticationMethod::UserIdp;
            } else if identify.subject_type == AuthenticatedEntityKind::User {
                // Mobile introspection does not provide a subject
                method = AuthenticationMethod::UserDevice;
            } else if identify.subject_type == AuthenticatedEntityKind::Device {
                method = AuthenticationMethod::Device;
            }

            Ok(AuthenticationResult {
                entity: identify.subject_type,
                iss,
                sub: identify.subject,
                aud: None,
                scope: introspection.scope,
                username: identify.username,
                client_id: Some(identify.client_id),
                method,
                idp_role: introspection.idp_role,
            })
        }
    }
}