package com.peachpasswords.sync

import kotlinx.coroutines.async
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.runBlocking
import org.bouncycastle.crypto.params.X25519PrivateKeyParameters
import org.junit.Assert.assertArrayEquals
import org.junit.Assert.assertEquals
import org.junit.Assert.assertTrue
import org.junit.Test
import java.security.SecureRandom

class NoiseKKTest {
    private fun keys(seed: Int? = null): NoiseKKKeyPair {
        val privateKey = if (seed == null) {
            X25519PrivateKeyParameters(SecureRandom())
        } else {
            X25519PrivateKeyParameters(ByteArray(32) { seed.toByte() }, 0)
        }
        return NoiseKKKeyPair(privateKey.generatePublicKey().encoded, privateKey.encoded)
    }

    @Test
    fun `KK mutually authenticates and carries transport frames`() = runBlocking {
        val initiator = keys()
        val responder = keys()
        val aToB = Channel<ByteArray>(Channel.UNLIMITED)
        val bToA = Channel<ByteArray>(Channel.UNLIMITED)
        val initiatorHandshake = async {
            noiseKKInitiator(
                initiator.secretKey,
                initiator.publicKey,
                responder.publicKey,
                { aToB.send(it) },
                { bToA.receive() },
            )
        }
        val responderHandshake = async {
            noiseKKResponder(
                responder.secretKey,
                responder.publicKey,
                initiator.publicKey,
                { bToA.send(it) },
                { aToB.receive() },
            )
        }
        val a = initiatorHandshake.await()
        val b = responderHandshake.await()

        assertArrayEquals(a.handshakeHash, b.handshakeHash)
        assertEquals("one", String(b.decrypt(a.encrypt("one".encodeToByteArray()))))
        assertEquals("two", String(a.decrypt(b.encrypt("two".encodeToByteArray()))))
    }

    @Test
    fun `KK handshake hash matches browser cross-platform vector`() = runBlocking {
        val initiator = keys(0x01)
        val responder = keys(0x02)
        val initiatorEphemeral = keys(0x03)
        val responderEphemeral = keys(0x04)
        val aToB = Channel<ByteArray>(Channel.UNLIMITED)
        val bToA = Channel<ByteArray>(Channel.UNLIMITED)

        val initiatorHandshake = async {
            noiseKKInitiator(
                initiator.secretKey,
                initiator.publicKey,
                responder.publicKey,
                { aToB.send(it) },
                { bToA.receive() },
                initiatorEphemeral,
            )
        }
        val responderHandshake = async {
            noiseKKResponder(
                responder.secretKey,
                responder.publicKey,
                initiator.publicKey,
                { bToA.send(it) },
                { aToB.receive() },
                responderEphemeral,
            )
        }

        val a = initiatorHandshake.await()
        val b = responderHandshake.await()
        assertEquals(
            "ac3dd2e4c40290fab2a0192b0b33b37bb8c31396ffec23271ffa2a99ef3a0f75",
            pubkeyToHex(a.handshakeHash),
        )
        assertArrayEquals(a.handshakeHash, b.handshakeHash)
    }

    @Test
    fun `KK rejects an initiator key mismatch without fallback`() = runBlocking {
        val initiator = keys()
        val attacker = keys()
        val responder = keys()
        val aToB = Channel<ByteArray>(Channel.UNLIMITED)
        val bToA = Channel<ByteArray>(Channel.UNLIMITED)

        val initiatorHandshake = async {
            runCatching {
                noiseKKInitiator(
                    initiator.secretKey,
                    initiator.publicKey,
                    responder.publicKey,
                    { aToB.send(it) },
                    { bToA.receive() },
                )
            }
        }
        val responderHandshake = async {
            runCatching {
                noiseKKResponder(
                    responder.secretKey,
                    responder.publicKey,
                    attacker.publicKey,
                    { bToA.send(it) },
                    { aToB.receive() },
                )
            }
        }

        val failure = responderHandshake.await().exceptionOrNull()
        assertTrue(failure?.message?.contains("decrypt", ignoreCase = true) == true)
        initiatorHandshake.cancel()
    }
}