package ch.threema.domain.protocol.multidevice

import ch.threema.base.crypto.NaCl
import ch.threema.domain.protocol.connection.data.D2dMessage
import ch.threema.protobuf.d2d.TransactionScope
import ch.threema.protobuf.d2d.transactionScope
import kotlin.test.Test
import kotlin.test.assertContentEquals
import kotlin.test.assertEquals

/**
 * Test vectors according to the `multidevice-kdf` repository
 */
class MultiDeviceKeysTest {

    private companion object {
        private val DGK1 = "1b35ed7e1ba9993171fe4a7eed30c2831905c3a583616d61e93782da900bf8ba".hexToByteArray()
        private val DGPK1 = "60c0ed9098902c0b6093be4a819bf344900bc7473504ccdb61004b1b58aa2233".hexToByteArray()
        private val DGRK1 = "a4bf34ff67ed3b731a2aa6e023335f7eba6c914e877da3d15bff41d84f7f75f8".hexToByteArray()
        private val DGDIK1 = "5953fb9775d8c23fae573e245534dac2c7aa16b62ea73b954ea4177d192e8c50".hexToByteArray()
        private val DGSDDK1 = "7a91ad36bc9537ccfd6fd76b23d7f5319e506e6a5294c988c9a75c8e34a3c9dd".hexToByteArray()
        private val DGTSK1 = "7c6b94affb171c564bfd375d50c4781d72ab2671ae4035ab4307dedce67ef30e".hexToByteArray()

        private val DGK2 = "7cf1c4847fb32d6c3702747018d0cccdc2f724c115bfca1036ae6208d2b7c68c".hexToByteArray()
        private val DGPK2 = "2b59e83e5a93ce4a09eeb4db91ec30e63cc3f173385742dd4a27ef83e5bef4f3".hexToByteArray()
        private val DGRK2 = "fa33785f528439c2bfd9b7220ff03ebca919c1c127aca91878f2038b65a79c73".hexToByteArray()
        private val DGDIK2 = "8a8199023dfeda5793e06552fb968282d7e7e29452c1229dcb212fb7d48f1e6a".hexToByteArray()
        private val DGSDDK2 = "bd96dd0f700c5d77da666990854eb4287d5f20a9be0deda6d88b875963c39b3f".hexToByteArray()
        private val DGTSK2 = "351e81590fb0115f8d1ab604c7eaf6996f3e079ea0a591db62caea4dcd7c6bda".hexToByteArray()

        private val DGK3 = "0000000000000000000000000000000000000000000000000000000000000000".hexToByteArray()
        private val DGPK3 = "4b464e5d33debe0d3f9be535b9a1449f79caac615c852da734b47ef3a23e14ca".hexToByteArray()
        private val DGRK3 = "d597f6380d1ecf6f1a7fd265c49bc53cff04c0efc1236a542dae338bba4bc6e9".hexToByteArray()
        private val DGDIK3 = "8bb45615b2982e8d6aefcfbe46a44bcea0df048df7a05d9d50afd01911ed10d8".hexToByteArray()
        private val DGSDDK3 = "86fa7962c7d85e01876de9a90cad95cf83b4605d40b83b6580b80618f56fa4d4".hexToByteArray()
        private val DGTSK3 = "7ab0a7c3239323e1b9f697eb59196d888747f027df356a305f58f55ebb8c23aa".hexToByteArray()
    }

    @Test
    fun testKeyDerivationDgk1() {
        val keys = MultiDeviceKeys(DGK1)

        assertContentEquals(DGPK1, keys.dgpk)
        assertContentEquals(DGRK1, keys.dgrk)
        assertContentEquals(DGDIK1, keys.dgdik)
        assertContentEquals(DGSDDK1, keys.dgsddk)
        assertContentEquals(DGTSK1, keys.dgtsk)
    }

    @Test
    fun testKeyDerivationDgk2() {
        val keys = MultiDeviceKeys(DGK2)

        assertContentEquals(DGPK2, keys.dgpk)
        assertContentEquals(DGRK2, keys.dgrk)
        assertContentEquals(DGDIK2, keys.dgdik)
        assertContentEquals(DGSDDK2, keys.dgsddk)
        assertContentEquals(DGTSK2, keys.dgtsk)
    }

    @Test
    fun testKeyDerivationDgk3() {
        val keys = MultiDeviceKeys(DGK3)

        assertContentEquals(DGPK3, keys.dgpk)
        assertContentEquals(DGRK3, keys.dgrk)
        assertContentEquals(DGDIK3, keys.dgdik)
        assertContentEquals(DGSDDK3, keys.dgsddk)
        assertContentEquals(DGTSK3, keys.dgtsk)
    }

    @Test
    fun testEncryptDecryptDeviceInfo() {
        val keys = MultiDeviceKeys(DGK1)
        val deviceInfo = D2dMessage.DeviceInfo(
            D2dMessage.DeviceInfo.Platform.ANDROID,
            "Unit Test",
            "1.2.3",
            "Test Client",
        )

        val encrypted = keys.encryptDeviceInfo(deviceInfo)
        val decrypted = keys.decryptDeviceInfo(encrypted)
        assertEquals(deviceInfo, decrypted)
    }

    @Test
    fun testEncryptTransactionScope() {
        val keys = MultiDeviceKeys(DGK1)

        listOf(
            TransactionScope.Scope.USER_PROFILE_SYNC,
            TransactionScope.Scope.CONTACT_SYNC,
            TransactionScope.Scope.GROUP_SYNC,
            TransactionScope.Scope.DISTRIBUTION_LIST_SYNC,
            TransactionScope.Scope.SETTINGS_SYNC,
            TransactionScope.Scope.MDM_PARAMETER_SYNC,
            TransactionScope.Scope.NEW_DEVICE_SYNC,
        ).forEach {
            val expected = transactionScope {
                scope = it
            }
            val encrypted = keys.encryptTransactionScope(it)
            val nonce = encrypted.copyOfRange(0, NaCl.NONCE_BYTES)
            val data = encrypted.copyOfRange(NaCl.NONCE_BYTES, encrypted.size)
            val decrypted = NaCl.symmetricDecryptData(
                data = data,
                key = keys.dgtsk,
                nonce = nonce,
            )
            assertEquals(expected, TransactionScope.parseFrom(decrypted))
        }
    }

    @Test
    fun testDecryptTransactionScope() {
        val keys = MultiDeviceKeys(DGK1)

        listOf(
            TransactionScope.Scope.USER_PROFILE_SYNC,
            TransactionScope.Scope.CONTACT_SYNC,
            TransactionScope.Scope.GROUP_SYNC,
            TransactionScope.Scope.DISTRIBUTION_LIST_SYNC,
            TransactionScope.Scope.SETTINGS_SYNC,
            TransactionScope.Scope.MDM_PARAMETER_SYNC,
            TransactionScope.Scope.NEW_DEVICE_SYNC,
        ).forEach { expectedScope ->
            val bytes = transactionScope {
                scope = expectedScope
            }.toByteArray()
            val nonce = ByteArray(24) { it.toByte() }
            val encrypted = nonce + NaCl.symmetricEncryptData(data = bytes, key = keys.dgtsk, nonce = nonce)
            val decrypted = keys.decryptTransactionScope(encrypted)

            assertEquals(
                expected = expectedScope,
                actual = decrypted,
            )
        }
    }
}
