
// 1. HTTPS
import java.net.URL;
import java.io.*;
import javax.net.ssl.HttpsURLConnection;

// 2. json deser
//import org.codehaus.jackson.map.ObjectMapper;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.JsonNodeFactory;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonAutoDetect;


// 3, extract PublicKey
import java.util.Base64;
import java.io.ByteArrayInputStream;
import java.security.GeneralSecurityException; 
import java.security.PublicKey; 
import java.security.Signature; 
import java.security.cert.CertificateFactory; 
import java.security.cert.X509Certificate; 

// 4, validate token
import java.security.spec.InvalidKeySpecException;
import java.security.NoSuchAlgorithmException;
import java.security.Key;
import java.security.PublicKey;
import java.security.interfaces.RSAPublicKey;
import io.jsonwebtoken.Header;
import io.jsonwebtoken.Claims;
import io.jsonwebtoken.Jwt;
import io.jsonwebtoken.Jws;
import io.jsonwebtoken.JwsHeader;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.jackson.io.JacksonDeserializer;
import io.jsonwebtoken.SigningKeyResolverAdapter;
import io.jsonwebtoken.security.Jwk;
import io.jsonwebtoken.security.Jwks;
// only dbg: when keys taken from file, not URL
import java.nio.file.Files;
import java.nio.file.Paths;

import java.util.logging.Logger;

public class IamSigningKeyResolver extends SigningKeyResolverAdapter
{
   private static final Logger LOGGER = Logger.getLogger(IamSigningKeyResolver.class.getName());
   private String keysURL;


   public IamSigningKeyResolver(String keysUrl) {this.keysURL = keysUrl;}

   @Override
   public Key resolveSigningKey(JwsHeader jwsHeader, Claims claims)
   {
      LOGGER.info( "IamSigningKeyResolver::resolveSigningKey" );

      //inspect the header or claims, lookup and return the signing key

      String keyId = jwsHeader.getKeyId(); //or any other field that you need to inspect

      Key key = null;
      try
      {
         key = lookupVerificationKey(keyId); //implement me
      }
      catch(Exception e)
      {
         e.printStackTrace();
      }

      return key;
   }



   private Key lookupVerificationKey(String keyId)
         throws Exception, GeneralSecurityException
      {
         LOGGER.info( "IamSigningKeyResolver::lookupVerificationKey" );

         String jsonKeys = doHttps();

         PublicKey pubKey = (PublicKey)getKeyFromKeys(jsonKeys, keyId);

         return pubKey;
      }


   private String doHttps() throws Exception
   {
      LOGGER.info("doHttps : " + keysURL);

      URL myUrl = new URL(keysURL);
      HttpsURLConnection conn = (HttpsURLConnection)myUrl.openConnection();
      InputStream is = conn.getInputStream();
      InputStreamReader isr = new InputStreamReader(is);
      BufferedReader br = new BufferedReader(isr);

      String inputLine;
      String jsonKeys = ""; 
      while ((inputLine = br.readLine()) != null) {
         jsonKeys = jsonKeys + inputLine;
      }

      br.close();

      return jsonKeys;
   }


   private Key getKeyFromKeys(String jsonKeys, String keyId)
         throws JsonProcessingException, GeneralSecurityException, IOException
      {
         LOGGER.info( "IamSigningKeyResolver::getKeyFromKeys");

         Key key = null;

         ObjectMapper mapper = new ObjectMapper();

         JsonNode keysNode = mapper.readTree(jsonKeys).get("keys");
         if(keysNode.isArray())
         {
            for (JsonNode node : keysNode)
            {
               String nodeContent = mapper.writeValueAsString(node);

               LOGGER.info("key: " + nodeContent);

               Jwk<?> jwk = Jwks.parser().build().parse(nodeContent);

               String jwkkid = jwk.getId();

               LOGGER.info("kid-token : " + keyId + "kid-store : " + jwkkid + " key-type: " + jwk.getType());

               if(keyId.equals(jwkkid))
               {
                  key = jwk.toKey();
               }
            }
         }

         return key;
      }

}

