package com.atlassian.asap.core.validator;
   
   import com.atlassian.asap.api.Jwt;
   import com.atlassian.asap.api.exception.CannotRetrieveKeyException;
   import com.atlassian.asap.api.exception.InvalidTokenException;
   import com.atlassian.asap.core.keys.KeyProvider;
   import com.atlassian.asap.core.parser.JwtParser;
   import com.atlassian.asap.core.parser.VerifiableJwt;
   import com.atlassian.asap.core.server.AuthenticationContext;
  import com.atlassian.asap.nimbus.parser.NimbusJwtParser;
  
  import java.security.PublicKey;
  import java.time.Clock;
  import java.util.Collections;
  import java.util.Objects;
  import java.util.Optional;
  import java.util.Set;
  import java.util.function.Supplier;
  
  /**
   * Orchestrates the parsing of the JWT using the provided {@link JwtParser}, the verification of its signature using
   * the provided {@link KeyProvider} and the validation of its claims using the
   * provided {@link JwtClaimsValidator}.
   */
  public class JwtValidatorImpl implements JwtValidator {
      private final KeyProvider<PublicKey> publicKeyProvider;
      private final JwtParser jwtParser;
      private final JwtClaimsValidator jwtClaimsValidator;
      private final Supplier<Set<String>> resourceServerAudiences;
  
      /**
       * Create a new instance of {@link JwtValidatorImpl}.
       *
       * @param publicKeyProvider       the key provider to use for retrieving public keys for signature verification
       * @param jwtParser               the parser to use for parsing a serialized jwt string
       * @param claimValidator          the validator to use for verifying the claims set contained in a JWT
       * @param resourceServerAudiences all JWT messages will need to have one of these audiences to be valid.
       *                                the supplier is called every time a JWT token is validated
       */
      public JwtValidatorImpl(KeyProvider<PublicKey> publicKeyProvider,
                              JwtParser jwtParser,
                              JwtClaimsValidator claimValidator,
                              Supplier<Set<String>> resourceServerAudiences) {
          this.publicKeyProvider = Objects.requireNonNull(publicKeyProvider);
          this.jwtParser = Objects.requireNonNull(jwtParser);
          this.jwtClaimsValidator = Objects.requireNonNull(claimValidator);
          this.resourceServerAudiences = Objects.requireNonNull(resourceServerAudiences);
      }
  
  
      /**
       * Create a new instance of {@link JwtValidatorImpl}.
       *
       * @param publicKeyProvider       the key provider to use for retrieving public keys for signature verification
       * @param jwtParser               the parser to use for parsing a serialized jwt string
       * @param claimValidator          the validator to use for verifying the claims set contained in a JWT
       * @param resourceServerAudiences all JWT messages will need to have one of these audiences to be valid
       */
      public JwtValidatorImpl(KeyProvider<PublicKey> publicKeyProvider,
                              JwtParser jwtParser,
                              JwtClaimsValidator claimValidator,
                              Set<String> resourceServerAudiences) {
          this(publicKeyProvider, jwtParser, claimValidator, () -> resourceServerAudiences);
      }
  
      /**
       * Create a new instance of {@link JwtValidatorImpl}.
       *
       * @param publicKeyProvider      the key provider to use for retrieving public keys for signature verification
       * @param jwtParser              the parser to use for parsing a serialized jwt string
       * @param claimValidator         the validator to use for verifying the claims set contained in a JWT
       * @param resourceServerAudience all JWT messages will need to have this audience to be valid
       */
      public JwtValidatorImpl(KeyProvider<PublicKey> publicKeyProvider,
                              JwtParser jwtParser,
                              JwtClaimsValidator claimValidator,
                              String resourceServerAudience) {
          this(publicKeyProvider, jwtParser, claimValidator, Collections.singleton(resourceServerAudience));
      }
  
      @Override
      public final Jwt readAndValidate(String serializedJwt)
              throws InvalidTokenException, CannotRetrieveKeyException {
          // Parse the serialized JWT without performing any verification of the signature or claims
          VerifiableJwt verifiableJwt = jwtParser.parse(serializedJwt);
  
          // Ensure we have a valid key id before the public key is fetched
          ValidatedKeyId validatedKeyId = ValidatedKeyId.validate(verifiableJwt.getHeader().getKeyId());
  
          // fetch the public key and verify the signature of the JWT
          PublicKey publicKey = publicKeyProvider.getKey(validatedKeyId);
          verifiableJwt.verifySignature(publicKey);
  
          // validate the claims of the JWT, and return it if validation is successful
          jwtClaimsValidator.validate(verifiableJwt, resourceServerAudiences.get());
  
          return verifiableJwt;
      }
  
     @Override
     public final Optional<String> determineUnverifiedIssuer(String serializedJwt) {
         return jwtParser.determineUnverifiedIssuer(serializedJwt);
     }
 
     /**
      * A factory method that returns an instance with a typical configuration.
      *
      * @param authenticationContext context of the authentication
      * @return a request authenticator for the given context
      */
     public static JwtValidator createDefault(AuthenticationContext authenticationContext) {
         return new JwtValidatorImpl(authenticationContext.getPublicKeyProvider(),
                 new NimbusJwtParser(),
                 new JwtClaimsValidator(Clock.systemUTC()),
                 authenticationContext.getResourceServerAudiences());
     }
 
     /**
      * A factory method that returns an instance with a typical configuration.
      *
      * @param audience             the audience this filter will accept requests for
      * @param publicKeyRepoBaseUrl the base URL of the public key repository
      * @return a web filter
      */
     public static JwtValidator createDefault(String audience, String publicKeyRepoBaseUrl) {
         AuthenticationContext authContext = new AuthenticationContext(audience, publicKeyRepoBaseUrl);
         return createDefault(authContext);
     }
 
     /**
      * A factory method that returns an instance with a typical configuration.
      *
      * @param audiences            the audiences this filter will accept requests for
      * @param publicKeyRepoBaseUrl the base URL of the public key repository
      * @return a web filter
      */
     public static JwtValidator createDefault(Set<String> audiences, String publicKeyRepoBaseUrl) {
         AuthenticationContext authContext = new AuthenticationContext(audiences, publicKeyRepoBaseUrl);
         return createDefault(authContext);
     }
 }