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

import com.airlock.iam.security.util.SeededSecureRandom;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.base.Strings;
import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;
import java.util.Arrays;
import java.util.regex.Pattern;
import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.PBEParametersGenerator;
import org.bouncycastle.crypto.generators.PKCS5S2ParametersGenerator;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.ParametersWithIV;
import org.bouncycastle.util.encoders.Hex;

public final class PasswordBasedKeyDerivationFunction2 implements KeyDerivationFuntion {
   @VisibleForTesting
   static final byte[] DEFAULT_PBE_SALT = Hex.decode(
      "fca0e420651ac42d78b950f60b5579f28484a7df76ec00a62de839d75def63de819875093381837ec04f2b76bdd5af0f5cada0dffb63fea4e0a6f75b94656a01d265ee73b5c5905571f7ee8b261374e45e5853809ccba0e49919349986584c9d0b8037f2724e562862e704cc41fa5d567d45b5354c625233e83c63528ec3b05ac812e25533c2cce756c0b3c387e8203b33a44e89d4b5ff9ded465833ed19500e408fb4ea54681b27fa8f9168b82c0c58f62a9f9838c867dcc00cf4ee180410c223fa4cda004716ea74a325d9e90aa23f92421621a3c07db533b69484a39e151909af66aa62eaa49e54f4e40489d56e03ec52941540fc1053cbe6b38e5ac7a533"
   );
   private static final int DEFAULT_ITERATION_COUNT = 16000;
   private static final int LEGACY_ITERATION_COUNT = 128;
   private static final int DEFAULT_IV_SIZE = 128;
   private static final int DEFAULT_KEY_SIZE = 256;
   private static final int SALT_LENGTH_IN_BYTES = 16;
   private static final String PBKDF2 = "PBKDF2";
   private static final char SEPARATOR = '|';
   private final int iterations;
   private final int maxPBKDF2Iterations;
   private final int requiredKeySize;
   private final int requiredIVSize;
   private final boolean useFixDefaultSalt;
   private final SecureRandom rand = SeededSecureRandom.getInstance();

   private PasswordBasedKeyDerivationFunction2(int iterations, int maxDecryptionIterations, int requiredIVSize, int requiredKeySize, boolean useFixDefaultSalt) {
      Preconditions.checkArgument(iterations >= 1, "Iteration Count must be at least 1, received " + iterations);
      Preconditions.checkArgument(maxDecryptionIterations >= 1, "Max Decryption Iteration Count must be at least 1, received " + maxDecryptionIterations);
      Preconditions.checkArgument(requiredIVSize >= 1, "IV size must be at least 1, received " + requiredIVSize);
      Preconditions.checkArgument(requiredKeySize >= 1, "Key size must be at least 1, received " + requiredKeySize);
      this.iterations = iterations;
      this.maxPBKDF2Iterations = maxDecryptionIterations;
      this.requiredIVSize = requiredIVSize;
      this.requiredKeySize = requiredKeySize;
      this.useFixDefaultSalt = useFixDefaultSalt;
   }

   public static PasswordBasedKeyDerivationFunction2 createWithInsecureFixDefaults(int maxCipherDecryptionIterations) {
      return new PasswordBasedKeyDerivationFunction2(128, maxCipherDecryptionIterations, 128, 256, true);
   }

   public static PasswordBasedKeyDerivationFunction2 createWithSecureValues(int maxCipherDecryptionIterations) {
      return new PasswordBasedKeyDerivationFunction2(16000, maxCipherDecryptionIterations, 128, 256, false);
   }

   @Override
   public KeyDerivationFunctionResult deriveParametersForEncryption(String password) throws KeyDerivationException {
      byte[] salt = this.computeSalt();
      KeyAndIV cipherParams = this.deriveParameters(password, salt, this.iterations);
      String prefix = this.getPrefix(salt);
      return new KeyDerivationFunctionResult(cipherParams, prefix);
   }

   private byte[] computeSalt() {
      if (this.useFixDefaultSalt) {
         return (byte[])DEFAULT_PBE_SALT.clone();
      }

      byte[] salt = new byte[16];
      this.rand.nextBytes(salt);
      return salt;
   }

   private KeyAndIV deriveParameters(String password, byte[] salt, int numberOfIterations) throws KeyDerivationException {
      PKCS5S2ParametersGenerator generator = new PKCS5S2ParametersGenerator();
      generator.init(PBEParametersGenerator.PKCS5PasswordToUTF8Bytes(password.toCharArray()), salt, numberOfIterations);
      CipherParameters params = generator.generateDerivedParameters(this.requiredKeySize, this.requiredIVSize);
      return this.convertToKeyAndIV(params);
   }

   private KeyAndIV convertToKeyAndIV(CipherParameters params) throws KeyDerivationException {
      if (params instanceof ParametersWithIV paramsWithIV) {
         byte[] iv = (byte[])paramsWithIV.getIV().clone();
         if (paramsWithIV.getParameters() instanceof KeyParameter keyParam) {
            byte[] key = (byte[])keyParam.getKey().clone();
            return new KeyAndIV(key, iv);
         }
      }

      throw new KeyDerivationException("The types used in the computed CipherParameters do not match the expectation.");
   }

   private String getPrefix(byte[] salt) {
      return "PBKDF2|" + this.iterations + "|" + saltToString(salt);
   }

   @Override
   public KeyAndIV deriveParametersForDecryption(String password, String prefix) throws KeyDerivationException {
      if (this.isSuitablePrefix(prefix)) {
         PasswordBasedKeyDerivationFunction2.ParseResult parsed = this.parsePrefix(prefix);
         return this.deriveParameters(password, parsed.getParsedSalt(), parsed.getParsedIterations());
      } else if (Strings.isNullOrEmpty(prefix)) {
         return this.deriveParameters(password, DEFAULT_PBE_SALT, 128);
      } else {
         throw new KeyDerivationException("Wrong KDF Prefix for PBKDF2: " + prefix);
      }
   }

   @Override
   public boolean isSuitablePrefix(String prefix) {
      return Strings.nullToEmpty(prefix).startsWith("PBKDF2");
   }

   private static String saltToString(byte[] salt) {
      return new String(Hex.encode(salt), StandardCharsets.ISO_8859_1);
   }

   private static byte[] stringToSalt(String salt) {
      return Hex.decode(salt.getBytes(StandardCharsets.ISO_8859_1));
   }

   private PasswordBasedKeyDerivationFunction2.ParseResult parsePrefix(String prefix) throws KeyDerivationException {
      try {
         String[] params = prefix.split(Pattern.quote(Character.toString('|')));
         Preconditions.checkArgument(params.length == 3, "Number of parameters found in prefix was " + params.length + ", expected 3. The prefix is " + prefix);
         int decryptionIterations = Integer.parseInt(params[1]);
         Preconditions.checkState(
            1 <= decryptionIterations && decryptionIterations <= this.maxPBKDF2Iterations,
            "The number of PBKDF2 iterations specified in the ticket ("
               + decryptionIterations
               + ") exceeds the maximum allowed iterations ("
               + this.maxPBKDF2Iterations
               + "). If you are using an AES256DecryptionTicketDecoder, you may change the maximum allowed iterations in the configuration."
         );
         byte[] salt = stringToSalt(params[2]);
         return new PasswordBasedKeyDerivationFunction2.ParseResult(decryptionIterations, salt);
      } catch (RuntimeException e) {
         throw new KeyDerivationException("Cannot parse KDF Prefix '" + prefix + "'. Cause: " + e.getMessage(), e);
      }
   }

   private static class ParseResult {
      private final int parsedIterations;
      private final byte[] parsedSalt;

      ParseResult(int iterations, byte[] salt) {
         this.parsedIterations = iterations;
         this.parsedSalt = Arrays.copyOf(salt, salt.length);
      }

      int getParsedIterations() {
         return this.parsedIterations;
      }

      byte[] getParsedSalt() {
         return this.parsedSalt;
      }
   }
}