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.Signature

@RunWith(AndroidJUnit4::class)
class AsymmetricSignatureBenchmark {

    @get:Rule
    val benchmarkRule = BenchmarkRule()

    private val testFile = "qwertyuiopasdfghjklzxcvbnm123456789".toByteArray()

    @Test
    fun benchmarkRsa512EcbPKCS1Padding() {
        benchmarkDigestWithRsa("MD5withRSA", 512, KeyProperties.DIGEST_MD5)
    }

    @Test
    fun benchmarkRsa768EcbPKCS1Padding() {
        benchmarkDigestWithRsa("MD5withRSA", 768, KeyProperties.DIGEST_MD5)
    }

    @Test
    fun benchmarkRsa1024EcbPKCS1Padding() {
        benchmarkDigestWithRsa("MD5withRSA", 1024, KeyProperties.DIGEST_MD5)
    }

    @Test
    fun benchmarkRsa2048EcbPKCS1Padding() {
        benchmarkDigestWithRsa("MD5withRSA", 2048, KeyProperties.DIGEST_MD5)
    }

    @Test
    fun benchmarkRsa3072EcbPKCS1Padding() {
        benchmarkDigestWithRsa("MD5withRSA", 3072, KeyProperties.DIGEST_MD5)
    }

    @Test
    fun benchmarkRsa4096EcbPKCS1Padding() {
        benchmarkDigestWithRsa("MD5withRSA", 4096, KeyProperties.DIGEST_MD5)
    }

    @Test
    fun benchmarkNoneWithECDSA224() {
        benchmarkDigestWithECDSA("NONEwithECDSA", 224, KeyProperties.DIGEST_NONE)
    }

    @Test
    fun benchmarkNoneWithECDSA256() {
        benchmarkDigestWithECDSA("NONEwithECDSA", 256, KeyProperties.DIGEST_NONE)
    }

    @Test
    fun benchmarkNoneWithECDSA384() {
        benchmarkDigestWithECDSA("NONEwithECDSA", 384, KeyProperties.DIGEST_NONE)
    }

    @Test
    fun benchmarkNoneWithECDSA521() {
        benchmarkDigestWithECDSA("NONEwithECDSA", 521, KeyProperties.DIGEST_NONE)
    }

    @Test
    fun benchmarkNoneWithRSA512() {
        benchmarkDigestWithRsa("NONEwithRSA", 512, KeyProperties.DIGEST_NONE)
    }

    @Test
    fun benchmarkNoneWithRSA768() {
        benchmarkDigestWithRsa("NONEwithRSA", 768, KeyProperties.DIGEST_NONE)
    }

    @Test
    fun benchmarkNoneWithRSA1024() {
        benchmarkDigestWithRsa("NONEwithRSA", 1024, KeyProperties.DIGEST_NONE)
    }

    @Test
    fun benchmarkNoneWithRSA2048() {
        benchmarkDigestWithRsa("NONEwithRSA", 2048, KeyProperties.DIGEST_NONE)
    }

    @Test
    fun benchmarkNoneWithRSA3072() {
        benchmarkDigestWithRsa("NONEwithRSA", 3072, KeyProperties.DIGEST_NONE)
    }

    @Test
    fun benchmarkNoneWithRSA4096() {
        benchmarkDigestWithRsa("NONEwithRSA", 4096, KeyProperties.DIGEST_NONE)
    }

    @Test
    fun benchmarkSHA1WithECDSA224() {
        benchmarkDigestWithECDSA("SHA1withECDSA", 224, KeyProperties.DIGEST_SHA1)
    }

    @Test
    fun benchmarkSHA1WithECDSA256() {
        benchmarkDigestWithECDSA("SHA1withECDSA", 256, KeyProperties.DIGEST_SHA1)
    }

    @Test
    fun benchmarkSHA1WithECDSA384() {
        benchmarkDigestWithECDSA("SHA1withECDSA", 384, KeyProperties.DIGEST_SHA1)
    }

    @Test
    fun benchmarkSHA1WithECDSA521() {
        benchmarkDigestWithECDSA("SHA1withECDSA", 521, KeyProperties.DIGEST_SHA1)
    }

    @Test
    fun benchmarkSHA1WithRSA512() {
        benchmarkDigestWithRsa("SHA1withRSA", 512, KeyProperties.DIGEST_SHA1)
    }

    @Test
    fun benchmarkSHA1WithRSA768() {
        benchmarkDigestWithRsa("SHA1withRSA", 768, KeyProperties.DIGEST_SHA1)
    }

    @Test
    fun benchmarkSHA1WithRSA1024() {
        benchmarkDigestWithRsa("SHA1withRSA", 1024, KeyProperties.DIGEST_SHA1)
    }

    @Test
    fun benchmarkSHA1WithRSA2048() {
        benchmarkDigestWithRsa("SHA1withRSA", 2048, KeyProperties.DIGEST_SHA1)
    }

    @Test
    fun benchmarkSHA1WithRSA3072() {
        benchmarkDigestWithRsa("SHA1withRSA", 3072, KeyProperties.DIGEST_SHA1)
    }

    @Test
    fun benchmarkSHA1WithRSA4096() {
        benchmarkDigestWithRsa("SHA1withRSA", 4096, KeyProperties.DIGEST_SHA1)
    }

    @Test
    fun benchmarkSHA1WithRSA512Pss() {
        benchmarkDigestWithRsa(
            "SHA1withRSA/PSS",
            512,
            KeyProperties.DIGEST_SHA1,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA1WithRSA768Pss() {
        benchmarkDigestWithRsa(
            "SHA1withRSA/PSS",
            768,
            KeyProperties.DIGEST_SHA1,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA1WithRSA1024Pss() {
        benchmarkDigestWithRsa(
            "SHA1withRSA/PSS",
            1024,
            KeyProperties.DIGEST_SHA1,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA1WithRSA2048Pss() {
        benchmarkDigestWithRsa(
            "SHA1withRSA/PSS",
            2048,
            KeyProperties.DIGEST_SHA1,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA1WithRSA3072Pss() {
        benchmarkDigestWithRsa(
            "SHA1withRSA/PSS",
            3072,
            KeyProperties.DIGEST_SHA1,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA1WithRSA4096Pss() {
        benchmarkDigestWithRsa(
            "SHA1withRSA/PSS",
            4096,
            KeyProperties.DIGEST_SHA1,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA224WithECDSA224() {
        benchmarkDigestWithECDSA("SHA224withECDSA", 224, KeyProperties.DIGEST_SHA224)
    }

    @Test
    fun benchmarkSHA224WithECDSA256() {
        benchmarkDigestWithECDSA("SHA224withECDSA", 256, KeyProperties.DIGEST_SHA224)
    }

    @Test
    fun benchmarkSHA224WithECDSA384() {
        benchmarkDigestWithECDSA("SHA224withECDSA", 384, KeyProperties.DIGEST_SHA224)
    }

    @Test
    fun benchmarkSHA224WithECDSA521() {
        benchmarkDigestWithECDSA("SHA224withECDSA", 521, KeyProperties.DIGEST_SHA224)
    }

    @Test
    fun benchmarkSHA224WithRSA512() {
        benchmarkDigestWithRsa("SHA224withRSA", 512, KeyProperties.DIGEST_SHA224)
    }

    @Test
    fun benchmarkSHA224WithRSA768() {
        benchmarkDigestWithRsa("SHA224withRSA", 768, KeyProperties.DIGEST_SHA224)
    }

    @Test
    fun benchmarkSHA224WithRSA1024() {
        benchmarkDigestWithRsa("SHA224withRSA", 1024, KeyProperties.DIGEST_SHA224)
    }

    @Test
    fun benchmarkSHA224WithRSA2048() {
        benchmarkDigestWithRsa("SHA224withRSA", 2048, KeyProperties.DIGEST_SHA224)
    }

    @Test
    fun benchmarkSHA224WithRSA3072() {
        benchmarkDigestWithRsa("SHA224withRSA", 3072, KeyProperties.DIGEST_SHA224)
    }

    @Test
    fun benchmarkSHA224WithRSA4096() {
        benchmarkDigestWithRsa("SHA224withRSA", 4096, KeyProperties.DIGEST_SHA224)
    }

    @Test
    fun benchmarkSHA224WithRSA512Pss() {
        benchmarkDigestWithRsa(
            "SHA224withRSA/PSS",
            512,
            KeyProperties.DIGEST_SHA224,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA224WithRSA768Pss() {
        benchmarkDigestWithRsa(
            "SHA224withRSA/PSS",
            768,
            KeyProperties.DIGEST_SHA224,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA224WithRSA1024Pss() {
        benchmarkDigestWithRsa(
            "SHA224withRSA/PSS",
            1024,
            KeyProperties.DIGEST_SHA224,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA224WithRSA2048Pss() {
        benchmarkDigestWithRsa(
            "SHA224withRSA/PSS",
            2048,
            KeyProperties.DIGEST_SHA224,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA224WithRSA3072Pss() {
        benchmarkDigestWithRsa(
            "SHA224withRSA/PSS",
            3072,
            KeyProperties.DIGEST_SHA224,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA224WithRSA4096Pss() {
        benchmarkDigestWithRsa(
            "SHA224withRSA/PSS",
            4096,
            KeyProperties.DIGEST_SHA224,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA256WithECDSA224() {
        benchmarkDigestWithECDSA("SHA256withECDSA", 224, KeyProperties.DIGEST_SHA256)
    }

    @Test
    fun benchmarkSHA256WithECDSA256() {
        benchmarkDigestWithECDSA("SHA256withECDSA", 256, KeyProperties.DIGEST_SHA256)
    }

    @Test
    fun benchmarkSHA256WithECDSA384() {
        benchmarkDigestWithECDSA("SHA256withECDSA", 384, KeyProperties.DIGEST_SHA256)
    }

    @Test
    fun benchmarkSHA256WithECDSA521() {
        benchmarkDigestWithECDSA("SHA256withECDSA", 521, KeyProperties.DIGEST_SHA256)
    }

    @Test
    fun benchmarkSHA256WithRSA512() {
        benchmarkDigestWithRsa("SHA256withRSA", 512, KeyProperties.DIGEST_SHA256)
    }

    @Test
    fun benchmarkSHA256WithRSA768() {
        benchmarkDigestWithRsa("SHA256withRSA", 768, KeyProperties.DIGEST_SHA256)
    }

    @Test
    fun benchmarkSHA256WithRSA1024() {
        benchmarkDigestWithRsa("SHA256withRSA", 1024, KeyProperties.DIGEST_SHA256)
    }

    @Test
    fun benchmarkSHA256WithRSA2048() {
        benchmarkDigestWithRsa("SHA256withRSA", 2048, KeyProperties.DIGEST_SHA256)
    }

    @Test
    fun benchmarkSHA256WithRSA3072() {
        benchmarkDigestWithRsa("SHA256withRSA", 3072, KeyProperties.DIGEST_SHA256)
    }

    @Test
    fun benchmarkSHA256WithRSA4096() {
        benchmarkDigestWithRsa("SHA256withRSA", 4096, KeyProperties.DIGEST_SHA256)
    }

    //@Test
    fun benchmarkSHA256WithRSA512Pss() {
        // Not working on LG G6
        benchmarkDigestWithRsa(
            "SHA256withRSA/PSS",
            512,
            KeyProperties.DIGEST_SHA256,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA256WithRSA768Pss() {
        benchmarkDigestWithRsa(
            "SHA256withRSA/PSS",
            768,
            KeyProperties.DIGEST_SHA256,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA256WithRSA1024Pss() {
        benchmarkDigestWithRsa(
            "SHA256withRSA/PSS",
            1024,
            KeyProperties.DIGEST_SHA256,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA256WithRSA2048Pss() {
        benchmarkDigestWithRsa(
            "SHA256withRSA/PSS",
            2048,
            KeyProperties.DIGEST_SHA256,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA256WithRSA3072Pss() {
        benchmarkDigestWithRsa(
            "SHA256withRSA/PSS",
            3072,
            KeyProperties.DIGEST_SHA256,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA256WithRSA4096Pss() {
        benchmarkDigestWithRsa(
            "SHA256withRSA/PSS",
            4096,
            KeyProperties.DIGEST_SHA256,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA384WithECDSA224() {
        benchmarkDigestWithECDSA("SHA384withECDSA", 224, KeyProperties.DIGEST_SHA384)
    }

    @Test
    fun benchmarkSHA384WithECDSA256() {
        benchmarkDigestWithECDSA("SHA384withECDSA", 256, KeyProperties.DIGEST_SHA384)
    }

    @Test
    fun benchmarkSHA384WithECDSA384() {
        benchmarkDigestWithECDSA("SHA384withECDSA", 384, KeyProperties.DIGEST_SHA384)
    }

    @Test
    fun benchmarkSHA384WithECDSA521() {
        benchmarkDigestWithECDSA("SHA384withECDSA", 521, KeyProperties.DIGEST_SHA384)
    }

    //@Test
    fun benchmarkSHA384WithRSA512() {
        // Not working on LG G6 9.0
        benchmarkDigestWithRsa("SHA384withRSA", 512, KeyProperties.DIGEST_SHA384)
    }

    @Test
    fun benchmarkSHA384WithRSA768() {
        benchmarkDigestWithRsa("SHA384withRSA", 768, KeyProperties.DIGEST_SHA384)
    }

    @Test
    fun benchmarkSHA384WithRSA1024() {
        benchmarkDigestWithRsa("SHA384withRSA", 1024, KeyProperties.DIGEST_SHA384)
    }

    @Test
    fun benchmarkSHA384WithRSA2048() {
        benchmarkDigestWithRsa("SHA384withRSA", 2048, KeyProperties.DIGEST_SHA384)
    }

    @Test
    fun benchmarkSHA384WithRSA3072() {
        benchmarkDigestWithRsa("SHA384withRSA", 3072, KeyProperties.DIGEST_SHA384)
    }

    @Test
    fun benchmarkSHA384WithRSA4096() {
        benchmarkDigestWithRsa("SHA384withRSA", 4096, KeyProperties.DIGEST_SHA384)
    }

    //@Test
    fun benchmarkSHA384WithRSA512Pss() {
        // Not working on LG G6 9.0
        benchmarkDigestWithRsa(
            "SHA384withRSA/PSS",
            512,
            KeyProperties.DIGEST_SHA384,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    //@Test
    fun benchmarkSHA384WithRSA768Pss() {
        // Not working on LG G6 9.0
        benchmarkDigestWithRsa(
            "SHA384withRSA/PSS",
            768,
            KeyProperties.DIGEST_SHA384,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA384WithRSA1024Pss() {
        benchmarkDigestWithRsa(
            "SHA384withRSA/PSS",
            1024,
            KeyProperties.DIGEST_SHA384,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA384WithRSA2048Pss() {
        benchmarkDigestWithRsa(
            "SHA384withRSA/PSS",
            2048,
            KeyProperties.DIGEST_SHA384,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA384WithRSA3072Pss() {
        benchmarkDigestWithRsa(
            "SHA384withRSA/PSS",
            3072,
            KeyProperties.DIGEST_SHA384,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA384WithRSA4096Pss() {
        benchmarkDigestWithRsa(
            "SHA384withRSA/PSS",
            4096,
            KeyProperties.DIGEST_SHA384,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA512WithECDSA224() {
        benchmarkDigestWithECDSA("SHA512withECDSA", 224, KeyProperties.DIGEST_SHA512)
    }

    @Test
    fun benchmarkSHA512WithECDSA256() {
        benchmarkDigestWithECDSA("SHA512withECDSA", 256, KeyProperties.DIGEST_SHA512)
    }

    @Test
    fun benchmarkSHA512WithECDSA384() {
        benchmarkDigestWithECDSA("SHA512withECDSA", 384, KeyProperties.DIGEST_SHA512)
    }

    @Test
    fun benchmarkSHA512WithECDSA521() {
        benchmarkDigestWithECDSA("SHA512withECDSA", 521, KeyProperties.DIGEST_SHA512)
    }

    //@Test
    fun benchmarkSHA512WithRSA512() {
        // Not working on LG G6 9.0
        benchmarkDigestWithRsa("SHA512withRSA", 512, KeyProperties.DIGEST_SHA512)
    }

    @Test
    fun benchmarkSHA512WithRSA768() {
        benchmarkDigestWithRsa("SHA512withRSA", 768, KeyProperties.DIGEST_SHA512)
    }

    @Test
    fun benchmarkSHA512WithRSA1024() {
        benchmarkDigestWithRsa("SHA512withRSA", 1024, KeyProperties.DIGEST_SHA512)
    }

    @Test
    fun benchmarkSHA512WithRSA2048() {
        benchmarkDigestWithRsa("SHA512withRSA", 2048, KeyProperties.DIGEST_SHA512)
    }

    @Test
    fun benchmarkSHA512WithRSA3072() {
        benchmarkDigestWithRsa("SHA512withRSA", 3072, KeyProperties.DIGEST_SHA512)
    }

    @Test
    fun benchmarkSHA512WithRSA4096() {
        benchmarkDigestWithRsa("SHA512withRSA", 4096, KeyProperties.DIGEST_SHA512)
    }

    //@Test
    fun benchmarkSHA512WithRSA512Pss() {
        // Not working on LG G6 9.0
        benchmarkDigestWithRsa(
            "SHA512withRSA/PSS",
            512,
            KeyProperties.DIGEST_SHA512,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    //@Test
    fun benchmarkSHA512WithRSA768Pss() {
        // Not working on LG G6 9.0
        benchmarkDigestWithRsa(
            "SHA512withRSA/PSS",
            768,
            KeyProperties.DIGEST_SHA512,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    //@Test
    fun benchmarkSHA512WithRSA1024Pss() {
        // Not working on LG G6 9.0
        benchmarkDigestWithRsa(
            "SHA512withRSA/PSS",
            1024,
            KeyProperties.DIGEST_SHA512,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA512WithRSA2048Pss() {
        benchmarkDigestWithRsa(
            "SHA512withRSA/PSS",
            2048,
            KeyProperties.DIGEST_SHA512,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    @Test
    fun benchmarkSHA512WithRSA3072Pss() {
        benchmarkDigestWithRsa(
            "SHA512withRSA/PSS",
            3072,
            KeyProperties.DIGEST_SHA512,
            KeyProperties.SIGNATURE_PADDING_RSA_PSS
        )
    }

    private fun benchmarkDigestWithRsa(
        mode: String,
        keySize: Int,
        digest: String,
        padding: String = KeyProperties.SIGNATURE_PADDING_RSA_PKCS1
    ) {
        val kpg = KeyPairGenerator.getInstance(KeyProperties.KEY_ALGORITHM_RSA, "AndroidKeyStore")
        val parameterSpec: KeyGenParameterSpec =
            KeyGenParameterSpec.Builder(
                "TEST_KEY",
                KeyProperties.PURPOSE_SIGN
                        or
                        KeyProperties.PURPOSE_VERIFY
            )
                .apply {
                    setKeySize(keySize)
                    setDigests(digest)
                    setSignaturePaddings(padding)
                }
                .build()
        kpg.apply {
            initialize(parameterSpec)
            generateKeyPair()
        }
        val kS = KeyStore.getInstance("AndroidKeyStore").apply { load(null) }
        val entry = kS.getEntry("TEST_KEY", null) as KeyStore.PrivateKeyEntry

        benchmarkRule.measureRepeated {
            Signature.getInstance(mode)
                .run {
                    initSign(entry.privateKey)
                    update(testFile)
                    sign()
                }
        }
    }

    private fun benchmarkDigestWithECDSA(mode: String, keySize: Int, digest: String) {
        val kpg = KeyPairGenerator.getInstance(KeyProperties.KEY_ALGORITHM_EC, "AndroidKeyStore")
        val parameterSpec: KeyGenParameterSpec =
            KeyGenParameterSpec.Builder(
                "TEST_KEY",
                KeyProperties.PURPOSE_SIGN
                        or
                        KeyProperties.PURPOSE_VERIFY
            )
                .apply {
                    setKeySize(keySize)
                    setDigests(digest)
                }
                .build()
        kpg.apply {
            initialize(parameterSpec)
            generateKeyPair()
        }
        val kS = KeyStore.getInstance("AndroidKeyStore").apply { load(null) }
        val entry = kS.getEntry("TEST_KEY", null) as KeyStore.PrivateKeyEntry

        benchmarkRule.measureRepeated {
            Signature.getInstance(mode)
                .run {
                    initSign(entry.privateKey)
                    update(testFile)
                    sign()
                }
        }
    }
}