package com.atlassian.asap.nimbus.parser;
   
   import com.atlassian.asap.api.JwsHeader;
   import com.atlassian.asap.api.Jwt;
   import com.atlassian.asap.api.JwtBuilder;
   import com.atlassian.asap.api.SigningAlgorithm;
   import com.atlassian.asap.core.SecurityProvider;
   import com.atlassian.asap.core.exception.SignatureMismatchException;
   import com.atlassian.asap.core.exception.UnsupportedAlgorithmException;
  import com.atlassian.asap.core.parser.VerifiableJwt;
  import com.atlassian.asap.nimbus.serializer.NimbusJwtSerializer;
  import com.nimbusds.jose.JWSObject;
  import com.nimbusds.jose.crypto.RSASSAVerifier;
  import com.nimbusds.jwt.JWTClaimsSet;
  import org.junit.Before;
  import org.junit.Test;
  import org.junit.runner.RunWith;
  import org.mockito.Mock;
  import org.mockito.runners.MockitoJUnitRunner;
  
  import java.security.KeyPair;
  import java.security.KeyPairGenerator;
  import java.security.Provider;
  import java.security.interfaces.RSAPublicKey;
  import java.time.Instant;
  import java.time.temporal.ChronoUnit;
  import java.util.Collections;
  import java.util.Optional;
  
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.fail;
  import static org.mockito.Matchers.any;
  import static org.mockito.Mockito.verify;
  import static org.mockito.Mockito.when;
  
  @RunWith(MockitoJUnitRunner.class)
  public class NimbusVerifiableJwtTest {
      private static final Provider PROVIDER = SecurityProvider.getProvider();
  
      @Mock
      private Jwt jwt;
      @Mock
      private JWSObject jwsObject;
      @Mock
      private RSAPublicKey rsaPublicKey;
  
      @Mock
      private JwsHeader jwsHeader;
  
      @Before
      public void setUpMocks() {
          when(jwt.getHeader()).thenReturn(jwsHeader);
      }
  
      @Test
      public void shouldVerifyValidRS256Signature() throws Exception {
          when(jwsHeader.getAlgorithm()).thenReturn(SigningAlgorithm.RS256);
          when(jwsObject.verify(any(RSASSAVerifier.class))).thenReturn(true); // valid signature
  
          new NimbusVerifiableJwt(jwt, jwsObject, PROVIDER).verifySignature(rsaPublicKey); // should not throw
  
          verify(jwsObject).verify(any(RSASSAVerifier.class));
      }
  
      @Test
      public void shouldThrowIfRS256SignatureIsInvalid() throws Exception {
          when(jwsHeader.getAlgorithm()).thenReturn(SigningAlgorithm.RS256);
          when(jwsObject.verify(any(RSASSAVerifier.class))).thenReturn(false); // invalid signature
  
          try {
              new NimbusVerifiableJwt(jwt, jwsObject, PROVIDER).verifySignature(rsaPublicKey);
              fail("Should have thrown");
          } catch (SignatureMismatchException ex) {
              verify(jwsObject).verify(any(RSASSAVerifier.class));
          }
      }
  
      @Test(expected = UnsupportedAlgorithmException.class)
      public void shouldThrowIfAlgorithmIsUnsupported() throws Exception {
          when(jwsHeader.getAlgorithm()).thenReturn(SigningAlgorithm.ES256);
  
          new NimbusVerifiableJwt(jwt, jwsObject, PROVIDER).verifySignature(rsaPublicKey);
      }
  
      @Test
      public void shouldReturnValidVerifiableJwt() throws Exception {
          // 1) Generate a test serialized JWT to parse
          KeyPair keyPair = KeyPairGenerator.getInstance("RSA").generateKeyPair();
  
          Instant now = Instant.now();
  
          Jwt unverifiedJwt = JwtBuilder.newJwt()
                  .algorithm(SigningAlgorithm.RS256)
                  .keyId("my_key")
                  .issuer("my_issuer")
                  .jwtId("my_jwtId")
                  .subject(Optional.of("my_subject"))
                  .audience("my_audience")
                  .expirationTime(now.plusSeconds(60))
                 .issuedAt(now)
                 .notBefore(Optional.of(now))
                 .build();
 
         String serializedJwt = new NimbusJwtSerializer().serialize(unverifiedJwt, keyPair.getPrivate());
         JWSObject jwsObject = JWSObject.parse(serializedJwt);
         JWTClaimsSet jwtClaimsSet = JWTClaimsSet.parse(jwsObject.getPayload().toJSONObject());
 
         // 2) Parse the test jwt and assert it's as expected
         VerifiableJwt result = NimbusVerifiableJwt.buildVerifiableJwt(jwsObject, jwtClaimsSet, PROVIDER);
 
         assertEquals(SigningAlgorithm.RS256, result.getHeader().getAlgorithm());
         assertEquals("my_key", result.getHeader().getKeyId());
         assertEquals("my_issuer", result.getClaims().getIssuer());
         assertEquals("my_jwtId", result.getClaims().getJwtId());
         assertEquals(Optional.of("my_subject"), result.getClaims().getSubject());
         assertEquals(Collections.singleton("my_audience"), result.getClaims().getAudience());
         assertEquals(now.plusSeconds(60).truncatedTo(ChronoUnit.SECONDS), result.getClaims().getExpiry());
         assertEquals(now.truncatedTo(ChronoUnit.SECONDS), result.getClaims().getIssuedAt());
         assertEquals(Optional.of(now.truncatedTo(ChronoUnit.SECONDS)), result.getClaims().getNotBefore());
     }
 
 
 }