package hsr.dsa.utils;

import java.nio.charset.StandardCharsets;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.Base64;
import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

public class AES {

  // 256bit needs extra encryption library from oracle Java Cryptography Extension
  // (JCE) Unlimited Strength
  // https://www.oracle.com/technetwork/java/javase/downloads/jce8-download-2133166.html
  private static final int AES_KEY_SIZE = 16; // 16*8 = 128 ... 32*8 = 256
  private static final int IV_SIZE = 16;
  private static final String ALGO = "AES/CBC/PKCS5Padding";

  public AES() {
  }

  public byte[] encrypt(byte[] clean, String key)
      throws NoSuchAlgorithmException, NoSuchPaddingException, BadPaddingException,
        IllegalBlockSizeException, InvalidAlgorithmParameterException, InvalidKeyException {

    // Generating IV.
    byte[] iv = new byte[IV_SIZE];
    SecureRandom random = new SecureRandom();
    random.nextBytes(iv);
    IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);

    // Hashing key.
    MessageDigest digest = MessageDigest.getInstance("SHA-256");
    digest.update(key.getBytes(StandardCharsets.UTF_8));
    byte[] keyBytes = new byte[AES_KEY_SIZE];
    System.arraycopy(digest.digest(), 0, keyBytes, 0, keyBytes.length);
    SecretKeySpec secretKeySpec = new SecretKeySpec(keyBytes, "AES");

    // Encrypt.
    Cipher cipher = Cipher.getInstance(ALGO);
    cipher.init(Cipher.ENCRYPT_MODE, secretKeySpec, ivParameterSpec);
    byte[] encrypted = cipher.doFinal(clean);

    // Combine IV and encrypted part.
    byte[] encryptedIVAndText = new byte[IV_SIZE + encrypted.length];
    System.arraycopy(iv, 0, encryptedIVAndText, 0, IV_SIZE);
    System.arraycopy(encrypted, 0, encryptedIVAndText, IV_SIZE, encrypted.length);

    return encryptedIVAndText;
  }

  public String encrypt(String plainText, String key)
      throws NoSuchPaddingException, InvalidKeyException, NoSuchAlgorithmException,
        IllegalBlockSizeException, BadPaddingException, InvalidAlgorithmParameterException {
    return Base64.getEncoder()
        .encodeToString(this.encrypt(plainText.getBytes(StandardCharsets.UTF_8), key));
  }

  public byte[] decrypt(byte[] encryptedIvTextBytes, String key)
      throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidAlgorithmParameterException,
        InvalidKeyException, BadPaddingException, IllegalBlockSizeException {
    // Extract IV.
    byte[] iv = new byte[IV_SIZE];
    System.arraycopy(encryptedIvTextBytes, 0, iv, 0, iv.length);
    IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);

    // Extract encrypted part.
    int encryptedSize = encryptedIvTextBytes.length - IV_SIZE;
    byte[] encryptedBytes = new byte[encryptedSize];
    System.arraycopy(encryptedIvTextBytes, IV_SIZE, encryptedBytes, 0, encryptedSize);

    // Hash key.
    byte[] keyBytes = new byte[AES_KEY_SIZE];
    MessageDigest md = MessageDigest.getInstance("SHA-256");
    md.update(key.getBytes());
    System.arraycopy(md.digest(), 0, keyBytes, 0, keyBytes.length);
    SecretKeySpec secretKeySpec = new SecretKeySpec(keyBytes, "AES");

    // Decrypt.
    Cipher cipherDecrypt = Cipher.getInstance(ALGO);
    cipherDecrypt.init(Cipher.DECRYPT_MODE, secretKeySpec, ivParameterSpec);
    return cipherDecrypt.doFinal(encryptedBytes);
  }

  public String decrypt(String encryptedIvTextBytes, String key)
      throws NoSuchPaddingException, InvalidKeyException, NoSuchAlgorithmException,
      IllegalBlockSizeException, BadPaddingException, InvalidAlgorithmParameterException {
    return new String(this.decrypt(Base64.getDecoder().decode(encryptedIvTextBytes), key));
  }
}
