package io.flamenet.e2e

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

/**
 * AES-256-GCM against NIST's published vectors, plus the layout the protocol
 * depends on.
 *
 * The vectors are from NIST CAVP `gcmEncryptExtIV256`, not from our other
 * engines — same reasoning as [HkdfTest]. The layout assertions ARE about our
 * protocol, and they are the ones that would otherwise be discovered by two
 * engines failing to talk to each other.
 */
class AeadTest {

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

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

    /** NIST CAVP, AES-256-GCM, 96-bit IV, empty plaintext and AAD. */
    @Test
    fun `nist vector - empty plaintext`() {
        val key = hex("b52c505a37d78eda5dd34f20c22540ea1b58963cf8e5bf8ffa85f9f2492505b4")
        val nonce = hex("516c33929df5a3284ff463d7")
        val box = Aead.seal(key, nonce, ByteArray(0))
        assertEquals("bdc1ac884d332457a1d2664f168c76f0", box.copyOfRange(12, box.size).hex())
    }

    /** NIST CAVP, AES-256-GCM, 96-bit IV, 128-bit plaintext, no AAD. */
    @Test
    fun `nist vector - one block`() {
        val key = hex("31bdadd96698c204aa9ce1448ea94ae1fb4a9a0b3c9d773b51bb1822666b8f22")
        val nonce = hex("0d18e06c7c725ac9e362e1ce")
        val plaintext = hex("2db5168e932556f8089a0622981d017d")
        val box = Aead.seal(key, nonce, plaintext)
        val body = box.copyOfRange(12, box.size)
        assertEquals(
            "fa4362189661d163fcd6a56d8bf0405ad636ac1bbedd5cc3ee727dc2ab4a9489",
            body.hex(),
            "ciphertext || tag should match the published vector",
        )
    }

    @Test
    fun `round trip returns the plaintext`() {
        val key = ByteArray(32) { it.toByte() }
        val nonce = ByteArray(12) { (it * 7).toByte() }
        val message = "the quick brown fox".toByteArray()
        assertContentEquals(message, Aead.open(key, Aead.seal(key, nonce, message)))
    }

    /**
     * The layout the other two engines assume. CryptoKit's `combined` emits
     * nonce first; a box assembled any other way decrypts to garbage on iOS
     * rather than failing, which is the worst kind of interop bug.
     */
    @Test
    fun `box is nonce then ciphertext then tag`() {
        val key = ByteArray(32) { 1 }
        val nonce = ByteArray(12) { 2 }
        val message = ByteArray(5) { 3 }
        val box = Aead.seal(key, nonce, message)
        assertContentEquals(nonce, box.copyOfRange(0, 12), "the nonce must come first")
        assertEquals(12 + message.size + 16, box.size, "nonce + ciphertext + 16-byte tag")
    }

    @Test
    fun `additional data is authenticated`() {
        val key = ByteArray(32) { 4 }
        val nonce = ByteArray(12) { 5 }
        val box = Aead.seal(key, nonce, "hello".toByteArray(), aad = "context".toByteArray())
        assertContentEquals("hello".toByteArray(),
                            Aead.open(key, box, aad = "context".toByteArray()))
        assertFailsWith<Aead.DecryptionFailed> {
            Aead.open(key, box, aad = "different".toByteArray())
        }
    }

    @Test
    fun `a flipped bit anywhere is refused`() {
        val key = ByteArray(32) { 6 }
        val nonce = ByteArray(12) { 7 }
        val box = Aead.seal(key, nonce, "tamper with me".toByteArray())
        // Every position: nonce, ciphertext and tag alike.
        for (index in box.indices) {
            val broken = box.copyOf()
            broken[index] = (broken[index].toInt() xor 0x01).toByte()
            assertFailsWith<Aead.DecryptionFailed>("byte $index was not authenticated") {
                Aead.open(key, broken)
            }
        }
    }

    @Test
    fun `the wrong key is refused`() {
        val box = Aead.seal(ByteArray(32) { 8 }, ByteArray(12) { 9 }, "secret".toByteArray())
        assertFailsWith<Aead.DecryptionFailed> { Aead.open(ByteArray(32) { 10 }, box) }
    }

    /** A truncated box must be a clean refusal, not an index-out-of-bounds. */
    @Test
    fun `a box too short to hold a nonce and tag is refused cleanly`() {
        val key = ByteArray(32) { 11 }
        for (size in 0..(Aead.NONCE_LEN + Aead.TAG_LEN - 1)) {
            assertFailsWith<Aead.DecryptionFailed> { Aead.open(key, ByteArray(size)) }
        }
    }

    @Test
    fun `wrong key or nonce size is rejected up front`() {
        assertFailsWith<IllegalArgumentException> {
            Aead.seal(ByteArray(16), ByteArray(12), ByteArray(0))
        }
        assertFailsWith<IllegalArgumentException> {
            Aead.seal(ByteArray(32), ByteArray(16), ByteArray(0))
        }
    }

    /** Failures must not tell an attacker which kind of failure it was. */
    @Test
    fun `every failure reports the same thing`() {
        val key = ByteArray(32) { 12 }
        val box = Aead.seal(key, ByteArray(12) { 13 }, "x".toByteArray())
        val badKey = runCatching { Aead.open(ByteArray(32) { 14 }, box) }.exceptionOrNull()
        val badTag = runCatching {
            val broken = box.copyOf(); broken[broken.size - 1] = 0
            Aead.open(key, broken)
        }.exceptionOrNull()
        assertTrue(badKey is Aead.DecryptionFailed && badTag is Aead.DecryptionFailed)
        assertEquals(badKey.message, badTag.message, "the failures are distinguishable")
    }
}