user-management/src/main/java/com/cleverthis/authservice/service/impl/KeycloakAdminClient.java

package com.cleverthis.authservice.service.impl;

import com.cleverthis.authservice.config.KeycloakClientConfigProperties;
import com.cleverthis.authservice.dto.KeycloakAdminTokenResponse;
import com.cleverthis.authservice.dto.KeycloakClientRepresentation;
import com.cleverthis.authservice.dto.KeycloakRoleRepresentation;
import com.cleverthis.authservice.exception.ServiceUnavailableException;
import java.net.URI;
import java.time.Duration;
import java.util.List;
import java.util.concurrent.atomic.AtomicReference;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.util.UriComponentsBuilder;
import reactor.core.publisher.Mono;
import reactor.util.retry.Retry;

/**
 * Service to interact with Keycloak Admin API. Handles obtaining an admin token for this
 * auth-service and making admin calls.
 */
@Service
@Slf4j
public class KeycloakAdminClient {

    private final WebClient webClient;
    private final AtomicReference<String> adminAccessToken = new AtomicReference<>();
    private final AtomicReference<Long> tokenExpiryTime = new AtomicReference<>(0L);

    private final String keycloakAdminHost;
    private final String realm;
    private final String adminAccountUsername;
    private final String adminAccountPassword;

    /**
     * Create a keycloak admin client.
     */
    public KeycloakAdminClient(
          @Qualifier("keycloakWebClient") final WebClient webClient,
          final KeycloakClientConfigProperties configProperties
    ) {
        this.webClient = webClient;
        this.keycloakAdminHost = configProperties.getKeycloakAdminHost();
        this.realm = configProperties.getRealm();
        this.adminAccountUsername = configProperties.getAdminAccountUsername();
        this.adminAccountPassword = configProperties.getAdminAccountPassword();
    }

    private String getAdminTokenEndpoint() {
        return UriComponentsBuilder.fromUriString(this.keycloakAdminHost)
              .path("/realms/{realm}/protocol/openid-connect/token")
              .buildAndExpand(this.realm)
              .encode().toUriString();
    }

    private String getClientsEndpoint() {
        return UriComponentsBuilder.fromUriString(this.keycloakAdminHost)
              .path("/admin/realms/{realm}/clients")
              .buildAndExpand(this.realm)
              .encode().toUriString();
    }

    private String getUserClientRolesEndpoint(String userId, String clientInternalId) {
        return UriComponentsBuilder.fromUriString(this.keycloakAdminHost)
              .path("/admin/realms/{realm}/users/{userId}"
                    + "/role-mappings/clients/{clientId}/composite")
              .buildAndExpand(this.realm, this.realm, userId, clientInternalId)
              .encode().toUriString();
    }

    /**
     * Retrieves a valid admin access token for this service, refreshing if necessary. Uses client
     * credentials grant.
     *
     * @return Mono emitting the admin access token.
     */
    Mono<String> getAdminAccessToken() {
        if (System.currentTimeMillis() < this.tokenExpiryTime.get()
                && this.adminAccessToken.get() != null) {
            log.debug("Using cached admin access token.");
            return Mono.just(this.adminAccessToken.get());
        }

        log.info("Fetching new admin access token for auth-service using account: {}",
              this.adminAccountUsername);
        MultiValueMap<String, String> formData = new LinkedMultiValueMap<>();
        formData.add("client_id", "admin-cli");
        formData.add("username", this.adminAccountUsername);
        formData.add("password", this.adminAccountPassword);
        formData.add("grant_type", "password");

        return this.webClient.post().uri(getAdminTokenEndpoint())
                .contentType(MediaType.APPLICATION_FORM_URLENCODED)
                .body(BodyInserters.fromFormData(formData)).retrieve()
                .onStatus(status -> status.is4xxClientError() || status.is5xxServerError(),
                        response -> response.bodyToMono(String.class).flatMap(errorBody -> {
                            log.error("Failed to get admin token. Status: {}, Body: {}",
                                    response.statusCode(), errorBody);
                            return Mono.error(new ServiceUnavailableException(
                                    "Could not obtain admin token from Keycloak. Status: "
                                            + response.statusCode()));
                        }))
                .bodyToMono(KeycloakAdminTokenResponse.class).map(tokenResponse -> {
                    this.adminAccessToken.set(tokenResponse.getAccessToken());
                    // Set expiry time slightly before actual expiry to be safe (e.g., 30 seconds
                    // buffer)
                    this.tokenExpiryTime.set(System.currentTimeMillis()
                            + (tokenResponse.getExpiresIn() - 30) * 1000L);
                    log.info("Successfully obtained new admin access token.");
                    return tokenResponse.getAccessToken();
                })
                .retryWhen(Retry.backoff(3, Duration.ofSeconds(2))
                        .maxBackoff(Duration.ofSeconds(10))
                        .filter(throwable -> throwable instanceof ServiceUnavailableException)
                        .doBeforeRetry(
                                retrySignal -> log.warn("Retrying admin token fetch attempt #{}",
                                        retrySignal.totalRetries() + 1)));
    }

    /**
     * Fetches the internal UUID of a Keycloak client given its public clientId.
     *
     * @param publicClientId The human-readable client_id (e.g., "cleverswarm-client").
     * @return Mono emitting the internal UUID string.
     */
    public Mono<String> getClientInternalId(String publicClientId) {
        log.debug("Fetching internal ID for client: {}", publicClientId);

        // We expect only one client with this ID
        URI targetUri = UriComponentsBuilder.fromUriString(getClientsEndpoint())
                .queryParam("clientId", publicClientId).queryParam("max", 1).build().toUri();

        log.debug("Constructed URI for getClientInternalId: {}", targetUri.toString());

        return getAdminAccessToken().flatMap(token -> this.webClient.get().uri(targetUri)
                .header(HttpHeaders.AUTHORIZATION, "Bearer " + token).retrieve()
                .onStatus(status -> status.is4xxClientError() || status.is5xxServerError(),
                        response -> response.bodyToMono(String.class).flatMap(errorBody -> {
                            log.error("Failed to get client '{}'. Status: {}, Body: {}",
                                    publicClientId, response.statusCode(), errorBody);
                            return Mono.error(new ServiceUnavailableException(
                                    "Could not fetch client details from Keycloak. Client: "
                                            + publicClientId));
                        }))
                .bodyToFlux(KeycloakClientRepresentation.class)// Expecting a list, even if one item
                .next() // Take the first element if present
                .map(KeycloakClientRepresentation::getId)
                .switchIfEmpty(Mono.error(new ServiceUnavailableException(
                        "Client not found in Keycloak: " + publicClientId)))
                .doOnSuccess(id -> log.debug("Found internal ID '{}' for client '{}'", id,
                        publicClientId)));
    }

    /**
     * Fetches the effective client roles for a given user and a specific client (internal ID).
     *
     * @param userId           the Keycloak user's 'sub' (subject ID).
     * @param clientInternalId The internal UUID of the Keycloak client.
     * @return Mono emitting a List of KeycloakRoleRepresentation.
     */
    public Mono<List<KeycloakRoleRepresentation>> getUserEffectiveClientRoles(
          String userId, String clientInternalId
    ) {
        log.debug("Fetching effective client roles for user '{}' on client internal ID '{}'",
              userId, clientInternalId);
        return getAdminAccessToken().flatMap(token -> this.webClient.get()
                .uri(getUserClientRolesEndpoint(userId, clientInternalId))
                .header(HttpHeaders.AUTHORIZATION, "Bearer " + token).retrieve()
                .onStatus(status -> status.is4xxClientError() || status.is5xxServerError(),
                        response -> response.bodyToMono(String.class).flatMap(errorBody -> {
                            log.error(
                                    "Failed to get user client roles. User:"
                                    + " {}, ClientInternalId: {}, Status: {}, Body: {}",
                                    userId, clientInternalId, response.statusCode(), errorBody);
                            return Mono.error(new ServiceUnavailableException(
                                    "Could not fetch user client roles from Keycloak."));
                        }))
                .bodyToFlux(KeycloakRoleRepresentation.class).collectList()
                .doOnSuccess(roles -> log.debug(
                        "Found {} effective client roles for user '{}' on client internal ID '{}'",
                        roles.size(), userId, clientInternalId)));
    }
}