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.CBCBlockCipher;
import org.bouncycastle.crypto.paddings.PaddedBufferedBlockCipher;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.ParametersWithIV;

public final class CBCAESCipher implements Cipher {
   private static final String CBC_CIPHER_PREFIX = "CBC";
   private final int blockSize = createAESEngine().getBlockSize();

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

   private CBCAESCipher() {
   }

   @Override
   public CipherText encrypt(byte[] plaintext, KeyAndIV parameters) throws EncryptionException {
      PaddedBufferedBlockCipher cipher = new PaddedBufferedBlockCipher(CBCBlockCipher.newInstance(createAESEngine()));

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

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

   @Override
   public byte[] decrypt(CipherText ciphertext, KeyAndIV parameters) throws DecryptionException {
      PaddedBufferedBlockCipher cipher = new PaddedBufferedBlockCipher(CBCBlockCipher.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 "CBC".equals(prefix);
   }

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

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