package org.radarcns.management.security.jwt;

import java.io.IOException;
import java.security.KeyPair;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.UnrecoverableKeyException;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.util.List;
import java.util.Locale;
import java.util.Objects;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.Resource;

/**
 * Similar to Spring's
 * {@link org.springframework.security.oauth2.provider.token.store.KeyStoreKeyFactory}. However
 * this class does not assume a specific key type, while the Spring factory assumes RSA keys.
 */
public class RadarKeyStoreKeyFactory {
    private static final Logger logger = LoggerFactory.getLogger(RadarKeyStoreKeyFactory.class);

    private final List<Resource> resources;

    private final char[] password;
    private final KeyStore store;

    public RadarKeyStoreKeyFactory(@Nonnull List<Resource> resources, @Nonnull char[] password) {
        this.resources = Objects.requireNonNull(resources);
        if (this.resources.isEmpty()) {
            throw new IllegalArgumentException("No JWT keystore resource paths specified.");
        }
        this.password = Objects.requireNonNull(password);
        this.store = loadStore();
    }

    private @Nonnull KeyStore loadStore() {
        for (Resource resource : resources) {
            if (!resource.exists()) {
                logger.trace("Ignoring non-existant JWT key store {}", resource);
                continue;
            }
            try {
                String fileName = resource.getFilename().toLowerCase(Locale.US);
                String type = fileName.endsWith(".pfx") || fileName.endsWith(".p12")
                        ? "PKCS12" : "jks";
                KeyStore localStore = KeyStore.getInstance(type);
                localStore.load(resource.getInputStream(), this.password);
                logger.debug("Loaded JWT key store {}", resource);
                return localStore;
            } catch (CertificateException | NoSuchAlgorithmException | KeyStoreException
                    | IOException ex) {
                logger.error("Cannot load JWT key store {}", ex);
            }
        }
        throw new IllegalArgumentException("Cannot load any of the given JWT key stores "
                + resources);
    }

    /**
     * Get a key pair from the store using the store password.
     * @param alias key pair alias
     * @return loaded key pair or {@code null} if the key store does not contain a loadable key with
     *         given alias.
     * @throws IllegalArgumentException if the key alias password is wrong or the key cannot
     *                                  loaded.
     */
    public @Nullable KeyPair getKeyPair(@Nullable String alias) {
        return getKeyPair(alias, password);
    }

    /**
     * Get a key pair from the store with a given alias and password.
     * @param alias key pair alias
     * @param password key pair password
     * @return loaded key pair or {@code null} if the key store does not contain a loadable key with
     *         given alias.
     * @throws IllegalArgumentException if the key alias password is wrong or the key cannot
     *                                  loaded.
     */
    public @Nullable KeyPair getKeyPair(@Nullable String alias, char[] password) {
        try {
            PrivateKey key = (PrivateKey) store.getKey(alias, password);
            if (key == null) {
                logger.warn("JWT key store {} does not contain private key pair for alias {}",
                        resources, alias);
                return null;
            }
            Certificate cert = store.getCertificate(alias);
            if (cert == null) {
                logger.warn("JWT key store {} does not contain certificate pair for alias {}",
                        resources, alias);
                return null;
            }
            PublicKey publicKey = cert.getPublicKey();
            if (publicKey == null) {
                logger.warn("JWT key store {} does not contain public key pair for alias {}",
                        resources, alias);
                return null;
            }
            return new KeyPair(publicKey, key);
        } catch (NoSuchAlgorithmException ex) {
            logger.warn("JWT key store {} contains unknown algorithm for key pair with alias {}: {}",
                    resources, alias, ex.toString());
            return null;
        } catch (UnrecoverableKeyException | KeyStoreException ex) {
            throw new IllegalArgumentException("JWT key store " + resources
                    + " contains unrecoverable key pair with alias "
                    + alias + " (the password may be wrong)", ex);
        }
    }
}

@@ -1022,53 +1022,8 @@
 );%0A%0A
-    private final List%3CResource%3E resources;%0A%0A
     
@@ -1086,16 +1086,53 @@
 e store;
+%0A    private Resource loadedResource;
 %0A%0A    pu
@@ -1236,25 +1236,24 @@
    this.
-resources
+password
  = Objec
@@ -1270,25 +1270,24 @@
 NonNull(
-resources
+password
 );%0A     
@@ -1293,235 +1293,68 @@
     
-if (
 this.
-resources.isEmpty()) %7B%0A            throw new IllegalArgumentException(%22No JWT keystore resource paths specified.%22);%0A        %7D%0A        this.password = Objects.requireNonNull(password);%0A        this.store = loadStore(
+store = loadStore(Objects.requireNonNull(resources)
 );%0A
+%0A
     
@@ -1396,16 +1396,40 @@
 adStore(
+List%3CResource%3E resources
 ) %7B%0A    
@@ -2058,32 +2058,80 @@
 %7B%7D%22, resource);%0A
+                this.loadedResource = resource;%0A
                 
@@ -3744,33 +3744,38 @@
             
-r
+loadedR
 esource
-s
 , alias);%0A  
@@ -4022,33 +4022,38 @@
             
-r
+loadedR
 esource
-s
 , alias);%0A  
@@ -4299,33 +4299,38 @@
             
-r
+loadedR
 esource
-s
 , alias);%0A  
@@ -4588,25 +4588,30 @@
         
-r
+loadedR
 esource
-s
 , alias,
@@ -4787,25 +4787,30 @@
 ore %22 + 
-r
+loadedR
 esource
-s
 %0A       

Fix error messages to read currenlty loaded resource