package com.atlassian.asap.core.server.springsecurity;
   
   import com.atlassian.asap.api.Jwt;
   import com.atlassian.asap.api.JwtBuilder;
   import com.atlassian.asap.api.exception.CannotRetrieveKeyException;
   import com.atlassian.asap.api.exception.InvalidTokenException;
   import com.atlassian.asap.core.exception.PublicKeyNotFoundException;
   import com.atlassian.asap.core.validator.JwtValidator;
   import com.google.common.collect.ImmutableSet;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  import org.springframework.security.authentication.AuthenticationProvider;
  import org.springframework.security.authentication.AuthenticationServiceException;
  import org.springframework.security.authentication.BadCredentialsException;
  import org.springframework.security.core.Authentication;
  import org.springframework.security.core.AuthenticationException;
  import org.springframework.security.core.GrantedAuthority;
  import org.springframework.security.web.authentication.preauth.PreAuthenticatedAuthenticationToken;
  
  import java.util.Collection;
  import java.util.Collections;
  import java.util.Objects;
  import java.util.Set;
  
  /**
   * An authentication provider that validates an {@link UnverifiedBearerToken} and if the token is valid,
   * transforms it into an authenticated token which includes the granted authorities.
   *
   * <p>In order for this authentication provider to be invoked with a {@link UnverifiedBearerToken}, be sure to register a
   * {@link BearerTokenAuthenticationProcessingFilter}, otherwise it will never get asked to authenticate.
   *
   * <p>By default, this class will accept any valid token, which means any ASAP key will result in the request being
   * authenticated. If you only want to allow a subset of issuers/subjects to be authenticated, implement
   * {@link #getGrantedAuthorities} to throw an exception for unknown issuers/subjects.
   *
   * <p>If your application is set up so that it requires authorization (e.g. specific roles) in addition to authentication,
   * you can use {@link IssuerAndSubjectWhitelistAsapAuthenticationProvider} or implement {@link #getGrantedAuthorities}
   * to grant specific authorities to tokens. Otherwise, {@link #defaultAuthorities} are granted to all the tokens.
   */
  public class AsapAuthenticationProvider implements AuthenticationProvider {
      private static final Logger logger = LoggerFactory.getLogger(AsapAuthenticationProvider.class);
  
      private final JwtValidator jwtValidator;
      private final Set<GrantedAuthority> defaultAuthorities;
  
      public AsapAuthenticationProvider(JwtValidator jwtValidator,
                                        Collection<GrantedAuthority> defaultAuthorities) {
          this.jwtValidator = Objects.requireNonNull(jwtValidator);
          this.defaultAuthorities = ImmutableSet.copyOf(defaultAuthorities);
      }
  
      public AsapAuthenticationProvider(JwtValidator jwtValidator) {
          this(jwtValidator, Collections.emptySet());
      }
  
      @Override
      public final Authentication authenticate(Authentication authentication) throws AuthenticationException {
          String serialisedJwt = (String) authentication.getCredentials();
          try {
              Jwt validJwt = jwtValidator.readAndValidate(serialisedJwt);
              Collection<GrantedAuthority> grantedAuthorities = getGrantedAuthorities(validJwt);
              return new PreAuthenticatedAuthenticationToken(
                      effectiveSubject(validJwt),
                      immutableAndSerializable(validJwt),
                      grantedAuthorities);
          } catch (PublicKeyNotFoundException e) {
              logger.debug("Public key not found", e);
              throw new BadCredentialsException("Unable to verify token");
          } catch (CannotRetrieveKeyException e) {
              logger.error("Failed to retrieve public key", e);
              throw new AuthenticationServiceException("Failed to retrieve public key");
          } catch (InvalidTokenException e) {
              logger.debug("Invalid token", e);
              throw new BadCredentialsException("Invalid token");
          }
      }
  
      private static Jwt immutableAndSerializable(Jwt jwt) {
          return JwtBuilder.copyJwt(jwt).build();
      }
  
      /**
       * Subclasses should implement the strategy to grant authorities to valid JWT tokens with the given combination
       * of issuer and effective subject.
       *
       * @return authorities granted to the JWT token
       * @throws AuthenticationException is the token, although valid, is unauthorized
       */
      protected Collection<GrantedAuthority> getGrantedAuthorities(Jwt validJwt) throws AuthenticationException {
          return defaultAuthorities;
      }
  
      /**
       * @param jwt a token
       * @return the effective subject of the token, which may be the explicit 'sub' claim or otherwise the 'iss'
       */
      protected static String effectiveSubject(Jwt jwt) {
          return jwt.getClaims().getSubject().orElse(jwt.getClaims().getIssuer());
      }
 
     @Override
     public final boolean supports(Class<?> authentication) {
         return UnverifiedBearerToken.class.isAssignableFrom(authentication);
     }
 
     protected Collection<GrantedAuthority> getDefaultAuthorities() {
         return defaultAuthorities;
     }
 }