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 javax.crypto.Cipher
import javax.crypto.KeyGenerator
import javax.crypto.SecretKey
import javax.crypto.spec.GCMParameterSpec
import javax.crypto.spec.IvParameterSpec


@RunWith(AndroidJUnit4::class)
class SymmetricEncryptionBenchmark {

    @get:Rule
    val benchmarkRule = BenchmarkRule()

    val testFile = "qwertyuiopasdfghjklzxcvbnm123456789".toByteArray()
    val testFile16Byte = "qwertyuiqwertyui".toByteArray()

    @Test
    fun benchmarkAES128CbcNoPadding() {
        val key = generateKey(
            keySize = 128,
            blockMode = KeyProperties.BLOCK_MODE_CBC,
            padding = KeyProperties.ENCRYPTION_PADDING_NONE
        )
        Assert.assertEquals(testFile16Byte.size, 16)

        benchmarkRule.measureRepeated {
            val cipher = runWithTimingDisabled {
                getCipher("AES/CBC/NoPadding", key)
            }
            cipher.doFinal(testFile16Byte)
        }
    }

    @Test
    fun benchmarkAES192CbcNoPadding() {
        val key = generateKey(
            keySize = 192,
            blockMode = KeyProperties.BLOCK_MODE_CBC,
            padding = KeyProperties.ENCRYPTION_PADDING_NONE
        )
        Assert.assertEquals(testFile16Byte.size, 16)

        benchmarkRule.measureRepeated {
            val cipher = runWithTimingDisabled {
                getCipher("AES/CBC/NoPadding", key)
            }
            cipher.doFinal(testFile16Byte)
        }
    }

    @Test
    fun benchmarkAES256CbcNoPadding() {
        val key = generateKey(
            keySize = 256,
            blockMode = KeyProperties.BLOCK_MODE_CBC,
            padding = KeyProperties.ENCRYPTION_PADDING_NONE
        )
        Assert.assertEquals(testFile16Byte.size, 16)

        benchmarkRule.measureRepeated {
            val cipher = runWithTimingDisabled {
                getCipher("AES/CBC/NoPadding", key)
            }
            cipher.doFinal(testFile16Byte)
        }
    }


    @Test
    fun benchmarkAES128CbcPKCS7Padding() {
        val key = generateKey(
            keySize = 128,
            blockMode = KeyProperties.BLOCK_MODE_CBC,
            padding = KeyProperties.ENCRYPTION_PADDING_PKCS7
        )
        benchmarkRule.measureRepeated {
            val cipher = runWithTimingDisabled {
                getCipher("AES/CBC/PKCS7Padding", key)
            }
            cipher.doFinal(testFile)
        }
    }

    @Test
    fun benchmarkAES192CbcPKCS7Padding() {
        val key = generateKey(
            keySize = 192,
            blockMode = KeyProperties.BLOCK_MODE_CBC,
            padding = KeyProperties.ENCRYPTION_PADDING_PKCS7
        )
        benchmarkRule.measureRepeated {
            val cipher = runWithTimingDisabled {
                getCipher("AES/CBC/PKCS7Padding", key)
            }
            cipher.doFinal(testFile)
        }
    }

    @Test
    fun benchmarkAES256CbcPKCS7Padding() {
        val key = generateKey(
            keySize = 256,
            blockMode = KeyProperties.BLOCK_MODE_CBC,
            padding = KeyProperties.ENCRYPTION_PADDING_PKCS7
        )
        benchmarkRule.measureRepeated {
            val cipher = runWithTimingDisabled {
                getCipher("AES/CBC/PKCS7Padding", key)
            }
            cipher.doFinal(testFile)
        }
    }

    @Test
    fun benchmarkAES128CTRNoPadding() {
        val key = generateKey(
            keySize = 128,
            blockMode = KeyProperties.BLOCK_MODE_CTR,
            padding = KeyProperties.ENCRYPTION_PADDING_NONE
        )
        benchmarkRule.measureRepeated {
            val cipher = runWithTimingDisabled {
                getCipher("AES/CTR/NoPadding", key)
            }
            cipher.doFinal(testFile)
        }
    }

    @Test
    fun benchmarkAES192CTRNoPadding() {
        val key = generateKey(
            keySize = 192,
            blockMode = KeyProperties.BLOCK_MODE_CTR,
            padding = KeyProperties.ENCRYPTION_PADDING_NONE
        )
        benchmarkRule.measureRepeated {
            val cipher = runWithTimingDisabled {
                getCipher("AES/CTR/NoPadding", key)
            }
            cipher.doFinal(testFile)
        }
    }

    @Test
    fun benchmarkAES256CTRNoPadding() {
        val key = generateKey(
            keySize = 256,
            blockMode = KeyProperties.BLOCK_MODE_CTR,
            padding = KeyProperties.ENCRYPTION_PADDING_NONE
        )
        benchmarkRule.measureRepeated {
            val cipher = runWithTimingDisabled {
                getCipher("AES/CTR/NoPadding", key)
            }
            cipher.doFinal(testFile)
        }
    }

    @Test
    fun benchmarkAES128ECBNoPadding() {
        val key = generateKey(
            keySize = 128,
            blockMode = KeyProperties.BLOCK_MODE_ECB,
            padding = KeyProperties.ENCRYPTION_PADDING_NONE,
            randomizedRequired = false
        )
        benchmarkRule.measureRepeated {
            val cipher = runWithTimingDisabled {
                getCipher("AES/ECB/NoPadding", key)
            }
            cipher.doFinal(testFile16Byte)
        }
    }

    @Test
    fun benchmarkAES192ECBNoPadding() {
        val key = generateKey(
            keySize = 192,
            blockMode = KeyProperties.BLOCK_MODE_ECB,
            padding = KeyProperties.ENCRYPTION_PADDING_NONE,
            randomizedRequired = false
        )
        benchmarkRule.measureRepeated {
            val cipher = runWithTimingDisabled {
                getCipher("AES/ECB/NoPadding", key)
            }
            cipher.doFinal(testFile16Byte)
        }
    }

    @Test
    fun benchmarkAES256ECBNoPadding() {
        val key = generateKey(
            keySize = 256,
            blockMode = KeyProperties.BLOCK_MODE_ECB,
            padding = KeyProperties.ENCRYPTION_PADDING_NONE,
            randomizedRequired = false
        )
        benchmarkRule.measureRepeated {
            val cipher = runWithTimingDisabled {
                getCipher("AES/ECB/NoPadding", key)
            }
            cipher.doFinal(testFile16Byte)
        }
    }

    @Test
    fun benchmarkAES128ECBPKCS7Padding() {
        val key = generateKey(
            keySize = 128,
            blockMode = KeyProperties.BLOCK_MODE_ECB,
            padding = KeyProperties.ENCRYPTION_PADDING_PKCS7,
            randomizedRequired = false
        )
        benchmarkRule.measureRepeated {
            val cipher = runWithTimingDisabled {
                getCipher("AES/ECB/PKCS7Padding", key)
            }
            cipher.doFinal(testFile)
        }
    }

    @Test
    fun benchmarkAES192ECBPKCS7Padding() {
        val key = generateKey(
            keySize = 192,
            blockMode = KeyProperties.BLOCK_MODE_ECB,
            padding = KeyProperties.ENCRYPTION_PADDING_PKCS7,
            randomizedRequired = false
        )
        benchmarkRule.measureRepeated {
            val cipher = runWithTimingDisabled {
                getCipher("AES/ECB/PKCS7Padding", key)
            }
            cipher.doFinal(testFile)
        }
    }

    @Test
    fun benchmarkAES256ECBPKCS7Padding() {
        val key = generateKey(
            keySize = 256,
            blockMode = KeyProperties.BLOCK_MODE_ECB,
            padding = KeyProperties.ENCRYPTION_PADDING_PKCS7,
            randomizedRequired = false
        )
        benchmarkRule.measureRepeated {
            val cipher = runWithTimingDisabled {
                getCipher("AES/ECB/PKCS7Padding", key)
            }
            cipher.doFinal(testFile)
        }
    }

    @Test
    fun benchmarkAES128GCMNoPadding() {
        val key = generateKey(
            keySize = 128,
            blockMode = KeyProperties.BLOCK_MODE_GCM,
            padding = KeyProperties.ENCRYPTION_PADDING_NONE
        )
        benchmarkRule.measureRepeated {
            val cipher = runWithTimingDisabled {
                getCipher("AES/GCM/NoPadding", key)
            }
            cipher.doFinal(testFile)
        }
    }

    @Test
    fun benchmarkAES192GCMNoPadding() {
        val key = generateKey(
            keySize = 192,
            blockMode = KeyProperties.BLOCK_MODE_GCM,
            padding = KeyProperties.ENCRYPTION_PADDING_NONE
        )
        benchmarkRule.measureRepeated {
            val cipher = runWithTimingDisabled {
                getCipher("AES/GCM/NoPadding", key)
            }
            cipher.doFinal(testFile)
        }
    }

    @Test
    fun benchmarkAES256GCMNoPadding() {
        val key = generateKey(
            keySize = 256,
            blockMode = KeyProperties.BLOCK_MODE_GCM,
            padding = KeyProperties.ENCRYPTION_PADDING_NONE
        )
        benchmarkRule.measureRepeated {
            val cipher = runWithTimingDisabled {
                getCipher("AES/GCM/NoPadding", key)
            }
            cipher.doFinal(testFile)
        }
    }

    private fun getCipher(type: String, key: SecretKey): Cipher {
        return Cipher.getInstance(type).apply {
            init(
                Cipher.ENCRYPT_MODE,
                key
            )
        }
    }

    private fun getDecipher(type: String, key: SecretKey, vector: ByteArray): Cipher {
        val spec = IvParameterSpec(vector)
        return Cipher.getInstance(type).apply {
            init(
                Cipher.DECRYPT_MODE,
                key, spec
            )
        }
    }

    private fun getGCMDecipher(type: String, key: SecretKey, vector: ByteArray): Cipher {
        val spec = GCMParameterSpec(128, vector)
        return Cipher.getInstance(type).apply {
            init(
                Cipher.DECRYPT_MODE,
                key, spec
            )
        }
    }

    private fun generateKey(
        keySize: Int,
        blockMode: String? = null,
        padding: String? = null,
        randomizedRequired: Boolean = true
    ): SecretKey {
        val keyGenerator = KeyGenerator
            .getInstance("AES", "AndroidKeyStore")
        val keyGenParameterSpec =
            KeyGenParameterSpec.Builder(
                "AES_KEY_ALIAS",
                KeyProperties.PURPOSE_ENCRYPT
                        or
                        KeyProperties.PURPOSE_DECRYPT
            ).apply {
                blockMode?.also { setBlockModes(it) }
                padding?.also { setEncryptionPaddings(it) }
                setRandomizedEncryptionRequired(randomizedRequired)
            }
                .setKeySize(keySize)
                .build()
        keyGenerator.init(keyGenParameterSpec)
        return keyGenerator.generateKey()
    }
}
