ApiResource.java

package dod.p1.keycloak.api;

import dod.p1.keycloak.api.cache.GroupMembershipCache;
import dod.p1.keycloak.api.cache.UsersCache;
import dod.p1.keycloak.api.responses.*;
import dod.p1.keycloak.common.ApiEndpointConfig;
import dod.p1.keycloak.common.CommonConfig;
import org.jboss.logging.Logger;
import org.jboss.resteasy.annotations.cache.NoCache;
import org.keycloak.connections.jpa.JpaConnectionProvider;
import org.keycloak.models.Constants;
import org.keycloak.models.GroupModel;
import org.keycloak.models.KeycloakSession;
import org.keycloak.models.KeycloakSessionFactory;
import org.keycloak.models.KeycloakSessionTask;
import org.keycloak.models.RealmModel;
import org.keycloak.models.jpa.entities.GroupEntity;
import org.keycloak.models.jpa.entities.UserGroupMembershipEntity;
import org.keycloak.models.utils.KeycloakModelUtils;
import org.keycloak.models.utils.ModelToRepresentation;
import org.keycloak.services.managers.AppAuthManager;
import org.keycloak.services.managers.AuthenticationManager;

import javax.persistence.EntityManager;
import javax.persistence.TypedQuery;
import javax.ws.rs.*;
import javax.ws.rs.core.MediaType;
import javax.ws.rs.core.Response;
import java.util.*;
import java.util.stream.Stream;

public class ApiResource {
    private static final Logger LOGGER = Logger.getLogger(ApiResource.class);

    private final AuthenticationManager.AuthResult auth;
    private final Map<String, ApiEndpointConfig> apiEndpoints;
    private final GroupMembershipCache groupMembershipCache;
    private final UsersCache usersCache;
    private final KeycloakSession session;

    public ApiResource(KeycloakSession session, GroupMembershipCache groupMembershipCache, UsersCache usersCache) {
        if (!groupMembershipCache.isInitialized()) {
            LOGGER.info("groupMembershipCache is not initialized, initializing..");
            groupMembershipCache.init(session);
        } else {
            LOGGER.info("groupMembershipCache is initialized.");
        }
        if (!usersCache.isInitialized()) {
            LOGGER.info("usersCache is not initialized, initializing..");
            usersCache.init(session);
        } else {
            LOGGER.info("usersCache is initialized.");
        }
        this.groupMembershipCache = groupMembershipCache;
        this.usersCache = usersCache;
        this.session = session;

        AppAuthManager appAuthManager = new AppAuthManager();
        //bearer token needs to come from a user that exists in the baby-yoda realm for authentication to work
        this.auth = appAuthManager.authenticateBearerToken(session);
        RealmModel realm = session.getContext().getRealm();
        CommonConfig commonConfig = CommonConfig.getInstance(realm);
        this.apiEndpoints = commonConfig.getApiConfigs();
    }

    /**
     * Build a list of group response for the given IL group id.
     * 
     * Use the `first` and `max` parameters to page through results.
     *      * `first` indicates the index of the first group result to return.
     *      * `max` indicates the total "page" size - if `first` is set and `max` is not, a default page size of 100 will be used.
     */
    public static GroupsResponse buildGroupsModelForIlGroup(KeycloakSession session, ApiEndpointConfig endpointConfig,
                                                            GroupMembershipCache groupMembershipCache, UsersCache usersCache, int first, int max) {
        if (endpointConfig == null) {
            throw new IllegalArgumentException("endpoint config cannot be null.");
        }

        // initialize user output list
        Set<UserResponse> userResponses = new HashSet<>();

        // fetch all users for group
        Set<String> ilUserIds = groupMembershipCache.getMemberUserIds(endpointConfig.getAuthorizedGroup().getId());

        Map<String, GroupResponseBuilder> productToGroupInfoMap = new HashMap<>();

        //set a default max if none is provided
        if(first > 0 && max == 0) {
            max = Constants.DEFAULT_MAX_RESULTS; // default = 100
        }

        //get relevent group IDs with a custom query
        String[] releventGroupIds = getProductGroupIds(session, groupMembershipCache, endpointConfig.getEnvironment(), first, max);

        //get the current realm
        RealmModel realm = session.getContext().getRealm();

        //for each product group
        for(String productGroupId : releventGroupIds) {
            //get the group
            GroupModel productGroup = session.realms().getGroupById(productGroupId, realm);
            if(productGroup == null) {
                LOGGER.infof("null GroupModel returned for ID %s - skipping", productGroupId);
                continue;
            }

            // get the group's children
            Set<GroupModel> subGroups = productGroup.getSubGroups();

            // for each product sub-group (should only be `developer` and `collaborator`) - add the appropriate users
            for (GroupModel roleGroup : subGroups) {
                // double check that this is a "role" / leaf node 
                if (!(roleGroup.getName().equals("developer") || roleGroup.getName().equals("collaborator"))) {
                    LOGGER.errorf("found a non-relevent group (%s) which is not a `developer` or `collaborator` - this shouldn't happen", roleGroup.getId());
                    continue;
                }

                // get the mattermost team names for the group from the cache
                String roleGroupPath = ModelToRepresentation.buildGroupPath(roleGroup);
                Set<String> mattermostTeamNames = endpointConfig.getMattermostTeamNames(roleGroupPath);
                if (mattermostTeamNames.isEmpty()) {
                    //LOGGER.infof("Skipping because no mattermost team name: ", groupPath);
                    LOGGER.errorf("found a non-relevent group (%s) with no mattermost teamnames in the cache - this shouldn't happen", roleGroup.getId());
                    continue;
                }

                // build the group path and parse product and role
                String productPath = roleGroupPath.substring(0, roleGroupPath.length() - roleGroup.getName().length() - 1);
                String[] productPathSplits = roleGroupPath.split("/");
                String productName = productPathSplits[productPathSplits.length - 2];
                String role = roleGroup.getName();

                Set<String> memberIds = groupMembershipCache.getMemberUserIds(roleGroup.getId());
                GroupResponseBuilder groupInfo = productToGroupInfoMap.computeIfAbsent(productPath,
                        path -> new GroupResponseBuilder(productName, path, new HashMap<>(), mattermostTeamNames));
                if (memberIds != null) {
                    memberIds.stream().filter(ilUserIds::contains).forEach(u -> {
                        UserResponse user = usersCache.getUserById(u);
                        if (user != null) {
                            Set<String> roles = groupInfo.getUserToRolesMap().computeIfAbsent(user.getUsername(), n -> new HashSet<>());
                            roles.add(role);
                            userResponses.add(user);
                        }
                    });
                }
            }
        }
        LOGGER.info("Done building map for all groups ..");

        List<GroupResponse> groupResponses = new ArrayList<>();
        for(Map.Entry<String, GroupResponseBuilder> entry : productToGroupInfoMap.entrySet()) {
            GroupResponseBuilder groupInfo = entry.getValue();
            Set<GroupMember> groupMembers = new HashSet<>();
            for(Map.Entry<String, Set<String>> entry1 : groupInfo.getUserToRolesMap().entrySet()) {
                groupMembers.add(new GroupMember(entry1.getKey(), entry1.getValue()));
            }
            GroupResponse groupResponse = new GroupResponse(groupInfo.getName(), groupInfo.getPath(), groupMembers, groupInfo.getMattermostTeams());
            groupResponses.add(groupResponse);
        }
        return new GroupsResponse(groupResponses, userResponses);
    }

    // Uses a custom query to find all relevent product group IDs for the specified page.
    // Relevent groups must have a `developer` or `collaborator` child node which belong to the mmGroupId cache (i.e. have a path defined in the config for the environment)
    private static String[] getProductGroupIds(KeycloakSession session,GroupMembershipCache groupCache, String env, int first, int max) {
        KeycloakSessionFactory sessionFactory = session.getKeycloakSessionFactory();
        final List<String> releventGroupIds = new ArrayList<>();

        //get mattermost group IDs and build a SQL formatted list in the format '('group1_ID','group2_ID',...)'
        String[] mmGroupIds = groupCache.getAllMattermostGroups(env);
        if (mmGroupIds.length == 0) {
            return mmGroupIds;
        }
        String mmGroupIdCSV = "'" + String.join("','", mmGroupIds) + "'";

        // build a query to page through "valid" product groups where the group has a `developer` or `collaborator` child group which is listed in mattermost group IDs (i.e. path is defined in the config for the environment)
        String queryString = String.format("SELECT group.id FROM GroupEntity group WHERE group.id IN (%s)",
            String.format("SELECT g.parentId FROM GroupEntity g WHERE g.name IN ('developer', 'collaborator') AND g.id IN (%s)", mmGroupIdCSV));
        
        // make a keycloak transaction
        KeycloakModelUtils.runJobInTransaction(sessionFactory, new KeycloakSessionTask() {
            @Override
            public void run(KeycloakSession session) {
                // query for group IDs for groups with the name `developer` or `collaborator` with mattermost teams
                EntityManager em = session.getProvider(JpaConnectionProvider.class).getEntityManager();
                LOGGER.infof("Querying for product groups.",  mmGroupIds.length);
                try {
                    TypedQuery<String> groupQuery = em.createQuery(queryString, String.class);
                    if(first > 0) {
                        groupQuery.setFirstResult(first);
                    }
                    if(max > 0) {
                        groupQuery.setMaxResults(max);
                    }
                    
                    // add results to the output set and log size
                    groupQuery.getResultStream().forEach(id -> releventGroupIds.add(id));
                } catch (Exception e) {
                    LOGGER.error("error performing product group query: " + e.getMessage());
                }
                LOGGER.infof("found %d relevent groups", releventGroupIds.size());
            }
        });
        return releventGroupIds.toArray(String[]::new);
    }

    /**
     * Validate user is authorized to access the endpoint. Only particular api user is allowed to access each realm and they must authenticate successfully.
     * @param auth the auth result
     * @param allowedUserName the allowed user name to access the endpoint
     */
    public static void validateUserAuthorization(AuthenticationManager.AuthResult auth, String allowedUserName) {
        if (auth == null || !auth.getUser().getUsername().equals(allowedUserName)) {
            Response.status(Response.Status.BAD_REQUEST).build();
            LOGGER.info("Access to groups API is not authorized");
            throw new NotAuthorizedException("Not Authorized");
        }
    }

    @GET
    @NoCache
    @Produces(MediaType.APPLICATION_JSON)
    public GroupsResponse getGroups(
            @QueryParam("env") String env,
            @QueryParam("first") int first,
            @QueryParam("max") int max) {
        // check input params
        if(first < 0) {
            throw new BadRequestException("Query parameter 'first' cannot be less than 0");
        }
        if(max < 0) {
            throw new BadRequestException("Query paramter 'max' cannot be less than 0.");
        }

        LOGGER.info("Processing groups api request for " + env);
        LOGGER.infof("{first: %d, max %d}", first, max);
        ApiEndpointConfig endpointConfig = this.apiEndpoints.get(env);
        if (endpointConfig == null) {
            throw new BadRequestException(env + " is not a valid value for the 'env' query param.");
        }
        validateUserAuthorization(auth, endpointConfig.getAllowedUser());

        GroupsResponse groups = buildGroupsModelForIlGroup(this.session, endpointConfig, this.groupMembershipCache, this.usersCache, first, max);
        LOGGER.info("total users: " + groups.getUsers().size());
        LOGGER.info("total groups: " + groups.getGroups().size());
        LOGGER.info("Done processing groups api request for " + env);
        return groups;
    }
}