package be.iminds.ilabt.jfed.util;

import org.slf4j.LoggerFactory;

import javax.annotation.Nullable;
import javax.crypto.*;
import javax.crypto.spec.SecretKeySpec;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.UnsupportedEncodingException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.security.PrivateKey;
import java.security.PublicKey;

/**
 * FileEncrypter
 */
public class FileEncrypter {
    private static final org.slf4j.Logger LOG = LoggerFactory.getLogger(FileEncrypter.class);


    private static final int AES_KEYSIZE_BIT = 128;
            //256bit is more secure, BUT not included by default in java. users would

    // need to install the "unlimited strength file". See also
    // http://stackoverflow.com/questions/6481627/java-security-illegal-key-size-or-default-parameters
    @Nullable
    public static byte[] decryptData(byte[] data, PrivateKey privateKey) {
        try {
            ByteBuffer bb = ByteBuffer.wrap(data).order(ByteOrder.BIG_ENDIAN);
            int keySize = bb.getInt();
            byte[] encryptedKey = new byte[keySize];
            byte[] encryptedData = new byte[data.length - (keySize + 4)];
            assert bb.remaining() == data.length - 4;
            bb.get(encryptedKey);
            assert bb.remaining() == data.length - (keySize + 4);
            bb.get(encryptedData);

            //data is in encryptedKey and encryptedData now
            Cipher pkCipher = Cipher.getInstance("RSA");
            pkCipher.init(Cipher.DECRYPT_MODE, privateKey);
            ByteArrayInputStream baisKey = new ByteArrayInputStream(encryptedKey);
            CipherInputStream keyIs = new CipherInputStream(baisKey, pkCipher);
            byte[] aesKey = new byte[AES_KEYSIZE_BIT / 8];
            final int read = keyIs.read(aesKey);
            assert read == aesKey.length;
            SecretKeySpec aeskeySpec = new SecretKeySpec(aesKey, "AES");

            ByteArrayInputStream bais = new ByteArrayInputStream(encryptedData);
            Cipher aesCipher = Cipher.getInstance("AES");
            aesCipher.init(Cipher.DECRYPT_MODE, aeskeySpec);
            CipherInputStream is = new CipherInputStream(bais, aesCipher);
            ByteArrayOutputStream baos = new ByteArrayOutputStream();

            int i;
            byte[] b = new byte[1024];
            while ((i = is.read(b)) != -1) baos.write(b, 0, i);

            return baos.toByteArray();
        } catch (Exception e) {
            LOG.error("Failed to decrypt data: " + e.getMessage(), e);
            return null;
        }
    }

    public static byte[] encryptData(byte[] data, PublicKey publicKey) {
        try {
            KeyGenerator kgen = KeyGenerator.getInstance("AES");
            kgen.init(AES_KEYSIZE_BIT);
            SecretKey key = kgen.generateKey();
            byte[] aesKey = key.getEncoded();
            assert aesKey.length == AES_KEYSIZE_BIT / 8;
            SecretKeySpec aeskeySpec = new SecretKeySpec(aesKey, "AES");

            Cipher pkCipher = Cipher.getInstance("RSA");
            pkCipher.init(Cipher.ENCRYPT_MODE, publicKey);
            ByteArrayOutputStream keyBaos = new ByteArrayOutputStream();
            CipherOutputStream keyCipherOutputStream = new CipherOutputStream(keyBaos, pkCipher);
            keyCipherOutputStream.write(aesKey);
            keyCipherOutputStream.close();
            byte[] encryptedKeyBytes = keyBaos.toByteArray();

            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            byte[] lenBytes = new byte[4];
            ByteBuffer bb = ByteBuffer.wrap(lenBytes).order(ByteOrder.BIG_ENDIAN);
            bb.clear();
            bb.putInt(encryptedKeyBytes.length);

            baos.write(lenBytes);
            baos.write(encryptedKeyBytes);

            Cipher aesCipher = Cipher.getInstance("AES");
            aesCipher.init(Cipher.ENCRYPT_MODE, aeskeySpec);
            CipherOutputStream cipherOutputStream = new CipherOutputStream(baos, aesCipher);
            cipherOutputStream.write(data);
            cipherOutputStream.close();
            baos.close();

            return baos.toByteArray();
        } catch (Exception e) {
            LOG.error("Failed to encrypt data: " + e.getMessage(), e);
            return null;
        }
    }

    public static String decryptString(byte[] input, PrivateKey privateKey) {
        byte[] data = decryptData(input, privateKey);
        if (data == null)
            return null;
        try {
            return new String(data, "UTF-8");
        } catch (UnsupportedEncodingException e) {
            throw new RuntimeException("Unexpected: UTF-8 not supported", e);
        }
    }

    public static byte[] encryptString(String input, PublicKey publicKey) {
        try {
            return encryptData(input.getBytes("UTF-8"), publicKey);
        } catch (UnsupportedEncodingException e) {
            throw new RuntimeException("Unexpected: UTF-8 not supported", e);
        }
    }
}
