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

package org.gcube.oidc.rest;

import java.io.BufferedReader;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
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("");

        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 {

        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();
        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 url, Map<String, List<String>> params,
            String authorization) throws IOException, ProtocolException, UnsupportedEncodingException {

        HttpURLConnection con = (HttpURLConnection) url.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) {
            logger.debug("Adding authorization header as: {}", authorization);
            con.setRequestProperty("Authorization", authorization);
        }
        OutputStream os = con.getOutputStream();

        String queryString = mapToQueryString(params);
        logger.debug("Parameters query string is: {}", queryString);
        os.write(queryString.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) {
            clientId = getClientIdFromToken(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);
    }

    protected static String getClientIdFromToken(JWTToken token) {
        String clientId;
        logger.debug("Client id not provided, using authorized party field (azp)");
        clientId = token.getAzp();
        if (clientId == null) {
            logger.debug("Authorized party field (azp) not present, getting one of the audience field (aud)");
            clientId = getFirstAudienceNoAccount(token);
        }
        return clientId;
    }

    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 = getClientIdFromToken(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.getAccessTokenAsBearer());
        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;
    }

    public static byte[] getUserAvatar(URL avatarURL, JWTToken token) {
        logger.debug("Getting user avatar from URL: {}", avatarURL);
        ByteArrayOutputStream buffer;
        try {
            HttpURLConnection conn = (HttpURLConnection) avatarURL.openConnection();
            conn.setRequestMethod("GET");
            conn.setDoOutput(false);
            conn.setDoInput(true);
            conn.setRequestProperty("Accept", "image/png, image/jpeg, image/gif");
            if (token != null) {
                String authorization = token.getAccessTokenAsBearer();
                logger.debug("Adding authorization header as: {}", authorization);
                conn.setRequestProperty("Authorization", authorization);
            }
            logger.debug("Getting the resource opening the stream");
            InputStream is = conn.getInputStream();
            buffer = new ByteArrayOutputStream();
            int nRead;
            byte[] data = new byte[1024];
            while ((nRead = is.read(data, 0, data.length)) != -1) {
                buffer.write(data, 0, nRead);
            }

            buffer.flush();
            return buffer.toByteArray();
        } catch (IOException e) {
            logger.debug("Can't download avatar", e);
        }
        return null;
    }
}