// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

package com.azure.spring.cloud.autoconfigure.aad.filter;

import com.azure.spring.cloud.autoconfigure.aad.implementation.constants.AadJwtClaimNames;
import com.azure.spring.cloud.autoconfigure.aad.properties.AadAuthenticationProperties;
import com.azure.spring.cloud.autoconfigure.aad.properties.AadAuthorizationServerEndpoints;
import com.nimbusds.jose.JOSEException;
import com.nimbusds.jose.JWSAlgorithm;
import com.nimbusds.jose.JWSObject;
import com.nimbusds.jose.jwk.source.JWKSetCache;
import com.nimbusds.jose.jwk.source.JWKSource;
import com.nimbusds.jose.jwk.source.RemoteJWKSet;
import com.nimbusds.jose.proc.BadJOSEException;
import com.nimbusds.jose.proc.JWSKeySelector;
import com.nimbusds.jose.proc.JWSVerificationKeySelector;
import com.nimbusds.jose.proc.SecurityContext;
import com.nimbusds.jose.util.ResourceRetriever;
import com.nimbusds.jwt.JWT;
import com.nimbusds.jwt.JWTClaimsSet;
import com.nimbusds.jwt.JWTParser;
import com.nimbusds.jwt.proc.BadJWTException;
import com.nimbusds.jwt.proc.ConfigurableJWTProcessor;
import com.nimbusds.jwt.proc.DefaultJWTClaimsVerifier;
import com.nimbusds.jwt.proc.DefaultJWTProcessor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.MalformedURLException;
import java.net.URL;
import java.text.ParseException;
import java.util.Collections;
import java.util.HashSet;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

/**
 * A user principal manager to load user info from JWT.
 */
public class UserPrincipalManager {

    private static final Logger LOGGER = LoggerFactory.getLogger(UserPrincipalManager.class);
    private static final String LOGIN_MICROSOFT_ONLINE_ISSUER = "https://login.microsoftonline.com/";
    private static final String STS_WINDOWS_ISSUER = "https://sts.windows.net/";
    private static final String STS_CHINA_CLOUD_API_ISSUER = "https://sts.chinacloudapi.cn/";

    private static final String MSG_MALFORMED_AD_KEY_DISCOVERY_URI = "Failed to parse active directory key discovery uri.";

    private final JWKSource<SecurityContext> keySource;
    private final AadAuthenticationProperties aadAuthenticationProperties;
    private final boolean explicitAudienceCheck;
    private final Set<String> validAudiences = new HashSet<>();

    /**
     * ø Creates a new {@link UserPrincipalManager} with a predefined {@link JWKSource}.
     * <p>
     * This is helpful in cases the JWK is not a remote JWKSet or for unit testing.
     *
     * @param keySource - {@link JWKSource} containing at least one key
     */
    public UserPrincipalManager(JWKSource<SecurityContext> keySource) {
        this.keySource = keySource;
        this.explicitAudienceCheck = false;
        this.aadAuthenticationProperties = null;
    }

    /**
     * Create a new {@link UserPrincipalManager} based of the
     * {@link AadAuthorizationServerEndpoints#getJwkSetEndpoint()}
     *
     * @param endpoints - used to retrieve the JWKS URL
     * @param aadAuthenticationProperties - used to retrieve the environment.
     * @param resourceRetriever - configures the {@link RemoteJWKSet} call.
     * @param explicitAudienceCheck Whether explicitly check the audience.
     * @throws IllegalArgumentException If AAD key discovery URI is malformed.
     */
    public UserPrincipalManager(AadAuthorizationServerEndpoints endpoints,
                                AadAuthenticationProperties aadAuthenticationProperties,
                                ResourceRetriever resourceRetriever,
                                boolean explicitAudienceCheck) {
        this.aadAuthenticationProperties = aadAuthenticationProperties;
        this.explicitAudienceCheck = explicitAudienceCheck;
        if (explicitAudienceCheck) {
            // client-id for "normal" check
            this.validAudiences.add(this.aadAuthenticationProperties.getCredential().getClientId());
            // app id uri for client credentials flow (server to server communication)
            this.validAudiences.add(this.aadAuthenticationProperties.getAppIdUri());
        }
        try {
            String jwkSetEndpoint =
                endpoints.getJwkSetEndpoint();
            keySource = new RemoteJWKSet<>(new URL(jwkSetEndpoint), resourceRetriever);
        } catch (MalformedURLException e) {
            throw new IllegalArgumentException(MSG_MALFORMED_AD_KEY_DISCOVERY_URI, e);
        }
    }

    /**
     * Create a new {@link UserPrincipalManager} based of the
     * {@link AadAuthorizationServerEndpoints#getJwkSetEndpoint()}
     * ()}
     *
     * @param endpoints - used to retrieve the JWKS URL
     * @param aadAuthenticationProperties - used to retrieve the environment.
     * @param resourceRetriever - configures the {@link RemoteJWKSet} call.
     * @param jwkSetCache - used to cache the JWK set for a finite time, default set to 5 minutes which matches
     * constructor above if no jwkSetCache is passed in
     * @param explicitAudienceCheck Whether explicitly check the audience.
     * @throws IllegalArgumentException If AAD key discovery URI is malformed.
     */
    public UserPrincipalManager(AadAuthorizationServerEndpoints endpoints,
                                AadAuthenticationProperties aadAuthenticationProperties,
                                ResourceRetriever resourceRetriever,
                                boolean explicitAudienceCheck,
                                JWKSetCache jwkSetCache) {
        this.aadAuthenticationProperties = aadAuthenticationProperties;
        this.explicitAudienceCheck = explicitAudienceCheck;
        if (explicitAudienceCheck) {
            // client-id for "normal" check
            this.validAudiences.add(this.aadAuthenticationProperties.getCredential().getClientId());
            // app id uri for client credentials flow (server to server communication)
            this.validAudiences.add(this.aadAuthenticationProperties.getAppIdUri());
        }
        try {
            String jwkSetEndpoint = endpoints.getJwkSetEndpoint();
            keySource = new RemoteJWKSet<>(new URL(jwkSetEndpoint), resourceRetriever, jwkSetCache);
        } catch (MalformedURLException e) {
            throw new IllegalArgumentException(MSG_MALFORMED_AD_KEY_DISCOVERY_URI, e);
        }
    }

    /**
     * Parse the id token to {@link UserPrincipal}.
     *
     * @param aadIssuedBearerToken The token issued by AAD.
     * @return The parsed {@link UserPrincipal}.
     * @throws ParseException If the token couldn't be parsed to a valid JWS object.
     * @throws JOSEException If an internal processing exception is encountered.
     * @throws BadJOSEException If the JWT is rejected.
     */
    public UserPrincipal buildUserPrincipal(String aadIssuedBearerToken) throws ParseException, JOSEException,
        BadJOSEException {
        final JWSObject jwsObject = JWSObject.parse(aadIssuedBearerToken);
        final ConfigurableJWTProcessor<SecurityContext> validator = getValidator(jwsObject.getHeader().getAlgorithm());
        final JWTClaimsSet jwtClaimsSet = validator.process(aadIssuedBearerToken, null);
        validator.getJWTClaimsSetVerifier().verify(jwtClaimsSet, null);
        UserPrincipal userPrincipal = new UserPrincipal(aadIssuedBearerToken, jwsObject, jwtClaimsSet);
        userPrincipal.setRoles(getRoles(jwtClaimsSet));
        return userPrincipal;
    }

    Set<String> getRoles(JWTClaimsSet set) {
        if (set == null) {
            return Collections.emptySet();
        }
        Object rolesClaim = set.getClaim(AadJwtClaimNames.ROLES);
        if (rolesClaim == null) {
            return Collections.emptySet();
        }
        if (rolesClaim instanceof Iterable<?>) {
            return StreamSupport.stream(((Iterable<?>) rolesClaim).spliterator(), false)
                    .map(Object::toString)
                    .collect(Collectors.toSet());
        }
        return Stream.of(rolesClaim)
                .map(Object::toString)
                .collect(Collectors.toSet());
    }

    /**
     * Whether the token was issued by AAD.
     *
     * @param token the token
     * @return whether the token was issued by AAD
     */
    public boolean isTokenIssuedByAad(String token) {
        try {
            final JWT jwt = JWTParser.parse(token);
            return isAadIssuer(jwt.getJWTClaimsSet().getIssuer());
        } catch (ParseException e) {
            LOGGER.info("Fail to parse JWT {}, exception {}", token, e);
        }
        return false;
    }

    private static boolean isAadIssuer(String issuer) {
        if (issuer == null) {
            return false;
        }
        return issuer.startsWith(LOGIN_MICROSOFT_ONLINE_ISSUER)
            || issuer.startsWith(STS_WINDOWS_ISSUER)
            || issuer.startsWith(STS_CHINA_CLOUD_API_ISSUER);
    }

    private ConfigurableJWTProcessor<SecurityContext> getValidator(JWSAlgorithm jwsAlgorithm) {
        final ConfigurableJWTProcessor<SecurityContext> jwtProcessor = new DefaultJWTProcessor<>();
        final JWSKeySelector<SecurityContext> keySelector = new JWSVerificationKeySelector<>(jwsAlgorithm, keySource);
        jwtProcessor.setJWSKeySelector(keySelector);
        //TODO: would it make sense to inject it? and make it configurable or even allow to provide own implementation
        jwtProcessor.setJWTClaimsSetVerifier(new DefaultJWTClaimsVerifier<SecurityContext>() {
            @Override
            public void verify(JWTClaimsSet claimsSet, SecurityContext ctx) throws BadJWTException {
                super.verify(claimsSet, ctx);
                final String issuer = claimsSet.getIssuer();
                if (!isAadIssuer(issuer)) {
                    throw new BadJWTException("Invalid token issuer");
                }
                if (explicitAudienceCheck) {
                    Optional<String> matchedAudience = claimsSet.getAudience()
                                                                .stream()
                                                                .filter(validAudiences::contains)
                                                                .findFirst();
                    if (matchedAudience.isPresent()) {
                        LOGGER.debug("Matched audience: [{}]", matchedAudience.get());
                    } else {
                        throw new BadJWTException("Invalid token audience. Provided value " + claimsSet.getAudience()
                            + "does not match neither client-id nor AppIdUri.");
                    }
                }
            }
        });
        return jwtProcessor;
    }
}
