package io.flamenet.e2e

import kotlin.test.Test
import kotlin.test.assertContentEquals
import kotlin.test.assertEquals
import kotlin.test.assertFailsWith
import kotlin.test.assertTrue

/**
 * X25519 against RFC 7748's published vectors, plus the raw/DER boundary the
 * wire depends on.
 */
class X25519Test {

    private fun hex(s: String): ByteArray =
        s.chunked(2).map { it.toInt(16).toByte() }.toByteArray()

    private fun ByteArray.hex(): String = joinToString("") { "%02x".format(it) }

    /**
     * RFC 7748 §6.1, the worked Diffie-Hellman example.
     *
     * Only the public halves are checked here: the JCA will not construct an
     * X25519 private key from raw scalar bytes without a PKCS#8 wrapper, and
     * building one by hand to check a number the RFC already publishes would be
     * testing the wrapper rather than the agreement. The agreement itself is
     * checked by the symmetry test below, which is the property that matters.
     */
    @Test
    fun `rfc 7748 public keys survive the raw to DER round trip`() {
        val alicePub = hex("8520f0098930a754748b7ddcb43ef75a0dbf3a0d26381af4eba4a98eaa9b4e6a")
        val bobPub = hex("de9edb7d7b7dc1b4d35b61c2ece435373f8343c85b78674dadfc7e146f882b4f")
        for (raw in listOf(alicePub, bobPub)) {
            val key = X25519.publicKeyFromRaw(raw)
            assertContentEquals(raw, X25519.rawFromPublicKey(key),
                                "raw -> DER -> raw must be the identity")
        }
    }

    /** The property the protocol actually rests on. */
    @Test
    fun `both sides derive the same secret`() {
        val alice = X25519.generate()
        val bob = X25519.generate()
        val fromAlice = X25519.agree(alice.private, X25519.rawFromPublicKey(bob.public))
        val fromBob = X25519.agree(bob.private, X25519.rawFromPublicKey(alice.public))
        assertContentEquals(fromAlice, fromBob)
        assertEquals(32, fromAlice.size)
    }

    @Test
    fun `a different peer gives a different secret`() {
        val alice = X25519.generate()
        val bob = X25519.generate()
        val eve = X25519.generate()
        val withBob = X25519.agree(alice.private, X25519.rawFromPublicKey(bob.public))
        val withEve = X25519.agree(alice.private, X25519.rawFromPublicKey(eve.public))
        assertTrue(!withBob.contentEquals(withEve))
    }

    @Test
    fun `public keys are exactly 32 bytes on the wire`() {
        assertEquals(32, X25519.rawFromPublicKey(X25519.generate().public).size)
    }

    @Test
    fun `a wrong-sized public key is refused rather than padded`() {
        assertFailsWith<IllegalArgumentException> { X25519.publicKeyFromRaw(ByteArray(31)) }
        assertFailsWith<IllegalArgumentException> { X25519.publicKeyFromRaw(ByteArray(33)) }
        assertFailsWith<IllegalArgumentException> { X25519.publicKeyFromRaw(ByteArray(0)) }
    }

    /**
     * RFC 7748 §6.1 lists the small-order points. A peer sending one forces the
     * shared secret to all zeros — which both sides would then happily derive
     * keys from, and the sender knows it in advance.
     *
     * The RFC says implementations MAY check. When the peer's key arrives over
     * a relay nobody trusts, that is not optional.
     */
    @Test
    fun `small-order peer keys are refused, not silently accepted`() {
        val alice = X25519.generate()
        val smallOrder = listOf(
            "0000000000000000000000000000000000000000000000000000000000000000",
            "0100000000000000000000000000000000000000000000000000000000000000",
            "e0eb7a7c3b41b8ae1656e3faf19fc46ada098deb9c32b1fd866205165f49b800",
            "5f9c95bca3508c24b1d0b1559c83ef5b04445cc4581c8e86d8224eddd09f1157",
        )
        for (point in smallOrder) {
            val thrown = runCatching { X25519.agree(alice.private, hex(point)) }.exceptionOrNull()
            assertTrue(
                thrown != null,
                "small-order point $point produced a secret instead of being refused",
            )
        }
    }
}