import javax.crypto.*;
import javax.crypto.spec.GCMParameterSpec;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.util.Base64;

public class AES {

    public String encrypt(String msg, SecretKey key, byte[] iv) {

        byte[] cipher;

        try {
            Cipher encryptionCipher = Cipher.getInstance("AES/GCM/NoPadding");

            GCMParameterSpec parameterSpec = new GCMParameterSpec(128, iv);

            System.out.format("Key bits: %d\n", key.getEncoded().length * 8);
            encryptionCipher.init(Cipher.ENCRYPT_MODE, key, parameterSpec);

            byte[] bytes = msg.getBytes();

            cipher = encryptionCipher.doFinal(bytes);
            ByteBuffer byteBuffer = ByteBuffer.allocate(4 + iv.length + cipher.length);
            byteBuffer.putInt(iv.length);
            byteBuffer.put(iv);
            byteBuffer.put(cipher);
            cipher = byteBuffer.array();
            cipher = Base64.getEncoder().encode(cipher);

            return new String(cipher, StandardCharsets.UTF_8);
        } catch (NoSuchPaddingException | NoSuchAlgorithmException | InvalidKeyException | IllegalBlockSizeException | BadPaddingException e) {

            System.out.format("Error: %s\n", e.getMessage());
            e.printStackTrace();
        } catch (InvalidAlgorithmParameterException e) {
            e.printStackTrace();
        }

        return null;
    }

    public String decrypt(String cipher, SecretKey key) {
        byte[] msg;

        try {
            byte[] bytes = Base64.getDecoder().decode(cipher);
            ByteBuffer byteBuffer = ByteBuffer.wrap(bytes);
            int ivLength = byteBuffer.getInt();
            if(ivLength < 12 || ivLength >= 16) { // check input parameter
                throw new IllegalArgumentException("invalid iv length");
            }
            byte[] iv = new byte[ivLength];
            byteBuffer.get(iv);
            byte[] cipherText = new byte[byteBuffer.remaining()];
            byteBuffer.get(cipherText);


            Cipher decryptionCipher = Cipher.getInstance("AES/GCM/NoPadding");
            GCMParameterSpec parameterSpec = new GCMParameterSpec(128, iv);
            decryptionCipher.init(Cipher.DECRYPT_MODE, key, parameterSpec);

            msg = decryptionCipher.doFinal(cipherText);

            return new String(msg, StandardCharsets.UTF_8);
        } catch (NoSuchPaddingException | NoSuchAlgorithmException | InvalidKeyException | BadPaddingException | IllegalBlockSizeException | InvalidAlgorithmParameterException e) {
            e.printStackTrace();
        }

        return null;
    }
}