package cz.vutbr.benchmark

import android.security.keystore.KeyGenParameterSpec
import android.security.keystore.KeyProperties
import androidx.benchmark.junit4.BenchmarkRule
import androidx.benchmark.junit4.measureRepeated
import androidx.test.ext.junit.runners.AndroidJUnit4
import org.junit.Rule
import org.junit.Test
import org.junit.runner.RunWith
import java.security.KeyPairGenerator


/**
 * Benchmark, which will execute on an Android device.
 *
 * The body of [BenchmarkRule.measureRepeated] is measured in a loop, and Studio will
 * output the result. Modify your code to see how it affects performance.
 */
@RunWith(AndroidJUnit4::class)
class KeyPairGenerationBenchmark {


    @get:Rule
    val benchmarkRule = BenchmarkRule()

    @Test
    fun benchmarkRSA512() {
        benchmarkRSA(512)
    }

    @Test
    fun benchmarkRSA768() {
        benchmarkRSA(768)
    }

    @Test
    fun benchmarkRSA1024() {
        benchmarkRSA(1024)
    }

    @Test
    fun benchmarkRSA2048() {
        benchmarkRSA(2048)
    }

    @Test
    fun benchmarkRSA3072() {
        benchmarkRSA(3072)
    }

    //@Test
    fun benchmarkRSA4096() {
        benchmarkRSA(4096)
    }

    @Test
    fun benchmarkEC224() {
        benchmarkEC(224)
    }

    @Test
    fun benchmarkEC256() {
        benchmarkEC(256)
    }

    @Test
    fun benchmarkEC384() {
        benchmarkEC(384)
    }

    @Test
    fun benchmarkEC521() {
        benchmarkEC(521)
    }

    private fun benchmarkRSA(keySize: Int) {
        val kpg = KeyPairGenerator.getInstance(
            KeyProperties.KEY_ALGORITHM_RSA,
            "AndroidKeyStore"
        )
        val parameterSpec: KeyGenParameterSpec =
            KeyGenParameterSpec.Builder(
                "RSA_KEY",
                KeyProperties.PURPOSE_ENCRYPT
                        or
                        KeyProperties.PURPOSE_DECRYPT
            )
                .setKeySize(keySize)
                .build()
        kpg.apply { initialize(parameterSpec) }
        benchmarkRule.measureRepeated {
            kpg.generateKeyPair()
        }
    }

    private fun benchmarkEC(keySize: Int) {
        val kpg = KeyPairGenerator.getInstance(
            KeyProperties.KEY_ALGORITHM_EC,
            "AndroidKeyStore"
        )
        val parameterSpec: KeyGenParameterSpec =
            KeyGenParameterSpec.Builder(
                "EC_KEY",
                KeyProperties.PURPOSE_ENCRYPT
                        or
                        KeyProperties.PURPOSE_DECRYPT
            )
                .setKeySize(keySize)
                .build()
        kpg.apply { initialize(parameterSpec) }
        benchmarkRule.measureRepeated {
            kpg.generateKeyPair()
        }
    }

}
