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

import com.google.common.annotations.VisibleForTesting;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

public final class PBEWithSHA1andAES {
   public static final int DEFAULT_MAX_KDF_ITERATIONS = 32000;
   private Cipher encryptionCipher;
   private final KeyDerivationFuntion encryptionKdf;
   private List<Cipher> decryptionCiphers;
   private final List<KeyDerivationFuntion> decryptionKDFs;
   private Cipher decryptionCipherNoPrefixFound;
   private final KeyDerivationFuntion decryptionKDFNoPrefixFound;
   private final boolean suppressKdfPrefix;
   private final boolean suppressCipherPrefix;
   private String modeName;

   public static PBEWithSHA1andAES createSecureCipher(int maxPBKDF2Iterations) {
      Cipher encryptionCipher = GCMAESCipher.create();
      List<Cipher> decryptionCiphers = new ArrayList<>();
      decryptionCiphers.add(encryptionCipher);
      KeyDerivationFuntion encryptionKdf = PasswordBasedKeyDerivationFunction2.createWithSecureValues(maxPBKDF2Iterations);
      List<KeyDerivationFuntion> decryptionKDFs = new ArrayList<>();
      decryptionKDFs.add(encryptionKdf);
      return new PBEWithSHA1andAES.Builder()
         .setEncryptionCipher(encryptionCipher)
         .setEncryptionKdf(encryptionKdf)
         .setDecryptionCiphers(decryptionCiphers)
         .setDecryptionKDFs(decryptionKDFs)
         .setModeName("Secure Cipher")
         .build();
   }

   public static PBEWithSHA1andAES createSecureCipherWithRandomBlock(int maxPBKDF2Iterations) {
      return createSecureCipher(maxPBKDF2Iterations).prependOneRandomBlock().setModeName("Secure Cipher with random block");
   }

   public static PBEWithSHA1andAES createSecureCipherWithFlexibleDecryption(int maxPBKDF2Iterations) {
      Cipher encryptionCipher = GCMAESCipher.create();
      List<Cipher> decryptionCiphers = new ArrayList<>();
      decryptionCiphers.add(encryptionCipher);
      Cipher legacyCBCCipher = CBCAESCipher.create();
      decryptionCiphers.add(legacyCBCCipher);
      KeyDerivationFuntion encryptionKdf = PasswordBasedKeyDerivationFunction2.createWithSecureValues(maxPBKDF2Iterations);
      List<KeyDerivationFuntion> decryptionKDFs = new ArrayList<>();
      decryptionKDFs.add(encryptionKdf);
      return new PBEWithSHA1andAES.Builder()
         .setEncryptionCipher(encryptionCipher)
         .setDecryptionCiphers(decryptionCiphers)
         .setDecryptionCipherNoPrefixFound(legacyCBCCipher)
         .setEncryptionKdf(encryptionKdf)
         .setDecryptionKDFs(decryptionKDFs)
         .setDecryptionKDFNoPrefixFound(encryptionKdf)
         .setModeName("Secure Cipher with Legacy Decryption")
         .build();
   }

   public static PBEWithSHA1andAES createSecureCipherWithFlexibleDecryptionAndRandomBlock(int maxPBKDF2Iterations) {
      return createSecureCipherWithFlexibleDecryption(maxPBKDF2Iterations)
         .prependOneRandomBlock()
         .setModeName("Secure Cipher with Legacy Decryption and a random block");
   }

   @VisibleForTesting
   static PBEWithSHA1andAES createAuthenticatedNonSecretCipher(int maxPBKDF2Iterations) {
      return createAuthenticatedNonSecretStaticCipher(maxPBKDF2Iterations).prependOneRandomBlock().setModeName("Authenticated, Non-secret Cipher");
   }

   @VisibleForTesting
   static PBEWithSHA1andAES createAuthenticatedNonSecretStaticCipher(int maxPBKDF2Iterations) {
      Cipher encryptionCipher = GCMAESCipher.create();
      List<Cipher> decryptionCiphers = new ArrayList<>();
      decryptionCiphers.add(encryptionCipher);
      Cipher cbcCipher = CBCAESCipher.create();
      decryptionCiphers.add(cbcCipher);
      KeyDerivationFuntion encryptionKdf = PasswordBasedKeyDerivationFunction2.createWithInsecureFixDefaults(maxPBKDF2Iterations);
      List<KeyDerivationFuntion> decryptionKDFs = new ArrayList<>();
      decryptionKDFs.add(encryptionKdf);
      return new PBEWithSHA1andAES.Builder()
         .setEncryptionCipher(encryptionCipher)
         .setDecryptionCiphers(decryptionCiphers)
         .setDecryptionCipherNoPrefixFound(cbcCipher)
         .setEncryptionKdf(encryptionKdf)
         .setSuppressKdfPrefix(true)
         .setDecryptionKDFs(decryptionKDFs)
         .setDecryptionKDFNoPrefixFound(encryptionKdf)
         .setModeName("Authenticated, Non-secret, static cipher")
         .build();
   }

   @VisibleForTesting
   static PBEWithSHA1andAES createAuthenticatedNonSecretCipherWithoutFlexibleDecryption(int maxPBKDF2Iterations) {
      return createAuthenticatedNonSecretStaticCipherWithoutFlexibleDecryption(maxPBKDF2Iterations)
         .prependOneRandomBlock()
         .setModeName("Authenticated, Non-secret cipher without flexible decryption");
   }

   @VisibleForTesting
   static PBEWithSHA1andAES createAuthenticatedNonSecretStaticCipherWithoutFlexibleDecryption(int maxPBKDF2Iterations) {
      Cipher encryptionCipher = GCMAESCipher.create();
      List<Cipher> decryptionCiphers = new ArrayList<>();
      decryptionCiphers.add(encryptionCipher);
      KeyDerivationFuntion encryptionKdf = PasswordBasedKeyDerivationFunction2.createWithInsecureFixDefaults(maxPBKDF2Iterations);
      List<KeyDerivationFuntion> decryptionKDFs = new ArrayList<>();
      decryptionKDFs.add(encryptionKdf);
      return new PBEWithSHA1andAES.Builder()
         .setEncryptionCipher(encryptionCipher)
         .setDecryptionCiphers(decryptionCiphers)
         .setEncryptionKdf(encryptionKdf)
         .setSuppressKdfPrefix(true)
         .setDecryptionKDFs(decryptionKDFs)
         .setDecryptionKDFNoPrefixFound(encryptionKdf)
         .setModeName("Authenticated, Non-Secret, Static Cipher Without Flexible Decryption")
         .build();
   }

   public static PBEWithSHA1andAES createStaticCipher(int maxPBKDF2Iterations) {
      Cipher encryptionCipher = CBCAESCipher.create();
      List<Cipher> decryptionCiphers = new ArrayList<>();
      decryptionCiphers.add(encryptionCipher);
      decryptionCiphers.add(GCMAESCipher.create());
      KeyDerivationFuntion encryptionKdf = PasswordBasedKeyDerivationFunction2.createWithInsecureFixDefaults(maxPBKDF2Iterations);
      List<KeyDerivationFuntion> decryptionKDFs = new ArrayList<>();
      decryptionKDFs.add(encryptionKdf);
      return new PBEWithSHA1andAES.Builder()
         .setEncryptionCipher(encryptionCipher)
         .setSuppressCipherPrefix(true)
         .setDecryptionCiphers(decryptionCiphers)
         .setDecryptionCipherNoPrefixFound(encryptionCipher)
         .setEncryptionKdf(encryptionKdf)
         .setSuppressKdfPrefix(true)
         .setDecryptionKDFs(decryptionKDFs)
         .setDecryptionKDFNoPrefixFound(encryptionKdf)
         .setModeName("Static Cipher")
         .build();
   }

   public static PBEWithSHA1andAES createUnauthenticatedCipher(int maxPBKDF2Iterations) {
      Cipher encryptionCipher = CBCAESCipher.create();
      List<Cipher> decryptionCiphers = new ArrayList<>();
      decryptionCiphers.add(encryptionCipher);
      decryptionCiphers.add(GCMAESCipher.create());
      KeyDerivationFuntion encryptionKdf = PasswordBasedKeyDerivationFunction2.createWithInsecureFixDefaults(maxPBKDF2Iterations);
      List<KeyDerivationFuntion> decryptionKDFs = new ArrayList<>();
      decryptionKDFs.add(encryptionKdf);
      return new PBEWithSHA1andAES.Builder()
         .setEncryptionCipher(encryptionCipher)
         .setSuppressCipherPrefix(true)
         .setDecryptionCiphers(decryptionCiphers)
         .setDecryptionCipherNoPrefixFound(encryptionCipher)
         .setEncryptionKdf(encryptionKdf)
         .setSuppressKdfPrefix(true)
         .setDecryptionKDFs(decryptionKDFs)
         .setDecryptionKDFNoPrefixFound(encryptionKdf)
         .setModeName("Unauthenticated Cipher")
         .build()
         .prependOneRandomBlock();
   }

   private PBEWithSHA1andAES prependOneRandomBlock() {
      this.encryptionCipher = new OneBlockRandomnessPrepender(this.encryptionCipher);
      List<Cipher> modifiedDecCiphers = new ArrayList<>();

      for (Cipher decCipher : this.decryptionCiphers) {
         modifiedDecCiphers.add(new OneBlockRandomnessPrepender(decCipher));
      }

      this.decryptionCiphers = modifiedDecCiphers;
      if (this.decryptionCipherNoPrefixFound != null) {
         this.decryptionCipherNoPrefixFound = new OneBlockRandomnessPrepender(this.decryptionCipherNoPrefixFound);
      }

      return this;
   }

   private PBEWithSHA1andAES(PBEWithSHA1andAES.Builder builder) {
      this.encryptionCipher = builder.getEncryptionCipher();
      this.encryptionKdf = builder.getEncryptionKdf();
      this.decryptionCiphers = builder.getDecryptionCiphers();
      this.decryptionKDFs = builder.getDecryptionKDFs();
      this.decryptionCipherNoPrefixFound = builder.getDecryptionCipherNoPrefixFound();
      this.decryptionKDFNoPrefixFound = builder.getDecryptionKDFNoPrefixFound();
      this.suppressKdfPrefix = builder.isSuppressKdfPrefix();
      this.suppressCipherPrefix = builder.isSuppressCipherPrefix();
      this.modeName = builder.getModeName();
   }

   public byte[] encrypt(String password, byte[] plaintext) throws EncryptionException {
      Objects.requireNonNull(password, "password null");
      Objects.requireNonNull(plaintext, "plaintext is null");
      KeyDerivationFunctionResult kdfResult = this.computeKdfParameters(password);
      CipherText ciphertext = this.encryptionCipher.encrypt(plaintext, kdfResult.getCipherParams());
      this.setNecessaryPrefixes(ciphertext, kdfResult.getPrefix());
      return ciphertext.getCiphertextWithPrefixes();
   }

   private KeyDerivationFunctionResult computeKdfParameters(String password) throws EncryptionException {
      try {
         return this.encryptionKdf.deriveParametersForEncryption(password);
      } catch (KeyDerivationException e) {
         throw new EncryptionException("Cannot parse KDF Prefix: " + e.getMessage(), e);
      }
   }

   private void setNecessaryPrefixes(CipherText ciphertext, String kdfPrefix) throws EncryptionException {
      if (!this.suppressKdfPrefix) {
         try {
            ciphertext.setKdfPrefix(kdfPrefix);
         } catch (CipherText.InvalidPrefixException e) {
            throw new EncryptionException("Kdf provided invalid prefix: " + kdfPrefix, e);
         }
      }

      if (this.suppressCipherPrefix) {
         try {
            ciphertext.setCipherPrefix(null);
         } catch (CipherText.InvalidPrefixException e) {
            throw new EncryptionException("Programming error, null should be accepted as cipher prefix.", e);
         }
      }
   }

   public byte[] decrypt(String password, byte[] ciphertext) throws DecryptionException {
      Objects.requireNonNull(password, "password null");
      Objects.requireNonNull(ciphertext, "ciphertext is null");

      for (int maxPossiblePrefixes = 2; maxPossiblePrefixes >= 0; maxPossiblePrefixes--) {
         try {
            CipherText parsedCipherText = CipherText.createFromCiphertextWithPotentialPrefixes(ciphertext, maxPossiblePrefixes);
            KeyAndIV cipherParams = this.computeCipherParams(password, parsedCipherText.getKdfPrefix());
            return this.decryptInternal(parsedCipherText, ciphertext, cipherParams);
         } catch (PBEWithSHA1andAES.NoSuitablePrimitiveFound | CipherText.InvalidPrefixException e) {
         }
      }

      throw new DecryptionException("Could not decrypt. Actually, this here should not happen.");
   }

   private KeyAndIV computeCipherParams(String password, String kdfPrefix) throws DecryptionException, PBEWithSHA1andAES.NoSuitablePrimitiveFound {
      if (kdfPrefix != null) {
         for (KeyDerivationFuntion kdf : this.decryptionKDFs) {
            if (kdf.isSuitablePrefix(kdfPrefix)) {
               try {
                  return kdf.deriveParametersForDecryption(password, kdfPrefix);
               } catch (KeyDerivationException e) {
                  throw new DecryptionException("Cannot parse KDF Prefix: " + e.getMessage(), e);
               }
            }
         }

         throw new PBEWithSHA1andAES.NoSuitablePrimitiveFound(
            "No suitable KDF found for prefix '" + kdfPrefix + "'. Probably a wrongly parsed prefix. Repeat without parsing this prefix."
         );
      } else if (this.decryptionKDFNoPrefixFound != null) {
         try {
            return this.decryptionKDFNoPrefixFound.deriveParametersForDecryption(password, null);
         } catch (KeyDerivationException e) {
            throw new DecryptionException("Cannot parse KDF Prefix: " + e.getMessage(), e);
         }
      } else {
         throw new DecryptionException("Cannot decrypt: KDF Prefix is null but no decryptionKDFNoPrefixFound is given.");
      }
   }

   private byte[] decryptInternal(CipherText parsedCipherText, byte[] rawCipherText, KeyAndIV cipherParams) throws DecryptionException, PBEWithSHA1andAES.NoSuitablePrimitiveFound {
      String cipherPrefix = parsedCipherText.getCipherPrefix();
      if (cipherPrefix != null) {
         for (Cipher decCiphers : this.decryptionCiphers) {
            if (decCiphers.isSuitablePrefix(cipherPrefix)) {
               return decCiphers.decrypt(parsedCipherText, cipherParams);
            }
         }

         throw new PBEWithSHA1andAES.NoSuitablePrimitiveFound(
            "No suitable Cipher found for prefix '" + cipherPrefix + "'. Probably a wrongly parsed prefix. Repeat without parsing this prefix."
         );
      } else if (this.decryptionCipherNoPrefixFound != null) {
         return this.decryptionCipherNoPrefixFound.decrypt(CipherText.createWithPlainCiphertext(rawCipherText), cipherParams);
      } else {
         throw new DecryptionException("Cannot decrypt: Cipher prefix is null but no decryptionCipherNoPrefixFound is given.");
      }
   }

   public String getModeName() {
      return this.modeName;
   }

   private PBEWithSHA1andAES setModeName(String name) {
      this.modeName = name;
      return this;
   }

   private static class Builder {
      private Cipher encryptionCipher;
      private KeyDerivationFuntion encryptionKdf;
      private List<Cipher> decryptionCiphers;
      private List<KeyDerivationFuntion> decryptionKDFs;
      private Cipher decryptionCipherNoPrefixFound;
      private KeyDerivationFuntion decryptionKDFNoPrefixFound;
      private boolean suppressKdfPrefix;
      private boolean suppressCipherPrefix;
      private String modeName;

      Builder() {
      }

      PBEWithSHA1andAES build() {
         return new PBEWithSHA1andAES(this);
      }

      Cipher getEncryptionCipher() {
         return this.encryptionCipher;
      }

      PBEWithSHA1andAES.Builder setEncryptionCipher(Cipher encryptionCipher) {
         this.encryptionCipher = encryptionCipher;
         return this;
      }

      KeyDerivationFuntion getEncryptionKdf() {
         return this.encryptionKdf;
      }

      PBEWithSHA1andAES.Builder setEncryptionKdf(KeyDerivationFuntion encryptionKdf) {
         this.encryptionKdf = encryptionKdf;
         return this;
      }

      List<Cipher> getDecryptionCiphers() {
         return this.decryptionCiphers;
      }

      PBEWithSHA1andAES.Builder setDecryptionCiphers(List<Cipher> decryptionCiphers) {
         this.decryptionCiphers = decryptionCiphers;
         return this;
      }

      List<KeyDerivationFuntion> getDecryptionKDFs() {
         return this.decryptionKDFs;
      }

      PBEWithSHA1andAES.Builder setDecryptionKDFs(List<KeyDerivationFuntion> decryptionKDFs) {
         this.decryptionKDFs = decryptionKDFs;
         return this;
      }

      Cipher getDecryptionCipherNoPrefixFound() {
         return this.decryptionCipherNoPrefixFound;
      }

      PBEWithSHA1andAES.Builder setDecryptionCipherNoPrefixFound(Cipher decryptionCipherNoPrefixFound) {
         this.decryptionCipherNoPrefixFound = decryptionCipherNoPrefixFound;
         return this;
      }

      KeyDerivationFuntion getDecryptionKDFNoPrefixFound() {
         return this.decryptionKDFNoPrefixFound;
      }

      PBEWithSHA1andAES.Builder setDecryptionKDFNoPrefixFound(KeyDerivationFuntion decryptionKDFNoPrefixFound) {
         this.decryptionKDFNoPrefixFound = decryptionKDFNoPrefixFound;
         return this;
      }

      boolean isSuppressKdfPrefix() {
         return this.suppressKdfPrefix;
      }

      PBEWithSHA1andAES.Builder setSuppressKdfPrefix(boolean suppressKdfPrefix) {
         this.suppressKdfPrefix = suppressKdfPrefix;
         return this;
      }

      boolean isSuppressCipherPrefix() {
         return this.suppressCipherPrefix;
      }

      PBEWithSHA1andAES.Builder setSuppressCipherPrefix(boolean suppressCipherPrefix) {
         this.suppressCipherPrefix = suppressCipherPrefix;
         return this;
      }

      String getModeName() {
         return this.modeName;
      }

      PBEWithSHA1andAES.Builder setModeName(String modeName) {
         this.modeName = modeName;
         return this;
      }
   }

   private static final class NoSuitablePrimitiveFound extends CipherException {
      private NoSuitablePrimitiveFound(String msg) {
         super(msg);
      }
   }
}