package com.secureline.server;

import android.content.Context;
import android.os.Environment;
import android.security.keystore.KeyGenParameterSpec;
import android.security.keystore.KeyProperties;
import android.util.Slog;

import java.io.File;
import java.io.FileOutputStream;
import java.nio.ByteBuffer;
import java.security.KeyStore;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.GCMParameterSpec;

public class SecureLineIdentityStore {

    private static final String TAG = "SecureLineIdentityStore";

    private static final String BOOTSTRAP_ALIAS =
            "secureline_identity_bootstrap_key";

    private static final String SECURE_ALIAS =
            "secureline_identity_secure_key";

    private static final String TRANSFORMATION =
            "AES/GCM/NoPadding";

    private final File mIdentityFile;

    public SecureLineIdentityStore(Context context) {

        File dir =
                new File(
                        Environment.getDataSystemDeDirectory(0),
                        "secureline"
                );

        if (!dir.exists()) {

            if (!dir.mkdirs()) {

                throw new IllegalStateException(
                        "Failed creating DE secureline dir"
                );
            }

            dir.setReadable(false, false);
            dir.setWritable(true, true);
        }

        mIdentityFile =
                new File(dir, "identity.bin");
    }

    // ============================================================
    // existence
    // ============================================================

    public boolean identityExists() {

        if (!mIdentityFile.exists())
            return false;

        try {

            KeyStore ks =
                    KeyStore.getInstance("AndroidKeyStore");

            ks.load(null);

            return
                    ks.containsAlias(BOOTSTRAP_ALIAS)
                    || ks.containsAlias(SECURE_ALIAS);

        } catch (Throwable t) {

            Slog.e(TAG,
                    "identityExists error",
                    t);

            return false;
        }
    }

    // ============================================================
    // key creation
    // ============================================================

    private void generateBootstrapKeyIfNeeded()
            throws Exception {

        KeyStore ks =
                KeyStore.getInstance("AndroidKeyStore");

        ks.load(null);

        if (ks.containsAlias(BOOTSTRAP_ALIAS))
            return;

        KeyGenerator kg =
                KeyGenerator.getInstance(
                        KeyProperties.KEY_ALGORITHM_AES,
                        "AndroidKeyStore"
                );

        kg.init(
                new KeyGenParameterSpec.Builder(
                        BOOTSTRAP_ALIAS,
                        KeyProperties.PURPOSE_ENCRYPT
                        | KeyProperties.PURPOSE_DECRYPT
                )
                        .setBlockModes(
                                KeyProperties.BLOCK_MODE_GCM
                        )
                        .setEncryptionPaddings(
                                KeyProperties.ENCRYPTION_PADDING_NONE
                        )
                        .setUnlockedDeviceRequired(false)
                        .setRandomizedEncryptionRequired(true)
                        .setUserAuthenticationRequired(false)
                        .setKeySize(256)
                        .build()
        );

        kg.generateKey();

        Slog.i(TAG,
                "Bootstrap identity key generated");
    }

    private void generateSecureKeyIfNeeded()
            throws Exception {

        KeyStore ks =
                KeyStore.getInstance("AndroidKeyStore");

        ks.load(null);

        if (ks.containsAlias(SECURE_ALIAS))
            return;

        KeyGenerator kg =
                KeyGenerator.getInstance(
                        KeyProperties.KEY_ALGORITHM_AES,
                        "AndroidKeyStore"
                );

        kg.init(
                new KeyGenParameterSpec.Builder(
                        SECURE_ALIAS,
                        KeyProperties.PURPOSE_ENCRYPT
                        | KeyProperties.PURPOSE_DECRYPT
                )
                        .setBlockModes(
                                KeyProperties.BLOCK_MODE_GCM
                        )
                        .setEncryptionPaddings(
                                KeyProperties.ENCRYPTION_PADDING_NONE
                        )
                        .setUnlockedDeviceRequired(true)
                        .setRandomizedEncryptionRequired(true)
                        .setUserAuthenticationRequired(false)
                        .setKeySize(256)
                        .build()
        );

        kg.generateKey();

        Slog.i(TAG,
                "Secure identity key generated");
    }

    // ============================================================
    // store identity (BFU compatible)
    // ============================================================

    public void storeIdentity(byte[] plaintext)
            throws Exception {

        generateBootstrapKeyIfNeeded();

        encryptAtomic(
                plaintext,
                getKey(BOOTSTRAP_ALIAS)
        );

        Slog.i(TAG,
                "Identity stored (bootstrap)");
    }

    // ============================================================
    // load identity
    // ============================================================

    public byte[] loadIdentity()
            throws Exception {

        if (!mIdentityFile.exists())
            return null;

        KeyStore ks =
                KeyStore.getInstance("AndroidKeyStore");

        ks.load(null);

        SecretKey key =
                ks.containsAlias(SECURE_ALIAS)
                ? getKey(SECURE_ALIAS)
                : getKey(BOOTSTRAP_ALIAS);

        return decrypt(key);
    }

    // ============================================================
    // migrate after first unlock (CE protection)
    // ============================================================

    public void migrateIdentityAfterFirstUnlock() {

        try {

            KeyStore ks =
                    KeyStore.getInstance("AndroidKeyStore");

            ks.load(null);

            if (!ks.containsAlias(BOOTSTRAP_ALIAS))
                return;

            byte[] plaintext =
                    decrypt(getKey(BOOTSTRAP_ALIAS));

            generateSecureKeyIfNeeded();

            encryptAtomic(
                    plaintext,
                    getKey(SECURE_ALIAS)
            );

            if (!ks.containsAlias(SECURE_ALIAS)) {

                throw new IllegalStateException(
                        "Secure key missing after migration"
                );
            }

            ks.deleteEntry(BOOTSTRAP_ALIAS);

            Slog.i(TAG,
                    "Identity migrated to CE key");

        } catch (Throwable t) {

            Slog.e(TAG,
                    "Identity migration failed",
                    t);
        }
    }

    // ============================================================
    // crypto helpers
    // ============================================================

    private SecretKey getKey(String alias)
            throws Exception {

        KeyStore ks =
                KeyStore.getInstance("AndroidKeyStore");

        ks.load(null);

        return (SecretKey)
                ks.getKey(alias, null);
    }

    private void encryptAtomic(
            byte[] plaintext,
            SecretKey key
    ) throws Exception {

        Cipher cipher =
                Cipher.getInstance(TRANSFORMATION);

        cipher.init(
                Cipher.ENCRYPT_MODE,
                key
        );

        byte[] iv =
                cipher.getIV();

        byte[] enc =
                cipher.doFinal(plaintext);

        ByteBuffer buf =
                ByteBuffer.allocate(
                        4 + iv.length + enc.length
                );

        buf.putInt(iv.length);
        buf.put(iv);
        buf.put(enc);

        File tmp =
                new File(
                        mIdentityFile.getAbsolutePath() + ".tmp"
                );

        try (FileOutputStream fos =
                     new FileOutputStream(tmp)) {

            fos.write(buf.array());

            fos.flush();

            fos.getFD().sync();
        }

        if (!tmp.renameTo(mIdentityFile)) {

            throw new IllegalStateException(
                    "identity rename failed"
            );
        }
    }

    private byte[] decrypt(SecretKey key)
            throws Exception {

        byte[] data =
                java.nio.file.Files.readAllBytes(
                        mIdentityFile.toPath()
                );

        ByteBuffer buf =
                ByteBuffer.wrap(data);

        int ivLen =
                buf.getInt();

        if (ivLen != 12)
            throw new SecurityException(
                    "invalid IV"
            );

        byte[] iv =
                new byte[ivLen];

        buf.get(iv);

        byte[] enc =
                new byte[buf.remaining()];

        buf.get(enc);

        Cipher cipher =
                Cipher.getInstance(TRANSFORMATION);

        cipher.init(
                Cipher.DECRYPT_MODE,
                key,
                new GCMParameterSpec(128, iv)
        );

        return cipher.doFinal(enc);
    }

    public void deleteIdentity() {

    try {

        if (mIdentityFile.exists()) {

            if (!mIdentityFile.delete()) {

                Slog.w(TAG,
                        "identity delete failed");
            }
        }

        KeyStore ks =
                KeyStore.getInstance("AndroidKeyStore");

        ks.load(null);

        if (ks.containsAlias(BOOTSTRAP_ALIAS)) {

            ks.deleteEntry(BOOTSTRAP_ALIAS);
        }

        if (ks.containsAlias(SECURE_ALIAS)) {

            ks.deleteEntry(SECURE_ALIAS);
        }

        Slog.w(TAG,
                "identity deleted");

    } catch (Throwable t) {

        Slog.e(TAG,
                "identity delete error",
                t);
    }
}

}