oidc-library/src/main/java/org/gcube/oidc/rest/OpenIdConnectRESTHelper.java

package org.gcube.oidc.rest;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.io.UnsupportedEncodingException;
import java.net.HttpURLConnection;
import java.net.ProtocolException;
import java.net.URL;
import java.net.URLEncoder;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class OpenIdConnectRESTHelper {

    protected static final Logger logger = LoggerFactory.getLogger(OpenIdConnectRESTHelper.class);

    public static String buildLoginRequestURL(URL loginURL, String clientId, String state, String redirectURI)
            throws UnsupportedEncodingException {

        Map<String, List<String>> params = new HashMap<String, List<String>>();
        params.put("client_id", Arrays.asList(URLEncoder.encode(clientId, "UTF-8")));
        params.put("response_type", Arrays.asList("code"));
        params.put("scope", Arrays.asList("openid"));
        params.put("state", Arrays.asList(URLEncoder.encode(state, "UTF-8")));
        params.put("redirect_uri", Arrays.asList(URLEncoder.encode(redirectURI, "UTF-8")));
        params.put("login", Arrays.asList("true"));
        String q = mapToQueryString(params);
        return loginURL + "?" + q;
    }

    public static String mapToQueryString(Map<String, List<String>> params) {
        String q = params.entrySet().stream().flatMap(p -> p.getValue().stream().map(v -> p.getKey() + "=" + v))
                .reduce((p1, p2) -> p1 + "&" + p2).orElse("");

        if (logger.isDebugEnabled()) {
            logger.debug("Query string is: " + q);
        }
        return q;
    }

    public static JWTToken queryClientToken(String clientId, String clientSecret, URL tokenURL) throws Exception {
        Map<String, List<String>> params = new HashMap<>();
        params.put("grant_type", Arrays.asList("client_credentials"));
        params.put("client_id", Arrays.asList(URLEncoder.encode(clientId, "UTF-8")));
        params.put("client_secret", Arrays.asList(URLEncoder.encode(clientSecret, "UTF-8")));
        return performQueryTokenWithPOST(tokenURL, null, params);
    }

    public static JWTToken queryToken(String clientId, URL tokenURL, String code, String scope,
            String redirectURI) throws Exception {

        Map<String, List<String>> params = new HashMap<>();
        params.put("client_id", Arrays.asList(URLEncoder.encode(clientId, "UTF-8")));
        params.put("grant_type", Arrays.asList("authorization_code"));
        params.put("scope", Arrays.asList(URLEncoder.encode(scope, "UTF-8")));
        params.put("code", Arrays.asList(URLEncoder.encode(code, "UTF-8")));
        params.put("redirect_uri", Arrays.asList(URLEncoder.encode(redirectURI, "UTF-8")));
        return performQueryTokenWithPOST(tokenURL, null, params);
    }

    public static JWTToken performQueryTokenWithPOST(URL tokenURL, String authorization,
            Map<String, List<String>> params)
            throws Exception {

        if (logger.isDebugEnabled()) {
            logger.debug("Querying access token from OIDC server with URL: " + tokenURL);
        }
        HttpURLConnection httpURLConnection = performURLEncodedPOSTSendData(tokenURL, params, authorization);

        StringBuilder sb = new StringBuilder();
        int httpResultCode = httpURLConnection.getResponseCode();
        if (logger.isTraceEnabled()) {
            logger.trace("HTTP Response code: " + httpResultCode);
        }
        if (httpResultCode != HttpURLConnection.HTTP_OK) {
            BufferedReader br = new BufferedReader(new InputStreamReader(httpURLConnection.getErrorStream(), "UTF-8"));
            String line = null;
            while ((line = br.readLine()) != null) {
                sb.append(line + "\n");
            }
            br.close();
            throw new Exception("Unable to get token " + sb);
        } else {
            BufferedReader br = new BufferedReader(new InputStreamReader(httpURLConnection.getInputStream(), "UTF-8"));
            String line = null;
            while ((line = br.readLine()) != null) {
                sb.append(line + "\n");
            }
            br.close();
        }
        return JWTToken.fromString(sb.toString());
    }

    protected static HttpURLConnection performURLEncodedPOSTSendData(URL tokenURL, Map<String, List<String>> params,
            String authorization) throws IOException, ProtocolException, UnsupportedEncodingException {

        HttpURLConnection con = (HttpURLConnection) tokenURL.openConnection();
        con.setRequestMethod("POST");
        con.setDoOutput(true);
        con.setDoInput(true);
        con.setRequestProperty("Content-Type", "application/x-www-form-urlencoded");
        con.setRequestProperty("Accept", "application/json");
        if (authorization != null) {
            if (logger.isDebugEnabled()) {
                logger.debug("Adding authorization header as: " + authorization);
            }
            con.setRequestProperty("Authorization", authorization);
        }
        OutputStream os = con.getOutputStream();
        if (logger.isDebugEnabled()) {
            logger.debug("Sending parameters: " + params);
        }
        os.write(mapToQueryString(params).getBytes("UTF-8"));
        os.close();
        return con;
    }

    public static JWTToken queryUMAToken(URL tokenUrl, String authorizationToken, String audience,
            List<String> permissions) throws Exception {

        Map<String, List<String>> params = new HashMap<>();
        params.put("grant_type", Arrays.asList("urn:ietf:params:oauth:grant-type:uma-ticket"));
        params.put("audience", Arrays.asList(URLEncoder.encode(audience, "UTF-8")));
        if (permissions != null && !permissions.isEmpty()) {
            params.put(
                    "permission", permissions.stream().map(s -> {
                        try {
                            return URLEncoder.encode(s, "UTF-8");
                        } catch (UnsupportedEncodingException e) {
                            return "";
                        }
                    }).collect(Collectors.toList()));
        }
        return performQueryTokenWithPOST(tokenUrl, authorizationToken, params);
    }

    public static JWTToken refreshToken(URL tokenURL, JWTToken token) throws Exception {
        return refreshToken(tokenURL, null, null, token);
    }

    public static JWTToken refreshToken(URL tokenURL, String clientId, JWTToken token) throws Exception {
        return refreshToken(tokenURL, clientId, null, token);
    }

    public static JWTToken refreshToken(URL tokenURL, String clientId, String clientSecret, JWTToken token)
            throws Exception {

        Map<String, List<String>> params = new HashMap<>();
        params.put("grant_type", Arrays.asList("refresh_token"));
        if (clientId == null) {
            if (logger.isDebugEnabled()) {
                logger.debug("Client id not provided, using authorized party field (azp)");
            }
            clientId = token.getAzp();
            if (clientId == null) {
                if (logger.isDebugEnabled()) {
                    logger.debug("Authorized party field (azp) not present, getting one of the audience field (aud)");
                }
                clientId = getFirstAudienceNoAccount(token);
            }
        }
        params.put("client_id", Arrays.asList(URLEncoder.encode(clientId, "UTF-8")));
        if (clientSecret != null) {
            params.put("client_secret", Arrays.asList(URLEncoder.encode(clientSecret, "UTF-8")));
        }
        params.put("refresh_token", Arrays.asList(token.getRefreshTokenString()));
        return performQueryTokenWithPOST(tokenURL, null, params);
    }

    private static String getFirstAudienceNoAccount(JWTToken token) {
        // Trying to get it from the token's audience ('aud' field), getting the first except the 'account'
        List<String> tokenAud = token.getAud();
        tokenAud.remove(JWTToken.ACCOUNT_RESOURCE);
        if (tokenAud.size() > 0) {
            return tokenAud.iterator().next();
        } else {
            // Setting it to empty string to avoid NPE in encoding
            return "";
        }
    }

    public static boolean logout(URL logoutUrl, JWTToken token) throws IOException {
        return logout(logoutUrl, null, token);
    }

    public static boolean logout(URL logoutUrl, String clientId, JWTToken token) throws IOException {
        Map<String, List<String>> params = new HashMap<>();
        if (clientId == null) {
            clientId = getFirstAudienceNoAccount(token);
        }
        params.put("client_id", Arrays.asList(URLEncoder.encode(clientId, "UTF-8")));
        params.put("refresh_token", Arrays.asList(token.getRefreshTokenString()));
        logger.info("Performing logut from OIDC server with URL: " + logoutUrl);
        HttpURLConnection httpURLConnection = performURLEncodedPOSTSendData(logoutUrl, params, token.getAsBearer());
        int responseCode = httpURLConnection.getResponseCode();
        if (responseCode == 204) {
            logger.info("Logout performed correctly");
            return true;
        } else {
            logger.error("Cannot perfrom logout: [" + responseCode + "] " + httpURLConnection.getResponseMessage());
        }
        return false;
    }
}