package com.atlassian.asap.nimbus.serializer;
   
   import com.atlassian.asap.api.AlgorithmType;
   import com.atlassian.asap.api.Jwt;
   import com.atlassian.asap.api.JwtClaims;
   import com.atlassian.asap.api.SigningAlgorithm;
   import com.atlassian.asap.core.SecurityProvider;
   import com.atlassian.asap.core.exception.SigningException;
   import com.atlassian.asap.core.exception.UnsupportedAlgorithmException;
  import com.atlassian.asap.core.serializer.JwtSerializer;
  import com.google.common.annotations.VisibleForTesting;
  import com.nimbusds.jose.JOSEException;
  import com.nimbusds.jose.JWSAlgorithm;
  import com.nimbusds.jose.JWSHeader;
  import com.nimbusds.jose.JWSObject;
  import com.nimbusds.jose.JWSSigner;
  import com.nimbusds.jose.Payload;
  import com.nimbusds.jose.crypto.ECDSASigner;
  import com.nimbusds.jose.crypto.RSASSASigner;
  import net.minidev.json.JSONObject;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.security.PrivateKey;
  import java.security.Provider;
  import java.security.interfaces.ECPrivateKey;
  import java.security.interfaces.RSAPrivateKey;
  
  /**
   * A serializer of JWT implemented using the Nimbus library.
   */
  public class NimbusJwtSerializer implements JwtSerializer {
      private static final Logger logger = LoggerFactory.getLogger(NimbusJwtSerializer.class);
  
      private final Provider provider;
  
      public NimbusJwtSerializer() {
          this(SecurityProvider.getProvider());
      }
  
      public NimbusJwtSerializer(Provider provider) {
          this.provider = provider;
      }
  
      @Override
      public String serialize(Jwt jwt, PrivateKey privateKey) throws SigningException, UnsupportedAlgorithmException {
          JWSObject jwsObject = getSignedJwsObject(jwt, privateKey);
  
          return jwsObject.serialize();
      }
  
      @VisibleForTesting
      JWSObject getSignedJwsObject(Jwt jwt, PrivateKey privateKey) throws UnsupportedAlgorithmException {
          SigningAlgorithm algorithm = jwt.getHeader().getAlgorithm();
          JWSHeader header = new JWSHeader.Builder(JWSAlgorithm.parse(algorithm.name())) // fails if algorithm is None
                  .keyID(jwt.getHeader().getKeyId())
                  .build();
          Payload payload = new Payload(toJsonPayload(jwt.getClaims()));
          JWSObject jwsObject = new JWSObject(header, payload);
          try {
              jwsObject.sign(getSigner(algorithm, privateKey));
          } catch (JOSEException e) {
              logger.error("Unexpected error when signing JWT token", e);
              throw new SigningException();
          }
          return jwsObject;
      }
  
      private JWSSigner getSigner(SigningAlgorithm algorithm, PrivateKey privateKey) throws UnsupportedAlgorithmException {
          if ((algorithm.type() == AlgorithmType.RSA || algorithm.type() == AlgorithmType.RSASSA_PSS) && privateKey instanceof RSAPrivateKey) {
              return createRSASSASignerForKey((RSAPrivateKey) privateKey);
          } else if (algorithm.type() == AlgorithmType.ECDSA && privateKey instanceof ECPrivateKey) {
              return createECDSASignerForKey((ECPrivateKey) privateKey);
          } else {
              throw new UnsupportedAlgorithmException(String.format("Unsupported algorithm %s or signing key type", algorithm.name()));
          }
      }
  
      @VisibleForTesting
      protected JWSSigner createRSASSASignerForKey(RSAPrivateKey privateKey) {
          final RSASSASigner rsassaSigner = new RSASSASigner(privateKey);
          rsassaSigner.getJCAContext().setProvider(provider);
          return rsassaSigner;
      }
  
      @VisibleForTesting
      protected JWSSigner createECDSASignerForKey(ECPrivateKey privateKey) throws UnsupportedAlgorithmException {
          try {
              final ECDSASigner ecdsaSigner = new ECDSASigner(privateKey);
              ecdsaSigner.getJCAContext().setProvider(provider);
              return ecdsaSigner;
          } catch (JOSEException e) {
              throw new UnsupportedAlgorithmException(String.format("Unsupported algorithm %s or signing key type", privateKey.getAlgorithm()));
          }
      }
  
      private static JSONObject toJsonPayload(JwtClaims claims) {
          return (JSONObject) Jsr353NimbusTranslator.jsr353ToNimbus(claims.getJson());
      }
 }