package dslab.crypto;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.logging.Logger;

import static java.util.Base64.getDecoder;
import static java.util.Base64.getEncoder;
import static javax.crypto.Cipher.DECRYPT_MODE;
import static javax.crypto.Cipher.ENCRYPT_MODE;

/**
 * Offers encryption and decryption of Strings via AES.
 */
public class AES {

    private static final Logger LOG = Logger.getLogger(AES.class.getSimpleName());
    private final SecretKey secretKey;

    private static final SecureRandom secureRandom = new SecureRandom();
    private final IvParameterSpec iv;
    private final byte[] clientChallenge;
    private Cipher encryptionCipher;
    private Cipher decryptionCipher;

    /**
     * For initialization on challenging side, i.e. client
     */
    public AES() {
        try {
            var aes = KeyGenerator.getInstance("AES");
            aes.init(256);
            secretKey = aes.generateKey();
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e); //never happens
        }

        var ivBytes = new byte[16];
        secureRandom.nextBytes(ivBytes);
        iv = new IvParameterSpec(ivBytes);

        clientChallenge = new byte[32];
        secureRandom.nextBytes(clientChallenge);

        initCiphers(new SecretKeySpec(secretKey.getEncoded(), "AES"), iv);
    }

    /**
     * For initialization on challenged side, i.e. server
     */
    public AES(ClientChallengeMessage challengeMessage) {
        clientChallenge = challengeMessage.getChallenge();
        secretKey = challengeMessage.getKey();
        iv = challengeMessage.getInitializationVector();

        initCiphers(challengeMessage.getKey(), challengeMessage.getInitializationVector());
    }

    private void initCiphers(SecretKeySpec key, IvParameterSpec initializationVector) {
        try {
            this.encryptionCipher = Cipher.getInstance("AES/CTR/NoPadding");
            encryptionCipher.init(ENCRYPT_MODE, key, initializationVector);

            this.decryptionCipher = Cipher.getInstance("AES/CTR/NoPadding");
            decryptionCipher.init(DECRYPT_MODE, key, initializationVector);

        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public String getBase64EncodedKey() {
        return getEncoder().encodeToString(secretKey.getEncoded());
    }

    public String getBase64ClientChallenge() {
        return getEncoder().encodeToString(clientChallenge);
    }

    public String getBase64InitializationVector() {
        return getEncoder().encodeToString(iv.getIV());
    }

    public String encrypt(String cleartext) {
        try {
            return getEncoder().encodeToString(encryptionCipher.doFinal(cleartext.getBytes()));
        } catch (IllegalBlockSizeException | BadPaddingException e) {
            throw new RuntimeException(e);
        }
    }

    public String decrypt(String ciphertext) {
        try {
            var base64Decoded = getDecoder().decode(ciphertext);
            var decryptedBytes = decryptionCipher.doFinal(base64Decoded);
            return new String(decryptedBytes);
        } catch (IllegalBlockSizeException | BadPaddingException e) {
            throw new RuntimeException(e);
        }
    }
}