package com.airlock.iam.security.util.crypto;

import java.io.IOException;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.security.AlgorithmParameters;
import java.security.GeneralSecurityException;
import java.security.Key;
import java.security.KeyFactory;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.KeySpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import javax.crypto.EncryptedPrivateKeyInfo;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.PBEKeySpec;

public final class RsaCipher {
   private static final Charset STRING_ENCODING = StandardCharsets.UTF_8;
   private static final String CIPHER = "RSA/ECB/PKCS1Padding";

   private RsaCipher() {
   }

   public static byte[] encrypt(String string, PublicKey key) throws GeneralSecurityException {
      if (string == null || string.length() == 0) {
         throw new IllegalArgumentException("String to encrypt is null or empty");
      }

      if (key == null) {
         throw new IllegalArgumentException("public key must not be null");
      }

      javax.crypto.Cipher cipher = javax.crypto.Cipher.getInstance("RSA/ECB/PKCS1Padding");
      cipher.init(1, key);
      return cipher.doFinal(string.getBytes(STRING_ENCODING));
   }

   public static String decrypt(byte[] encrypted, PrivateKey key) throws GeneralSecurityException {
      if (encrypted == null || encrypted.length == 0) {
         throw new IllegalArgumentException("data to decrypt is null or empty");
      }

      if (key == null) {
         throw new IllegalArgumentException("private key must not be null");
      }

      javax.crypto.Cipher cipher = javax.crypto.Cipher.getInstance("RSA/ECB/PKCS1Padding");
      cipher.init(2, key);
      return new String(cipher.doFinal(encrypted), STRING_ENCODING);
   }

   public static RSAPublicKey toPublicKey(byte[] bytes) throws GeneralSecurityException {
      KeyFactory rsaKeyFac = KeyFactory.getInstance("RSA");
      X509EncodedKeySpec keySpec = new X509EncodedKeySpec(bytes);
      return (RSAPublicKey)rsaKeyFac.generatePublic(keySpec);
   }

   public static RSAPrivateKey toPrivateKey(byte[] bytes) throws GeneralSecurityException {
      KeyFactory rsaKeyFac = KeyFactory.getInstance("RSA");
      PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(bytes);
      return (RSAPrivateKey)rsaKeyFac.generatePrivate(keySpec);
   }

   public static RSAPrivateKey toPrivateKey(byte[] privateKeyBytes, String password) throws GeneralSecurityException, IOException {
      EncryptedPrivateKeyInfo encryptPKInfo = new EncryptedPrivateKeyInfo(privateKeyBytes);
      javax.crypto.Cipher cipher = javax.crypto.Cipher.getInstance(encryptPKInfo.getAlgName());
      PBEKeySpec pbeKeySpec = new PBEKeySpec(password.toCharArray());
      SecretKeyFactory secFac = SecretKeyFactory.getInstance(encryptPKInfo.getAlgName());
      Key pbeKey = secFac.generateSecret(pbeKeySpec);
      AlgorithmParameters algParams = encryptPKInfo.getAlgParameters();
      cipher.init(2, pbeKey, algParams);
      KeySpec pkcs8KeySpec = encryptPKInfo.getKeySpec(cipher);
      KeyFactory kf = KeyFactory.getInstance("RSA");
      return (RSAPrivateKey)kf.generatePrivate(pkcs8KeySpec);
   }

   public static byte[] wrapKey(SecretKey secretKey, PublicKey publicKey) throws GeneralSecurityException {
      javax.crypto.Cipher cipher = javax.crypto.Cipher.getInstance("RSA/ECB/PKCS1Padding");
      cipher.init(3, publicKey);
      return cipher.wrap(secretKey);
   }

   public static SecretKey unwrapKey(byte[] encryptedKey, String wrappedKeyAlgorithm, PrivateKey privateKey) throws GeneralSecurityException {
      javax.crypto.Cipher cipher = javax.crypto.Cipher.getInstance("RSA/ECB/PKCS1Padding");
      cipher.init(4, privateKey);
      return (SecretKey)cipher.unwrap(encryptedKey, wrappedKeyAlgorithm, 3);
   }
}