Microsoft Azure Active Directory custom group authentication

package info.magnolia.sso.auth.azure;

import info.magnolia.objectfactory.Components;
import info.magnolia.sso.config.Pac4jConfigProvider;
import net.minidev.json.JSONArray;
import net.minidev.json.JSONObject;
import net.minidev.json.JSONValue;
import org.pac4j.core.authorization.generator.AuthorizationGenerator;
import org.pac4j.core.config.Config;
import org.pac4j.core.context.HttpConstants;
import org.pac4j.core.context.WebContext;
import org.pac4j.core.context.session.SessionStore;
import org.pac4j.core.profile.UserProfile;
import org.pac4j.core.util.HttpUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.HttpURLConnection;
import java.net.URL;
import java.util.*;

import org.apache.commons.lang3.StringUtils;

import org.pac4j.oidc.profile.OidcProfileDefinition;
import com.nimbusds.oauth2.sdk.token.BearerAccessToken;

import static java.util.function.Predicate.not;
import static java.util.stream.Collectors.toList;

public class CustomAzureAdAuthorizationGenerator implements AuthorizationGenerator {

    private static final Logger log = LoggerFactory.getLogger(CustomAzureAdAuthorizationGenerator.class);

    // if the user groups property is delivered with the token payload, it usually contains the UIDs only, not the names
    private static final String GROUPS_PROPERTY = "groups";
    // used to resolve groups in SSOAuthenticationModule
    private static final String MGNL_GROUPS_ATTR_NAME = "mgnlGroups";

    // property for resolving group memberships when there is no group or role attribute in the token payload
    private static final String AZURE_USER_GROUPS_SERVICE_URL = "https://graph.microsoft.com/v1.0/me/memberOf";
    // properties for resolving group UIDs
    // https://graph.microsoft.com/v1.0/<tenant>/groups/<group-uid>
    private static final String AZURE_GROUP_ID_PREFIX = "https://graph.microsoft.com/v1.0/";
    private static final String AZURE_GROUP_ID_SUFFIX = "/groups/";
    // common properties
    private final static String MS_GRAPH_VALUE_PROPERTY = "value";
    private final static String MS_GRAPH_GROUP_NAME_KEY = "displayName";

    // TODO Replace with your own tenant ID!
    private final static String tenant = "f7c33569-d9fg-87e6-a2af-3e4feq02310c"; (1)

    @Override
    public Optional<UserProfile> generate(WebContext context, SessionStore sessionStore, UserProfile profile) {
        log.debug("CustomAzureAdAuthorizationGenerator - generate ...");

        List<String> userGroups = new ArrayList<>();
        try {
            Optional<Object> groupsAttribute = Optional.ofNullable(profile.getAttribute(GROUPS_PROPERTY));

            // take the access token from the logged-in user
            // it is needed to query MS Graph APIs
            BearerAccessToken bearerAccessToken = ((BearerAccessToken) profile.getAttribute(OidcProfileDefinition.ACCESS_TOKEN));
            String bearerToken = StringUtils.defaultString(bearerAccessToken.getValue(), "");
            // in production, do not show this in the logs
            log.debug("Bearer token: {}", bearerToken);

            if (groupsAttribute.isEmpty()) {
                // there is no roles and no groups element in the token payload, just general user data
                // we will query MS graph for the user's group memberships
                log.debug("There is no groups/roles claim in the token payload. Group memberships will be queried from MS Graph.");
                JSONObject queryResult = queryMsGraph(bearerToken, AZURE_USER_GROUPS_SERVICE_URL);
                log.debug("Query result: {}", queryResult);

                // parse the result and add group names to the user's group list
                userGroups = getAzureUserGroups(queryResult);
            } else {
                // the groups claim is present in the token payload, but it does only contain a list of the group UIDs, not the names
                log.debug("The groups claim with group UIDs is present in the token payload. Group names will be queried from MS Graph.");

                Object groupsAttributeValue = groupsAttribute.get();
                log.debug("Attempting to gather authorization with {} property value: {}", GROUPS_PROPERTY, groupsAttributeValue);

                List<String> groupIDs = new ArrayList<>();

                if (groupsAttributeValue instanceof ArrayList) {
                    // parse a json array using their parent type ArrayList
                    List<Object> jsonArray = (ArrayList) groupsAttributeValue;
                    groupIDs.addAll(jsonArray.stream()
                            .map(Object::toString)
                            .filter(not(String::isEmpty))
                            .collect(toList()));
                } else if (groupsAttributeValue instanceof String) {
                    // parse a comma-separated list of groups?
                    groupIDs.addAll(Arrays.stream(((String) groupsAttributeValue).split(","))
                            .filter(not(String::isEmpty))
                            .collect(toList()));
                } else {
                    log.error("Could not parse groups from user profile: unexpected \"{}\" attribute type: {} ({})", GROUPS_PROPERTY, groupsAttributeValue, groupsAttributeValue.getClass().getName());
                    return Optional.of(profile);
                }
                log.debug("Group UIDs: {}", groupIDs);
                getAzureAdGroupNamesFromId(bearerToken, tenant, groupIDs, userGroups);
            }

            log.debug("Resulting group list for the logged in user: {}", userGroups);

            if (!userGroups.isEmpty()) {
                profile.addAttribute(MGNL_GROUPS_ATTR_NAME, userGroups);
            }

        } catch (Exception e) {
            log.info("An attempt to retrieve roles or groups for the logged in user failed.", e);
        }

        log.debug("User profile after assigning group membership: {}", profile);
        return Optional.of(profile);
    }

    /**
     * Parse group display names from the MS graph query result.
     *
     * @param jsonObject Query result as JSON object.
     * @return List with user's group names or empty list.
     */
    protected static List<String> getAzureUserGroups(JSONObject jsonObject) {
        List<String> userGroupList = new ArrayList<>();

        Optional<Object> groupArray = Optional.ofNullable(jsonObject.get(MS_GRAPH_VALUE_PROPERTY));
        if (groupArray.isEmpty()) {
            log.error("Could not find the value where the groups are contained in the JSON object.");
        } else {
            JSONArray jsonArray = (JSONArray) groupArray.get();
            jsonArray.forEach(element -> {
                JSONObject obj = (JSONObject) element;
                String displayName = StringUtils.defaultString((String) obj.get(MS_GRAPH_GROUP_NAME_KEY), "");
                if (StringUtils.isNotBlank(displayName)) {
                    log.debug("Fetched group with name {}.", displayName);
                    if (StringUtils.isNotBlank(displayName)) {
                        userGroupList.add(displayName);
                    }
                }
            });
        }
        return userGroupList;
    }

    /**
     * Resolve group names from a list with group UIDs retrieved in the token payload.
     *
     * @param bearerToken Access token.
     * @param tenant      Azure AD tenant provided in config.yaml.
     * @param groupIdList Groups attribute from the token payload containing group UIDs.
     * @param groupList   Resolved group names are added to the resulting group list.
     */
    protected static void getAzureAdGroupNamesFromId(String bearerToken, String tenant, List<String> groupIdList, List<String> groupList) {
        String msGraphUrl = AZURE_GROUP_ID_PREFIX + tenant + AZURE_GROUP_ID_SUFFIX;

        for (String groupId : groupIdList) {
            log.debug("Trying to resolve group with ID {}.", groupId);
            // GET https://graph.microsoft.com/v1.0/<tenant>/groups/<group-uid>
            String queryUrl = msGraphUrl + groupId;
            log.debug("Querying resource {}", queryUrl);
            JSONObject queryResult = queryMsGraph(bearerToken, queryUrl);
            log.debug("Query result: {}", queryResult);

            String displayName = StringUtils.defaultString((String) queryResult.get(MS_GRAPH_GROUP_NAME_KEY), "");
            if (StringUtils.isNotBlank(displayName)) {
                log.debug("Fetched group with name {}.", displayName);
                if (StringUtils.isNotBlank(displayName)) {
                    groupList.add(displayName);
                }
            }
        }
    }

    /**
     * Query MS Graph URL using the access token we retrieved after a successful login.
     *
     * @param bearerToken Access token.
     * @param msGraphUrl  URL to query.
     * @return The response retrieved from Azure as JSON object.
     */
    private static JSONObject queryMsGraph(String bearerToken, String msGraphUrl) {
        JSONObject jsonObjectResult = new JSONObject();

        HttpURLConnection connection = null;
        final Map<String, String> headers = new HashMap<>();
        headers.put(HttpConstants.ACCEPT_HEADER, HttpConstants.APPLICATION_JSON);
        headers.put(HttpConstants.AUTHORIZATION_HEADER, HttpConstants.BEARER_HEADER_PREFIX + bearerToken);

        try {
            connection = openConnection(new URL(msGraphUrl), HttpConstants.HTTP_METHOD.GET.name(), headers);
            final int responseCode = connection.getResponseCode();
            log.debug("Response Code received: {}", responseCode);
            String body = HttpUtils.readBody(connection);
            log.debug("Response body: {}", body);
            log.debug(body);

            jsonObjectResult = (JSONObject) JSONValue.parse(body);

            return jsonObjectResult;
        } catch (Exception e) {
            log.error("Problem while querying MS Graph: ", e);
        } finally {
            HttpUtils.closeConnection(connection);
        }

        return jsonObjectResult;
    }

    // method taken from the PAC4J HttpUtils class (has protected access)
    protected static HttpURLConnection openConnection(final URL url, final String requestMethod, final Map<String, String> headers) throws IOException {
        final HttpURLConnection connection = (HttpURLConnection) url.openConnection();
        connection.setDoInput(true);
        connection.setDoOutput(true);
        connection.setRequestMethod(requestMethod);
        connection.setConnectTimeout(HttpConstants.DEFAULT_CONNECT_TIMEOUT);
        connection.setReadTimeout(HttpConstants.DEFAULT_READ_TIMEOUT);
        if (headers != null) {
            for (final Map.Entry<String, String> entry : headers.entrySet()) {
                connection.setRequestProperty(entry.getKey(), entry.getValue());
            }
        }
        return connection;
    }
}