package com.atlassian.asap.nimbus.parser;
   
   import com.atlassian.asap.api.AlgorithmType;
   import com.atlassian.asap.api.JwsHeader;
   import com.atlassian.asap.api.Jwt;
   import com.atlassian.asap.api.JwtBuilder;
   import com.atlassian.asap.api.JwtClaims;
   import com.atlassian.asap.api.SigningAlgorithm;
   import com.atlassian.asap.core.exception.SignatureMismatchException;
  import com.atlassian.asap.core.exception.UnsupportedAlgorithmException;
  import com.atlassian.asap.core.parser.VerifiableJwt;
  import com.google.common.collect.Maps;
  import com.nimbusds.jose.JOSEException;
  import com.nimbusds.jose.JWSObject;
  import com.nimbusds.jose.JWSVerifier;
  import com.nimbusds.jose.crypto.ECDSAVerifier;
  import com.nimbusds.jose.crypto.RSASSAVerifier;
  import com.nimbusds.jwt.JWTClaimsSet;
  import net.minidev.json.JSONObject;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import javax.json.JsonObject;
  import java.security.Provider;
  import java.security.PublicKey;
  import java.security.interfaces.ECPublicKey;
  import java.security.interfaces.RSAPublicKey;
  import java.util.Date;
  import java.util.Locale;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Set;
  
  import static com.google.common.base.Predicates.in;
  import static com.google.common.base.Predicates.not;
  
  public class NimbusVerifiableJwt implements VerifiableJwt {
      private static final Logger logger = LoggerFactory.getLogger(NimbusVerifiableJwt.class);
      private static final Set<String> REGISTERED_CLAIM_NAMES = JWTClaimsSet.getRegisteredNames();
  
      private final Jwt unverifiedJwt;
      private final JWSObject jwsObject;
  
      private final Provider provider;
  
      public NimbusVerifiableJwt(final Jwt unverifiedJwt, final JWSObject jwsObject, final Provider provider) {
          this.unverifiedJwt = unverifiedJwt;
          this.jwsObject = jwsObject;
          this.provider = provider;
      }
  
      /**
       * Factory method to create a signature verifiable jwt.
       *
       * @param jwsObject a json web signature object
       * @param claims    jwt claims set
       * @param provider  Java Security provider
       * @return a signature verifiable jwt
       * @throws UnsupportedAlgorithmException if the signing algorithm is not supported
       */
      public static VerifiableJwt buildVerifiableJwt(final JWSObject jwsObject, final JWTClaimsSet claims, final Provider provider)
              throws UnsupportedAlgorithmException {
          Map<String, Object> customClaimsMap = Maps.filterKeys(claims.getClaims(), not(in(REGISTERED_CLAIM_NAMES)));
          JsonObject customClaims = (JsonObject) NimbusJsr353Translator.nimbusToJsr353(new JSONObject(customClaimsMap));
  
          Jwt unverifiedJwt = JwtBuilder.newJwt()
                  .algorithm(getSigningAlgorithm(jwsObject.getHeader().getAlgorithm().getName()))
                  .keyId(jwsObject.getHeader().getKeyID())
                  .issuer(claims.getIssuer())
                  .jwtId(claims.getJWTID())
                  .subject(Optional.ofNullable(claims.getSubject()))
                  .audience(claims.getAudience())
                  .expirationTime(claims.getExpirationTime().toInstant())
                  .issuedAt(claims.getIssueTime().toInstant())
                  .notBefore(Optional.ofNullable(claims.getNotBeforeTime()).map(Date::toInstant))
                  .customClaims(customClaims)
                  .build();
  
          return new NimbusVerifiableJwt(unverifiedJwt, jwsObject, provider);
      }
  
      @Override
      public void verifySignature(final PublicKey publicKey) throws SignatureMismatchException, UnsupportedAlgorithmException {
          try {
              if (!jwsObject.verify(verifierFor(unverifiedJwt.getHeader().getAlgorithm(), publicKey, provider))) {
                  // log at debug level because this can be caused by invalid input
                  logger.debug("Invalid JWT signature");
                  throw new SignatureMismatchException("Invalid JWT signature");
              }
          } catch (JOSEException e) {
              logger.error("Unexpected error when verifying a JWT signature", e);
              throw new SignatureMismatchException("Unexpected error when verifying JWT signature");
          }
      }
  
      private static JWSVerifier verifierFor(final SigningAlgorithm algorithm, final PublicKey publicKey, final Provider provider)
              throws UnsupportedAlgorithmException {
          if ((algorithm.type() == AlgorithmType.RSA || algorithm.type() == AlgorithmType.RSASSA_PSS) &&
                  publicKey instanceof RSAPublicKey) {
             RSASSAVerifier rsassaVerifier = new RSASSAVerifier((RSAPublicKey) publicKey);
             rsassaVerifier.getJCAContext().setProvider(provider);
             return rsassaVerifier;
         } else if (algorithm.type() == AlgorithmType.ECDSA && publicKey instanceof ECPublicKey) {
             try {
                 ECPublicKey ecPublicKey = (ECPublicKey) publicKey;
                 ECDSAVerifier ecdsaVerifier = new ECDSAVerifier(ecPublicKey);
                 ecdsaVerifier.getJCAContext().setProvider(provider);
                 return ecdsaVerifier;
             } catch (JOSEException e) {
                 // log at debug level since this can be caused by invalid input
                 logger.debug("Unsupported signing algorithm {} or public key algorithm {}", algorithm, publicKey.getAlgorithm());
                 throw new UnsupportedAlgorithmException(algorithm.name(), e);
             }
         } else {
             // log at debug level since this can be caused by invalid input
             logger.debug("Unsupported signing algorithm {} or public key algorithm {}", algorithm, publicKey.getAlgorithm());
             throw new UnsupportedAlgorithmException(algorithm.name());
         }
     }
 
 
     private static SigningAlgorithm getSigningAlgorithm(String algorithm) throws UnsupportedAlgorithmException {
         try {
             return SigningAlgorithm.valueOf(algorithm.toUpperCase(Locale.ROOT));
         } catch (IllegalArgumentException e) {
             throw new UnsupportedAlgorithmException(algorithm + " is not a supported asymmetric JWS algorithm");
         }
     }
 
     @Override
     public JwsHeader getHeader() {
         return unverifiedJwt.getHeader();
     }
 
     @Override
     public JwtClaims getClaims() {
         return unverifiedJwt.getClaims();
     }
 }