package com.atlassian.asap.core.validator;
   
   import com.atlassian.asap.api.Jwt;
   import com.atlassian.asap.api.JwtClaims;
   import com.atlassian.asap.api.exception.InvalidTokenException;
   import com.atlassian.asap.core.exception.InvalidClaimException;
   import com.atlassian.asap.core.exception.TokenExpiredException;
   import com.atlassian.asap.core.exception.TokenTooEarlyException;
   import org.apache.commons.lang3.StringUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.time.Clock;
  import java.time.Duration;
  import java.time.Instant;
  import java.util.Collections;
  import java.util.Objects;
  import java.util.Optional;
  import java.util.Set;
  
  
  /**
   * Validates the claims contained in a JWT.
   */
  public class JwtClaimsValidator {
      /**
       * The JWT spec says that implementers "MAY provide for some small leeway, usually no more than a few minutes, to account for clock skew".
       * Calculations of the current time for the purposes of accepting or rejecting time-based claims (e.g. "exp" and "nbf") will allow for the current time
       * being plus or minus this leeway, resulting in some time-based claims that are marginally before or after the current time being accepted instead of rejected.
       */
      public static final Duration TIME_CLAIM_LEEWAY = Duration.ofSeconds(Long.parseLong(
              System.getProperty("asap.resource.server.leeway.seconds", "30")));
  
      /**
       * No matter what the claims say, the server should reject tokens that are too long-lived.
       */
      public static final Duration DEFAULT_MAX_LIFETIME = Duration.ofHours(1);
  
      private static final Logger logger = LoggerFactory.getLogger(JwtClaimsValidator.class);
  
      private final Clock clock;
      private final Duration maxLifetime;
  
      public JwtClaimsValidator(Clock clock) {
          this(clock, DEFAULT_MAX_LIFETIME);
      }
  
      public JwtClaimsValidator(Clock clock, Duration maxTokenLifetime) {
          this.clock = Objects.requireNonNull(clock);
          this.maxLifetime = Objects.requireNonNull(maxTokenLifetime);
      }
  
      /**
       * Checks the validity of the claims contained in a JWT in a given authentication context.
       *
       * @param jwt                     a JWT token
       * @param resourceServerAudiences the JWT token must be addressed to one of these audiences
       * @throws InvalidTokenException if the claims are invalid or could not be verified
       */
      public void validate(Jwt jwt,
                           Set<String> resourceServerAudiences)
              throws InvalidTokenException {
          final JwtClaims claims = jwt.getClaims();
          final String issuer = claims.getIssuer();
          final Set<String> audience = claims.getAudience();
          final Instant issuedAt = claims.getIssuedAt();
          final Instant expiry = claims.getExpiry();
          final Optional<Instant> mayBeNotBefore = claims.getNotBefore();
          final String keyId = jwt.getHeader().getKeyId();
  
          issuerValidation(issuer, keyId);
          audienceValidation(audience, resourceServerAudiences);
          formalTimeClaimsValidation(issuedAt, expiry, mayBeNotBefore);
          relativeTimeValidation(issuedAt, expiry, mayBeNotBefore);
          // This implementation does not currently validate that the JWT id has not been seen before
      }
  
      private void issuerValidation(String issuer, String keyId)
              throws InvalidClaimException {
          // Since we use the issuer as a prefix, it cannot be blank
          if (StringUtils.isBlank(issuer)) {
              logger.debug("Rejecting blank issuer");
              throw new InvalidClaimException(JwtClaims.RegisteredClaim.ISSUER, "Issuer cannot be blank");
          }
  
          // The "iss" claim value must be the prefix of the key id. This is required to ensure that the issuer
          // owns the key used to sign the token. In the future we may have other mechanisms to ensure that and
          // we may be able to remove this restriction.
          if (!keyId.startsWith(issuer + "/")) {
              logger.debug("The issuer {} does not match the key id {}", issuer, keyId);
              throw new InvalidClaimException(JwtClaims.RegisteredClaim.ISSUER, "The issuer claim does not match the key id");
          }
      }
  
      private void audienceValidation(Set<String> audience,
                                      Set<String> resourceServerAudiences) throws InvalidClaimException {
          // The audience must contain one of the audiences in this resource server
          if (Collections.disjoint(audience, resourceServerAudiences)) {
              logger.debug("Rejected unrecognised audience {}", audience);
             throw new InvalidClaimException(JwtClaims.RegisteredClaim.AUDIENCE, "Unrecognised audience");
         }
     }
 
     /**
      * Formal validation of the time claims that are not relative to the current time.
      */
     private void formalTimeClaimsValidation(Instant issuedAt, Instant expiry, Optional<Instant> mayBeNotBefore)
             throws InvalidClaimException {
         // The token must have been issued before its expiry
         if (expiry.isBefore(issuedAt)) {
             logger.debug("Expiry time {} set before issue time {}", expiry, issuedAt);
             throw new InvalidClaimException(JwtClaims.RegisteredClaim.ISSUED_AT, "Expiry time set before issue time");
         }
 
         // The lifetime of the tokens is limited
         if (issuedAt.plus(maxLifetime).isBefore(expiry)) {
             logger.debug("Token exceeds lifetime limit, issued at {} and expires at {}", issuedAt, expiry);
             throw new InvalidClaimException(JwtClaims.RegisteredClaim.EXPIRY, "Token exceeds lifetime limit");
         }
 
         if (mayBeNotBefore.isPresent()) {
             Instant nbf = mayBeNotBefore.get();
 
             // Sanity check: the token must be valid at some point in time
             if (nbf.isAfter(expiry)) {
                 logger.debug("The expiry time {} must be after the not-before time {}", expiry, nbf);
                 throw new InvalidClaimException(JwtClaims.RegisteredClaim.NOT_BEFORE, "The expiry time must be after the not-before time");
             }
 
             // Reject tokens that were valid strictly before they were issued
             if (nbf.isBefore(issuedAt)) {
                 logger.debug("The token was valid since {} but was issued later at {}", nbf, issuedAt);
                 throw new InvalidClaimException(JwtClaims.RegisteredClaim.NOT_BEFORE, "The token must not be valid before it was issued");
             }
         }
     }
 
     /**
      * Validations that are relative to the current time.
      */
     private void relativeTimeValidation(Instant issuedAt, Instant expiry, Optional<Instant> mayBeNotBefore)
             throws TokenExpiredException, TokenTooEarlyException {
         Instant now = Instant.now(clock);
 
         // Calculate leeways for the jwt processing
         final Instant nowMinusLeeway = now.minus(TIME_CLAIM_LEEWAY);
         final Instant nowPlusLeeway = now.plus(TIME_CLAIM_LEEWAY);
 
         // The token must not have expired (with some tolerance)
         if (expiry.isBefore(nowMinusLeeway)) {
             logger.info("Rejecting expired token, now={}, expiry={}, leeway={}", now, expiry, TIME_CLAIM_LEEWAY);
             throw new TokenExpiredException(expiry, now);
         }
 
         // If "nbf" claim is optional, and it defaults to the "iat" value if absent.
         Instant effectiveNbf = mayBeNotBefore.orElse(issuedAt);
         if (effectiveNbf.isAfter(nowPlusLeeway)) {
             logger.info("Rejecting token that arrives too early, now={}, not before={}, leeway={}", now, effectiveNbf, TIME_CLAIM_LEEWAY);
             throw new TokenTooEarlyException(effectiveNbf, now);
         }
     }
 }