package com.adi.plugin;

import android.content.Context;
import android.security.KeyPairGeneratorSpec;
import android.util.Base64;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.math.BigInteger;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.UnrecoverableEntryException;
import java.security.cert.CertificateException;
import java.security.interfaces.RSAPublicKey;
import java.util.ArrayList;
import java.util.Calendar;

import javax.crypto.Cipher;
import javax.crypto.CipherInputStream;
import javax.crypto.CipherOutputStream;
import javax.crypto.NoSuchPaddingException;
import javax.security.auth.x500.X500Principal;

/**
 * Wrapper for Android Keystore Functionality
 */
public final class KeyStorage
{
    private static final String TAG = KeyStorage.class.getSimpleName();
    private static final String CIPHER_TYPE = "RSA/ECB/PKCS1Padding";
    private static final String ANDROID_KEY_STORE = "AndroidKeyStore";
    private static final String DISTINGUISHED_NAME = "CN=Sample Name, O=Android Authority";
    private static final String RSA = "RSA";
    private static final String UTF_8 = "UTF-8";
    private static final int CERT_YEARS = 25;

    private final RSAPublicKey publicKey;
    private final Key privateKey;

    /**
     * Constructor - checks if the given alias already exists and creates it if not.
     * Initialises public and private keys generated by keystore
     *
     * @param context   Context used by the wrapper
     * @param alias     Keystore Alias to use
     */
    public KeyStorage(final Context context, final String alias, final boolean encryptionRequired) throws KeyStoreException, CertificateException,
            NoSuchAlgorithmException, IOException, NoSuchProviderException,
            InvalidAlgorithmParameterException, UnrecoverableEntryException
    {
        final KeyStore keyStore = KeyStore.getInstance(ANDROID_KEY_STORE);
        keyStore.load(null);

        if (!keyStore.containsAlias(alias))
        {
            final Calendar start = Calendar.getInstance();
            final Calendar end = Calendar.getInstance();
            end.add(Calendar.YEAR, CERT_YEARS);
            final KeyPairGeneratorSpec.Builder builder = new KeyPairGeneratorSpec.Builder(context)
                    .setAlias(alias)
                    .setSubject(new X500Principal(DISTINGUISHED_NAME))
                    .setSerialNumber(BigInteger.ONE)
                    .setStartDate(start.getTime())
                    .setEndDate(end.getTime());
            if (encryptionRequired)
            {
                builder.setEncryptionRequired();
            }
            final KeyPairGeneratorSpec spec = builder.build();
            final KeyPairGenerator generator = KeyPairGenerator.getInstance(RSA, ANDROID_KEY_STORE);
            generator.initialize(spec);

            final KeyPair keyPair = generator.generateKeyPair();
            publicKey = (RSAPublicKey)keyPair.getPublic();
            privateKey = keyPair.getPrivate();
        }
        else
        {
            final KeyStore.PrivateKeyEntry privateKeyEntry = (KeyStore.PrivateKeyEntry) keyStore.getEntry(alias, null);
            publicKey = (RSAPublicKey) privateKeyEntry.getCertificate().getPublicKey();
            privateKey = privateKeyEntry.getPrivateKey();
        }
    }

    /**
     * Encrypt the given string value using the public key from the Keystore
     *
     * @param value     The string to encrypt
     * @return          The encrypted string
     */
    public String encryptString(final String value)
    {
        try
        {
            final Cipher inCipher = Cipher.getInstance(CIPHER_TYPE);
            inCipher.init(Cipher.ENCRYPT_MODE, publicKey);

            final ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
            final CipherOutputStream cipherOutputStream = new CipherOutputStream(
                    outputStream, inCipher);
            cipherOutputStream.write(value.getBytes(UTF_8));
            cipherOutputStream.close();

            final byte[] vals = outputStream.toByteArray();
            return Base64.encodeToString(vals, Base64.DEFAULT);
        }
        catch(NoSuchPaddingException e)
        {
            throw new IllegalArgumentException(e);
        }
        catch(NoSuchAlgorithmException e)
        {
            throw new IllegalArgumentException(e);
        }
        catch(InvalidKeyException e)
        {
            throw new IllegalArgumentException(e);
        }
        catch(IOException e)
        {
            throw new IllegalArgumentException(e);
        }
    }

    /**
     * Decrypt the given string using the private key from the Keystore
     *
     * @param cipherText    The Encrypted string
     * @return              The Decrypted string
     */
    public String decryptString(final String cipherText)
    {
        try
        {
            Cipher output = Cipher.getInstance(CIPHER_TYPE);
            output.init(Cipher.DECRYPT_MODE, privateKey);

            CipherInputStream cipherInputStream = new CipherInputStream(
                    new ByteArrayInputStream(Base64.decode(cipherText, Base64.DEFAULT)), output);
            ArrayList<Byte> values = new ArrayList<Byte>();
            int nextByte;
            while ((nextByte = cipherInputStream.read()) != -1)
            {
                values.add((byte) nextByte);
            }

            byte[] bytes = new byte[values.size()];
            for (int i = 0; i < bytes.length; i++)
            {
                bytes[i] = values.get(i);
            }

            return new String(bytes, 0, bytes.length, UTF_8);
        }
        catch(NoSuchPaddingException e)
        {
            throw new IllegalArgumentException(e);
        }
        catch(NoSuchAlgorithmException e)
        {
            throw new IllegalArgumentException(e);
        }
        catch(InvalidKeyException e)
        {
            throw new IllegalArgumentException(e);
        }
        catch(IOException e)
        {
            throw new IllegalArgumentException(e);
        }
    }
}