package com.sangcx.keystoretest;

import android.content.Context;
import android.content.SharedPreferences;
import android.os.Bundle;
import android.security.keystore.KeyGenParameterSpec;
import android.security.keystore.KeyProperties;
import android.util.Base64;

import androidx.appcompat.app.AppCompatActivity;

import java.io.IOException;
import java.security.InvalidAlgorithmParameterException;
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.PrivateKey;
import java.security.PublicKey;
import java.security.UnrecoverableEntryException;
import java.security.cert.CertificateException;

import javax.crypto.Cipher;

public class KeystoreTest extends AppCompatActivity {

    public static class KeyStoreHelper {
        private static final String KEY_ALIAS = "my_key_alias";
        private static final String ANDROID_KEYSTORE = "AndroidKeyStore";

        public static void generateKeyPair() throws NoSuchAlgorithmException, NoSuchProviderException, InvalidAlgorithmParameterException {
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(KeyProperties.KEY_ALGORITHM_RSA, ANDROID_KEYSTORE);
            keyPairGenerator.initialize(
                    new KeyGenParameterSpec.Builder(KEY_ALIAS, KeyProperties.PURPOSE_ENCRYPT | KeyProperties.PURPOSE_DECRYPT)
                            .setEncryptionPaddings(KeyProperties.ENCRYPTION_PADDING_RSA_PKCS1)
                            .build());
            keyPairGenerator.generateKeyPair();
        }

        public static KeyPair getKeyPair() throws KeyStoreException, CertificateException, NoSuchAlgorithmException, IOException, UnrecoverableEntryException {
            KeyStore keyStore = KeyStore.getInstance(ANDROID_KEYSTORE);
            keyStore.load(null);
            KeyStore.PrivateKeyEntry privateKeyEntry = (KeyStore.PrivateKeyEntry) keyStore.getEntry(KEY_ALIAS, null);
            if (privateKeyEntry != null) {
                return new KeyPair(privateKeyEntry.getCertificate().getPublicKey(), privateKeyEntry.getPrivateKey());
            }
            return null;
        }
    }

    public static class EncryptionHelper {

        public static byte[] encryptData(String data, PublicKey publicKey) throws Exception {
            Cipher cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
            cipher.init(Cipher.ENCRYPT_MODE, publicKey);
            return cipher.doFinal(data.getBytes("UTF-8"));
        }

        public static String decryptData(byte[] encryptedData, PrivateKey privateKey) throws Exception {
            Cipher cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
            cipher.init(Cipher.DECRYPT_MODE, privateKey);
            byte[] decryptedData = cipher.doFinal(encryptedData);
            return new String(decryptedData, "UTF-8");
        }
    }

    public static class StorageHelper {

        private static final String PREFS_NAME = "my_prefs";
        private static final String ENCRYPTED_DATA_KEY = "encrypted_data";

        public static void saveEncryptedData(Context context, byte[] encryptedData) {
            SharedPreferences sharedPreferences = context.getSharedPreferences(PREFS_NAME, Context.MODE_PRIVATE);
            SharedPreferences.Editor editor = sharedPreferences.edit();
            String encryptedDataString = Base64.encodeToString(encryptedData, Base64.DEFAULT);
            editor.putString(ENCRYPTED_DATA_KEY, encryptedDataString);
            editor.apply();
        }

        public static String getEncryptedData(Context context) {
            SharedPreferences sharedPreferences = context.getSharedPreferences(PREFS_NAME, Context.MODE_PRIVATE);
            return sharedPreferences.getString(ENCRYPTED_DATA_KEY, null);
        }
    }

    @Override
    protected void onCreate(Bundle savedInstanceState) {
        super.onCreate(savedInstanceState);
        setContentView(R.layout.activity_main);

        try {
            KeyPair keyPair = KeyStoreHelper.getKeyPair();
            if (keyPair == null) {
                KeyStoreHelper.generateKeyPair();
                keyPair = KeyStoreHelper.getKeyPair();
            }
            String dataToEncrypt = "abcdef";
            byte[] encryptedData = EncryptionHelper.encryptData(dataToEncrypt, keyPair.getPublic());
            StorageHelper.saveEncryptedData(this, encryptedData);
            String encryptedDataString = StorageHelper.getEncryptedData(this);
            byte[] encryptedDataRetrieved = Base64.decode(encryptedDataString, Base64.DEFAULT);
            String decryptedData = EncryptionHelper.decryptData(encryptedDataRetrieved, keyPair.getPrivate());
            System.out.println("Decrypted Data: " + decryptedData);

        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}