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.Assert
import org.junit.Rule
import org.junit.Test
import org.junit.runner.RunWith
import java.security.KeyPairGenerator
import java.security.KeyStore
import java.security.PublicKey
import java.security.spec.MGF1ParameterSpec
import javax.crypto.Cipher
import javax.crypto.spec.OAEPParameterSpec
import javax.crypto.spec.PSource


@RunWith(AndroidJUnit4::class)
class AsymmetricEncryptionBenchmark {

    @get:Rule
    val benchmarkRule = BenchmarkRule()

    val testFile = "qwertyuiopasdfghjklzxcvbnm123456789".toByteArray()

    @Test
    fun benchmarkRsa512EcbPKCS1Padding() {
        benchmarkRsa(
            "RSA/ECB/PKCS1Padding",
            512,
            KeyProperties.ENCRYPTION_PADDING_RSA_PKCS1
        )
    }

    @Test
    fun benchmarkRsa768EcbPKCS1Padding() {
        benchmarkRsa(
            "RSA/ECB/PKCS1Padding",
            768,
            KeyProperties.ENCRYPTION_PADDING_RSA_PKCS1
        )
    }

    @Test
    fun benchmarkRsa1024EcbPKCS1Padding() {
        benchmarkRsa(
            "RSA/ECB/PKCS1Padding",
            1024,
            KeyProperties.ENCRYPTION_PADDING_RSA_PKCS1
        )
    }

    @Test
    fun benchmarkRsa2048EcbPKCS1Padding() {
        benchmarkRsa(
            "RSA/ECB/PKCS1Padding",
            2048,
            KeyProperties.ENCRYPTION_PADDING_RSA_PKCS1
        )
    }

    @Test
    fun benchmarkRsa3072EcbPKCS1Padding() {
        benchmarkRsa(
            "RSA/ECB/PKCS1Padding",
            3072,
            KeyProperties.ENCRYPTION_PADDING_RSA_PKCS1
        )
    }

    @Test
    fun benchmarkRsa4096EcbPKCS1Padding() {
        benchmarkRsa(
            "RSA/ECB/PKCS1Padding",
            4096,
            KeyProperties.ENCRYPTION_PADDING_RSA_PKCS1
        )
    }

    //@Test
    fun benchmarkRsa512EcbOaep() {
        benchmarkRsa(
            "RSA/ECB/OAEPPadding",
            512,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA1
        )
    }

    @Test
    fun benchmarkRsa768EcbOaep() {
        benchmarkRsa(
            "RSA/ECB/OAEPPadding",
            768,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA1
        )
    }

    @Test
    fun benchmarkRsa1024EcbOaep() {
        benchmarkRsa(
            "RSA/ECB/OAEPPadding",
            1024,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA1
        )
    }

    @Test
    fun benchmarkRsa2048EcbOaep() {
        benchmarkRsa(
            "RSA/ECB/OAEPPadding",
            2048,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA1
        )
    }

    @Test
    fun benchmarkRsa3072EcbOaep() {
        benchmarkRsa(
            "RSA/ECB/OAEPPadding",
            3072,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA1
        )
    }

    @Test
    fun benchmarkRsa4096EcbOaep() {
        benchmarkRsa(
            "RSA/ECB/OAEPPadding",
            4096,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA1
        )
    }

    //@Test
    fun benchmarkRsa512EcbOaepSHA1AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-1AndMGF1Padding",
            512,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA1,
            OAEPParameterSpec(
                "SHA-1",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    @Test
    fun benchmarkRsa1024EcbOaepSHA1AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-1AndMGF1Padding",
            1024,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA1,
            OAEPParameterSpec(
                "SHA-1",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    @Test
    fun benchmarkRsa2048EcbOaepSHA1AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-1AndMGF1Padding",
            2048,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA1,
            OAEPParameterSpec(
                "SHA-1",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    @Test
    fun benchmarkRsa3072EcbOaepSHA1AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-1AndMGF1Padding",
            3072,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA1,
            OAEPParameterSpec(
                "SHA-1",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    @Test
    fun benchmarkRsa4096EcbOaepSHA1AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-1AndMGF1Padding",
            4096,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA1,
            OAEPParameterSpec(
                "SHA-1",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    //@Test
    fun benchmarkRsa512EcbOaepSHA224AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-224AndMGF1Padding",
            512,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA224,
            OAEPParameterSpec(
                "SHA-224",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    @Test
    fun benchmarkRsa1024EcbOaepSHA224AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-224AndMGF1Padding",
            1024,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA224,
            OAEPParameterSpec(
                "SHA-224",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    @Test
    fun benchmarkRsa2048EcbOaepSHA224AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-224AndMGF1Padding",
            2048,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA224,
            OAEPParameterSpec(
                "SHA-224",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    @Test
    fun benchmarkRsa3072EcbOaepSHA224AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-224AndMGF1Padding",
            3072,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA224,
            OAEPParameterSpec(
                "SHA-224",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    @Test
    fun benchmarkRsa4096EcbOaepSHA224AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-224AndMGF1Padding",
            4096,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA224,
            OAEPParameterSpec(
                "SHA-224",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    //@Test
    fun benchmarkRsa512EcbOaepSHA256AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-256AndMGF1Padding",
            512,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA256,
            OAEPParameterSpec(
                "SHA-256",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    @Test
    fun benchmarkRsa1024EcbOaepSHA256AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-256AndMGF1Padding",
            1024,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA256,
            OAEPParameterSpec(
                "SHA-256",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    @Test
    fun benchmarkRsa2048EcbOaepSHA256AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-256AndMGF1Padding",
            2048,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA256,
            OAEPParameterSpec(
                "SHA-256",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    @Test
    fun benchmarkRsa3072EcbOaepSHA256AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-256AndMGF1Padding",
            3072,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA256,
            OAEPParameterSpec(
                "SHA-256",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    @Test
    fun benchmarkRsa4096EcbOaepSHA256AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-256AndMGF1Padding",
            4096,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA256,
            OAEPParameterSpec(
                "SHA-256",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    //@Test
    fun benchmarkRsa512EcbOaepSHA384AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-384AndMGF1Padding",
            512,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA384,
            OAEPParameterSpec(
                "SHA-384",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    //@Test
    fun benchmarkRsa1024EcbOaepSHA384AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-384AndMGF1Padding",
            1024,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA384,
            OAEPParameterSpec(
                "SHA-384",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    @Test
    fun benchmarkRsa2048EcbOaepSHA384AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-384AndMGF1Padding",
            2048,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA384,
            OAEPParameterSpec(
                "SHA-384",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    @Test
    fun benchmarkRsa3072EcbOaepSHA384AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-384AndMGF1Padding",
            3072,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA384,
            OAEPParameterSpec(
                "SHA-384",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    @Test
    fun benchmarkRsa4096EcbOaepSHA384AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-384AndMGF1Padding",
            4096,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA384,
            OAEPParameterSpec(
                "SHA-384",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    //@Test
    fun benchmarkRsa512EcbOaepSHA512AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-512AndMGF1Padding",
            512,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA512,
            OAEPParameterSpec(
                "SHA-512",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    //@Test
    fun benchmarkRsa1024EcbOaepSHA512AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-512AndMGF1Padding",
            2048,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA512,
            OAEPParameterSpec(
                "SHA-512",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    @Test
    fun benchmarkRsa2048EcbOaepSHA512AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-512AndMGF1Padding",
            2048,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA512,
            OAEPParameterSpec(
                "SHA-512",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    @Test
    fun benchmarkRsa3072EcbOaepSHA512AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-512AndMGF1Padding",
            3072,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA512,
            OAEPParameterSpec(
                "SHA-512",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    @Test
    fun benchmarkRsa4096EcbOaepSHA512AndMgf1Padding() {
        benchmarkRsaAndMgf1Padding(
            "RSA/ECB/OAEPWithSHA-512AndMGF1Padding",
            4096,
            KeyProperties.BLOCK_MODE_ECB,
            KeyProperties.ENCRYPTION_PADDING_RSA_OAEP,
            KeyProperties.DIGEST_SHA512,
            OAEPParameterSpec(
                "SHA-512",
                "MGF1",
                MGF1ParameterSpec("SHA-1"),
                PSource.PSpecified.DEFAULT
            )
        )
    }

    private fun generateKey(
        keyType: String,
        blockMode: String? = null,
        keySize: Int,
        padding: String? = null,
        digest: String? = null
    ): KeyStore.PrivateKeyEntry {
        val kpg = KeyPairGenerator.getInstance(
            keyType,
            "AndroidKeyStore"
        )
        val parameterSpec: KeyGenParameterSpec =
            KeyGenParameterSpec.Builder(
                "RSA_KEY",
                KeyProperties.PURPOSE_ENCRYPT
                        or
                        KeyProperties.PURPOSE_DECRYPT
            )
                .apply {
                    setKeySize(keySize)
                    blockMode?.let { setBlockModes(it) }
                    padding?.let { setEncryptionPaddings(padding) }
                    digest?.let { setDigests(it) }
                }
                .build()
        kpg.apply { initialize(parameterSpec) }
        kpg.generateKeyPair()
        val kS = KeyStore.getInstance("AndroidKeyStore").apply { load(null) }
        return kS.getEntry("RSA_KEY", null) as KeyStore.PrivateKeyEntry
    }
    private fun benchmarkRsa(mode: String, keySize: Int, padding: String, digest: String? = null) {
        val key = generateKey(
            KeyProperties.KEY_ALGORITHM_RSA,
            KeyProperties.BLOCK_MODE_ECB,
            keySize,
            padding,
            digest
        )
        val cipher = Cipher.getInstance(mode).apply {
            init(
                Cipher.ENCRYPT_MODE,
                key.certificate.publicKey
            )
        }
        benchmarkRule.measureRepeated {
            cipher.doFinal(testFile)
        }
    }

    fun benchmarkRsaAndMgf1Padding(
        mode: String,
        keySize: Int,
        blockMode: String,
        padding: String,
        digest: String,
        spec: OAEPParameterSpec
    ) {
        val key = generateKey(
            KeyProperties.KEY_ALGORITHM_RSA,
            blockMode,
            keySize,
            padding,
            digest
        )

        val cipher = Cipher.getInstance(mode).apply {
            init(Cipher.ENCRYPT_MODE, key.certificate.publicKey, spec)
        }

        benchmarkRule.measureRepeated {
            cipher.doFinal(testFile)
        }
    }
}