package com.airlock.iam.security.certificate;

import com.airlock.iam.security.util.SecurityProvider;
import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.security.GeneralSecurityException;
import java.security.InvalidAlgorithmParameterException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.Provider;
import java.security.PublicKey;
import java.security.interfaces.ECPrivateKey;
import java.security.interfaces.ECPublicKey;
import java.security.interfaces.RSAPrivateCrtKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.ECGenParameterSpec;
import java.security.spec.EncodedKeySpec;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.RSAPublicKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.Optional;
import java.util.regex.Pattern;
import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.openssl.PEMKeyPair;
import org.bouncycastle.openssl.PEMParser;

public final class AsymmetricKeyPairUtils {
   private static final Provider PROVIDER = SecurityProvider.getProvider();
   private static final Pattern HEADER_PATTERN = Pattern.compile("-----BEGIN [ A-Z]+-----");
   private static final Pattern FOOTER_PATTERN = Pattern.compile("-----END [ A-Z]+-----");
   private static final Pattern WHITESPACE_PATTERN = Pattern.compile("\\s+");

   public static PublicKey publicKeyFromPem(String rawPublicKey) {
      return Optional.of(rawPublicKey)
         .map(pubKey -> removeMatchedPattern(pubKey, HEADER_PATTERN))
         .map(pubKey -> removeMatchedPattern(pubKey, FOOTER_PATTERN))
         .map(pubKey -> removeMatchedPattern(pubKey, WHITESPACE_PATTERN))
         .map(AsymmetricKeyPairUtils::decodePlainBase64)
         .map(AsymmetricKeyPairUtils::generatePublicKey)
         .orElseThrow(IllegalArgumentException::new);
   }

   public static RSAPublicKey rsaPublicKeyFromMimeBase64(String key) {
      try {
         X509EncodedKeySpec keySpec = new X509EncodedKeySpec(Base64.getMimeDecoder().decode(key));
         KeyFactory fact = KeyFactory.getInstance("RSA");
         return (RSAPublicKey)fact.generatePublic(keySpec);
      } catch (GeneralSecurityException e) {
         throw new IllegalArgumentException(e);
      }
   }

   public static RSAPublicKey rsaPublicKey(RSAPrivateKey key) {
      try {
         RSAPublicKeySpec publicKeySpec = new RSAPublicKeySpec(key.getModulus(), ((RSAPrivateCrtKey)key).getPublicExponent());
         KeyFactory keyFactory = KeyFactory.getInstance("RSA");
         return (RSAPublicKey)keyFactory.generatePublic(publicKeySpec);
      } catch (GeneralSecurityException e) {
         throw new IllegalArgumentException(e);
      }
   }

   public static RSAPrivateKey rsaPrivateKeyFromMimeBase64(String key) {
      return rsaPrivateKey(Base64.getMimeDecoder().decode(key));
   }

   public static RSAPrivateKey rsaPrivateKey(byte[] key) {
      try {
         PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(key);
         KeyFactory keyFactory = KeyFactory.getInstance("RSA");
         return (RSAPrivateKey)keyFactory.generatePrivate(keySpec);
      } catch (GeneralSecurityException e) {
         throw new IllegalArgumentException(e);
      }
   }

   public static ECPublicKey ecPublicKey(String key) {
      try {
         X509EncodedKeySpec keySpec = new X509EncodedKeySpec(Base64.getMimeDecoder().decode(key));
         KeyFactory fact = KeyFactory.getInstance("EC");
         return (ECPublicKey)fact.generatePublic(keySpec);
      } catch (GeneralSecurityException e) {
         throw new IllegalArgumentException(e);
      }
   }

   public static ECPrivateKey ecPrivateKey(String key) {
      try {
         PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(Base64.getMimeDecoder().decode(key));
         KeyFactory fact = KeyFactory.getInstance("EC");
         return (ECPrivateKey)fact.generatePrivate(keySpec);
      } catch (GeneralSecurityException e) {
         throw new IllegalArgumentException(e);
      }
   }

   public static KeyPair generateECKeyPair() {
      return generateECKeyPair(AsymmetricKeyPairUtils.ECCurve.P_256);
   }

   public static KeyPair generateECKeyPair(AsymmetricKeyPairUtils.ECCurve curve) {
      try {
         ECGenParameterSpec spec = new ECGenParameterSpec(curve.name);
         KeyPairGenerator ecKeyPairGenerator = KeyPairGenerator.getInstance("EC");
         ecKeyPairGenerator.initialize(spec);
         return ecKeyPairGenerator.generateKeyPair();
      } catch (NoSuchAlgorithmException | InvalidAlgorithmParameterException e) {
         throw new IllegalStateException(e);
      }
   }

   public static KeyPair generateRSAKeyPair() {
      return generateRSAKeyPair(2048);
   }

   public static KeyPair generateRSAKeyPair(int keysize) {
      try {
         KeyPairGenerator rsaKeyPairGenerator = KeyPairGenerator.getInstance("RSA");
         rsaKeyPairGenerator.initialize(keysize);
         return rsaKeyPairGenerator.generateKeyPair();
      } catch (NoSuchAlgorithmException e) {
         throw new IllegalStateException(e);
      }
   }

   @SuppressFBWarnings("REC_CATCH_EXCEPTION")
   public static KeyPair fromPEMInputStream(InputStream inputStream) {
      try {
         PEMParser pemParser;
         label40: {
            pemParser = new PEMParser(new InputStreamReader(inputStream, StandardCharsets.UTF_8));

            KeyPair var4;
            try {
               while (true) {
                  Object pemObject;
                  if ((pemObject = pemParser.readObject()) == null) {
                     break label40;
                  }

                  if (pemObject instanceof PEMKeyPair pemKeyPair) {
                     var4 = parseKeyPair(pemKeyPair);
                     break;
                  }
               }
            } catch (Throwable var6) {
               try {
                  pemParser.close();
               } catch (Throwable var5) {
                  var6.addSuppressed(var5);
               }

               throw var6;
            }

            pemParser.close();
            return var4;
         }

         pemParser.close();
      } catch (Exception e) {
         throw new IllegalArgumentException("Error parsing PEM input stream: " + e.getMessage(), e);
      }

      throw new IllegalArgumentException("Either public key or private key not found in PEM");
   }

   private static KeyPair parseKeyPair(PEMKeyPair pemKeyPair) throws IOException, NoSuchAlgorithmException, InvalidKeySpecException {
      SubjectPublicKeyInfo subjectPublicKeyInfo = pemKeyPair.getPublicKeyInfo();
      PrivateKeyInfo privateKeyInfo = pemKeyPair.getPrivateKeyInfo();
      EncodedKeySpec publicKeySpec = new X509EncodedKeySpec(subjectPublicKeyInfo.getEncoded());
      EncodedKeySpec privateKeySpec = new PKCS8EncodedKeySpec(privateKeyInfo.getEncoded());
      AlgorithmIdentifier algorithm = subjectPublicKeyInfo.getAlgorithm();
      KeyFactory keyFactory = KeyFactory.getInstance(algorithm.getAlgorithm().getId(), PROVIDER);
      return new KeyPair(keyFactory.generatePublic(publicKeySpec), keyFactory.generatePrivate(privateKeySpec));
   }

   private static PublicKey generatePublicKey(byte[] encodedPublicKey) {
      X509EncodedKeySpec spec = new X509EncodedKeySpec(encodedPublicKey);

      try {
         return KeyFactory.getInstance("RSA").generatePublic(spec);
      } catch (GeneralSecurityException e) {
         try {
            return KeyFactory.getInstance("EC").generatePublic(spec);
         } catch (GeneralSecurityException e2) {
            throw new IllegalArgumentException(e2);
         }
      }
   }

   private static byte[] decodePlainBase64(String publicKey) {
      return Base64.getDecoder().decode(publicKey);
   }

   private static String removeMatchedPattern(String publicKey, Pattern pattern) {
      return pattern.matcher(publicKey).replaceAll("");
   }

   private AsymmetricKeyPairUtils() {
      throw new AssertionError();
   }

   public enum ECCurve {
      P_256("secp256r1"),
      P_256K("secp256k1"),
      P_384("secp384r1"),
      P_521("secp521r1");

      private final String name;

      ECCurve(String name) {
         this.name = name;
      }
   }
}