authorization-utils/src/main/java/org/gcube/common/authorization/utils/secret/JWTSecret.java

package org.gcube.common.authorization.utils.secret;

import java.net.URLDecoder;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Base64;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.TimeUnit;

import org.gcube.com.fasterxml.jackson.databind.ObjectMapper;
import org.gcube.common.authorization.library.provider.AccessTokenProvider;
import org.gcube.common.authorization.library.provider.ClientInfo;
import org.gcube.common.authorization.library.provider.ExternalServiceInfo;
import org.gcube.common.authorization.library.provider.UserInfo;
import org.gcube.common.authorization.library.utils.Caller;
import org.gcube.common.authorization.utils.clientid.RenewalProvider;
import org.gcube.common.authorization.utils.user.KeycloakUser;
import org.gcube.common.authorization.utils.user.User;
import org.gcube.common.keycloak.KeycloakClientFactory;
import org.gcube.common.keycloak.model.AccessToken;
import org.gcube.common.keycloak.model.AccessToken.Access;
import org.gcube.common.keycloak.model.ModelUtils;
import org.gcube.common.keycloak.model.RefreshToken;
import org.gcube.common.keycloak.model.TokenResponse;
import org.gcube.common.keycloak.model.util.Time;
import org.gcube.common.scope.impl.ScopeBean;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * @author Luca Frosini (ISTI - CNR)
 */
public class JWTSecret extends Secret {

	private static final Logger logger = LoggerFactory.getLogger(JWTSecret.class);
	
	/**
	 * The interval of time expressed in milliseconds used as guard to refresh the token before that it expires .
	 * TimeUnit has been used to in place of just 
	 * using the number to have a clearer code 
	 */
	public static final long TOLERANCE = TimeUnit.MILLISECONDS.toMillis(200);
	
	protected AccessToken accessToken;
	protected TokenResponse tokenResponse;
	protected RenewalProvider renewalProvider; 
	protected Set<String> roles;
	protected ClientInfo clientInfo;
	protected Caller caller;
	protected String context;
	
	public JWTSecret(String token) {
		super(10, token);
	}
	
	private String getTokenString() {
		try {
			boolean expired = false;
			getAccessToken();
			
			if(Time.currentTimeMillis()>=(accessToken.getExp()-TOLERANCE)) {
				expired = true;
				if(tokenResponse!=null) {
					try {
						KeycloakClientFactory.newInstance().refreshToken(getUsername(), tokenResponse);
						expired = false;
					}catch (Exception e) {
						logger.warn("Unable to refresh the token with RefreshToken. Going to try to renew it if possible.", e);
					}
				}
			}
			
			if(expired && renewalProvider!=null) {
				try {
					JWTSecret renewed = (JWTSecret) renewalProvider.renew(getContext());
					this.token = renewed.token;
					this.accessToken = getAccessToken();
				}catch (Exception e) {
					logger.warn("Unable to renew the token with the RenewalProvider. I'll continue using the old token.", e);
				}
			}
		}catch (Exception e) {
			logger.error("Unexpected error in the procedure to evaluate/refresh the current token. I'll continue using the old token.", e);
		}
		return token;
	}
	
	@Override
	public void setToken() throws Exception {
		AccessTokenProvider.instance.set(getTokenString());
	}

	@Override
	public void resetToken() throws Exception {
		AccessTokenProvider.instance.reset();
	}
	
	protected AccessToken getAccessToken() {
		if(accessToken==null) {
			String realUmaTokenEncoded = token.split("\\.")[1];
			String realUmaToken = new String(Base64.getDecoder().decode(realUmaTokenEncoded.getBytes()));
			ObjectMapper mapper = new ObjectMapper();
			try {
				 accessToken = mapper.readValue(realUmaToken, AccessToken.class);
			}catch(Exception e){
				logger.error("Error parsing JWT token",e);
				throw new RuntimeException("Error parsing JWT token", e);
			}
		}
		return accessToken;
	}
	
	protected Set<String> getRoles() throws Exception{
		if(roles == null) {
			Map<String,Access> accesses = getAccessToken().getResourceAccess();
			String context = getContext(); 
			Access access = accesses.get(URLEncoder.encode(context, StandardCharsets.UTF_8.toString()));
			if(access != null) {
				roles = access.getRoles();
			}else {
				roles = new HashSet<>();
			}
		}
		return roles;
	}

	@Override
	public ClientInfo getClientInfo() throws Exception {
		if(clientInfo==null) {
			User user = getUser();
			if(user.isApplication()) {
				clientInfo = new ExternalServiceInfo(user.getUsername(), "unknown");
			}else {
				clientInfo = new UserInfo(user.getUsername(), new ArrayList<>(user.getRoles()), user.getEmail(), user.getGivenName(), user.getFamilyName());
			}
		}
		return clientInfo;
	}

	@Override
	public Caller getCaller() throws Exception {
		if(caller==null) {
			caller = new Caller(getClientInfo(), "token");
		}
		return caller;
	}
	
	@Override
	public String getContext() throws Exception {
		if(context==null) {
			String[] audience = getAccessToken().getAudience();
			for (String aud : audience) {
				if (aud != null && aud.compareTo("") != 0) {
					try {
						String contextToBeValidated = URLDecoder.decode(aud, StandardCharsets.UTF_8.toString());
						ScopeBean scopeBean = new ScopeBean(contextToBeValidated);
						context = scopeBean.toString();
						return context;
					} catch (Exception e) {
						logger.error("Invalid context name for audience {} in access token. Trying next one if any.", aud, e);
					}
				}
			}
			throw new Exception("Invalid context in access token");
		}
		return context;
	}

	@Override
	public String getUsername() throws Exception {
		return getAccessToken().getPreferredUsername();
	}
	
	@Override
	public Map<String, String> getHTTPAuthorizationHeaders() {
		Map<String, String> authorizationHeaders = new HashMap<>();
		authorizationHeaders.put("Authorization", "Bearer " + getTokenString());
		return authorizationHeaders;
	}

	public void setRenewalProvider(RenewalProvider renewalProvider) {
		this.renewalProvider = renewalProvider;
	}

	public void setTokenResponse(TokenResponse tokenResponse) {
		this.tokenResponse = tokenResponse;
	}
	
	protected boolean isExpired(AccessToken accessToken) {
		return Time.currentTimeMillis()>accessToken.getExp();
	}
	
	@Override
	public boolean isExpired() {
		return isExpired(getAccessToken());
	}

	@Override
	public boolean isRefreshable() {
		if(tokenResponse!=null) {
			try {
				RefreshToken  refreshToken = ModelUtils.getRefreshTokenFrom(tokenResponse);
				return isExpired(refreshToken);
			} catch (Exception e) {
				return false;
			}
		}
		return false;
	}
	
	@Override
	public User getUser() {
		if(user==null) {
			try {
				ObjectMapper objectMapper = new ObjectMapper();
				String accessTokenString = objectMapper.writeValueAsString(getAccessToken());
				user = objectMapper.readValue(accessTokenString, KeycloakUser.class);
				user.setRoles(getRoles());
			} catch (Exception e) {
				throw new RuntimeException();
			}
		}
		return user;
	}
	
}