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.Test
import java.security.SecureRandom

class PairingAuthTest {
    private data class Keys(val privateKey: ByteArray, val publicKey: ByteArray)

    private fun keys(): Keys {
        val privateKey = X25519PrivateKeyParameters(SecureRandom())
        return Keys(privateKey.encoded, privateKey.generatePublicKey().encoded)
    }

    private fun session(hash: ByteArray) = object : NoiseNKSession {
        override fun encrypt(plaintext: ByteArray): ByteArray = plaintext
        override fun decrypt(ciphertext: ByteArray): ByteArray = ciphertext
        override val handshakeHash: ByteArray = hash
    }

    @Test
    fun `auth proof matches browser cross-platform vector`() {
        val proof = computeInitiatorAuthProof(
            initiatorPrivateKey = ByteArray(32) { 0x01 },
            initiatorPublicKey = hexToPubkey(
                "a4e09292b651c278b9772c569f5fa9bb13d906b46ab68c9df9dc2b4409f8a209",
            ),
            responderPublicKey = hexToPubkey(
                "ce8d3ad1ccb633ec7b70c17814a5c76ecd029685050d344745ba05870e587d59",
            ),
            nonce = ByteArray(32) { 0x03 },
            handshakeHash = ByteArray(32) { 0x04 },
        )

        assertEquals(
            "d947d11a82892a65fdde94da5bbe6ca6034cf1f7e3268fc3ce5851b55f5813ce",
            pubkeyToHex(proof),
        )
    }

    @Test
    fun `initiator proof round trips and returns authenticated key`() = runBlocking {
        val initiator = keys()
        val responder = keys()
        val hash = ByteArray(32) { 0x07 }
        val initToResponder = Channel<ByteArray>(Channel.UNLIMITED)
        val responderToInit = Channel<ByteArray>(Channel.UNLIMITED)

        val prove = async {
            performInitiatorPairingAuth(
                session(hash),
                initiator.privateKey,
                initiator.publicKey,
                responder.publicKey,
                { initToResponder.send(it) },
                { responderToInit.receive() },
            )
        }
        val verify = async {
            performResponderPairingAuth(
                session(hash),
                responder.privateKey,
                responder.publicKey,
                { responderToInit.send(it) },
                { initToResponder.receive() },
            )
        }

        prove.await()
        assertArrayEquals(initiator.publicKey, verify.await())
    }

    @Test
    fun `proof from a different handshake hash is rejected`() = runBlocking {
        val initiator = keys()
        val responder = keys()
        val initToResponder = Channel<ByteArray>(Channel.UNLIMITED)
        val responderToInit = Channel<ByteArray>(Channel.UNLIMITED)

        val prove = async {
            performInitiatorPairingAuth(
                session(ByteArray(32) { 0x07 }),
                initiator.privateKey,
                initiator.publicKey,
                responder.publicKey,
                { initToResponder.send(it) },
                { responderToInit.receive() },
            )
        }
        val verify = async {
            runCatching {
                performResponderPairingAuth(
                    session(ByteArray(32) { 0x09 }),
                    responder.privateKey,
                    responder.publicKey,
                    { responderToInit.send(it) },
                    { initToResponder.receive() },
                )
            }
        }

        prove.await()
        val failure = verify.await().exceptionOrNull()
        assertEquals("auth proof did not verify", failure?.message)
    }
}