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

import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.CryptoException;
import org.bouncycastle.crypto.DataLengthException;
import org.bouncycastle.crypto.InvalidCipherTextException;
import org.bouncycastle.crypto.MultiBlockCipher;
import org.bouncycastle.crypto.engines.AESEngine;
import org.bouncycastle.crypto.modes.AEADBlockCipher;
import org.bouncycastle.crypto.modes.GCMBlockCipher;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.ParametersWithIV;

public final class GCMAESCipher implements Cipher {
   private static final String GCM_CIPHER_PREFIX = "GCM";
   private final int computedBlockSizeInBytes = createAESEngine().getBlockSize();

   public static GCMAESCipher create() {
      return new GCMAESCipher();
   }

   private GCMAESCipher() {
   }

   @Override
   public CipherText encrypt(byte[] plaintext, KeyAndIV parameters) throws EncryptionException {
      if (parameters.getIV().length < this.blockSizeInBytes()) {
         throw new EncryptionException("The IV is too short, it should have at least length " + this.blockSizeInBytes());
      }

      AEADBlockCipher aeCipher = GCMBlockCipher.newInstance(createAESEngine());

      try {
         aeCipher.init(true, this.convertToCipherParameters(parameters));
         byte[] ciphertext = new byte[aeCipher.getOutputSize(plaintext.length)];
         int len = aeCipher.processBytes(plaintext, 0, plaintext.length, ciphertext, 0);
         aeCipher.doFinal(ciphertext, len);
         return CipherText.createWithPlainCiphertext(ciphertext).setCipherPrefix("GCM");
      } catch (CryptoException | IllegalArgumentException | CipherText.InvalidPrefixException e) {
         throw new EncryptionException("Cannot encrypt: " + e.getMessage(), e);
      }
   }

   @Override
   public byte[] decrypt(CipherText ciphertext, KeyAndIV parameters) throws DecryptionException {
      if (!this.isSuitablePrefix(ciphertext.getCipherPrefix())) {
         throw new DecryptionException("Not a suitable ciphertext, prefix wrong.");
      }

      AEADBlockCipher cipher = GCMBlockCipher.newInstance(createAESEngine());
      cipher.init(false, this.convertToCipherParameters(parameters));

      try {
         byte[] ciphertextBytes = ciphertext.getCiphertextPlain();
         byte[] plaintext = new byte[cipher.getOutputSize(ciphertextBytes.length)];
         int outputLen = cipher.processBytes(ciphertextBytes, 0, ciphertextBytes.length, plaintext, 0);
         outputLen += cipher.doFinal(plaintext, outputLen);
         byte[] output = new byte[outputLen];
         System.arraycopy(plaintext, 0, output, 0, outputLen);
         return output;
      } catch (DataLengthException | IllegalStateException | InvalidCipherTextException e) {
         throw new DecryptionException("failed to decrypt: " + e.getMessage(), e);
      }
   }

   @Override
   public boolean isSuitablePrefix(String prefix) {
      return "GCM".equals(prefix);
   }

   @Override
   public int blockSizeInBytes() {
      return this.computedBlockSizeInBytes;
   }

   private static MultiBlockCipher createAESEngine() {
      return AESEngine.newInstance();
   }

   private CipherParameters convertToCipherParameters(KeyAndIV keyAndIv) {
      return new ParametersWithIV(new KeyParameter(keyAndIv.getKey()), keyAndIv.getIV());
   }
}