flamenet-messenger-android/engine/src/test/kotlin/io/flamenet/e2e/PQVectorTest.kt

package io.flamenet.e2e

import org.json.JSONObject
import java.io.File
import kotlin.test.Test
import kotlin.test.assertContentEquals
import kotlin.test.assertEquals
import kotlin.test.assertNotEquals

/**
 * The hybrid handshake against the published known-answer vector.
 */
class PQVectorTest {
    private val vector: JSONObject by lazy {
        // Resolved rather than hardcoded, for the same reason
        // PublishedVectorsTest does it: this module lives both as a subtree of
        // the app repo and as a repo of its own beside the engine.
        val candidates =
            listOfNotNull(
                System.getenv("FLAMENET_PQ_VECTORS"),
                "../../engine/test/vectors/pqx3dh-v2.json",
                "../../../flamenet-e2e/test/vectors/pqx3dh-v2.json",
                "../../flamenet-e2e/test/vectors/pqx3dh-v2.json",
            )
        val found =
            candidates.map(::File).firstOrNull { it.isFile }
                ?: error(
                    "no pqx3dh-v2.json found; looked in $candidates. " +
                        "Set FLAMENET_PQ_VECTORS to point at it.",
                )
        JSONObject(found.readText())
    }

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

    /**
     * The cross-implementation check.
     */
    @Test
    fun `a published seed derives the published public key`() {
        val archive = vector.getJSONObject("bobArchive")
        for (which in listOf("pqSpk", "pqOtp")) {
            val entry =
                if (which == "pqSpk") {
                    archive.getJSONObject("pqSpk")
                } else {
                    archive.getJSONObject("pqOtp").getJSONObject("1")
                }
            val seed = Bytes.unb64(entry.getString("seed"))
            val published = Bytes.unb64(entry.getString("pub"))

            assertEquals(PQKEM.SEED_LEN, seed.size, "$which seed is d‖z")
            assertContentEquals(
                published,
                PQKEM.fromSeed(seed).publicKey,
                "$which: this engine derives a different key from the same seed",
            )
        }
    }

    /**
     * The whole responder side, end to end, against the published answer.
     */
    @Test
    fun `the responder derives the published session key and associated data`() {
        val archive = vector.getJSONObject("bobArchive")
        val env = vector.getJSONObject("env")
        val x3dh = env.getJSONObject("x3dh")
        val expect = vector.getJSONObject("expect")

        fun priv(path: JSONObject) = X25519.privateKeyFromPkcs8(Bytes.unb64(path.getString("priv")))

        val ciphertext = Bytes.unb64(env.getJSONObject("pq").getString("ct"))
        val pqSeed = Bytes.unb64(archive.getJSONObject("pqOtp").getJSONObject("1").getString("seed"))
        assertEquals(
            "opk",
            env.getJSONObject("pq").getString("kind"),
            "this vector uses the one-time PQ prekey",
        )

        val sk =
            X3DH.respond(
                identityPrivateDH = priv(archive.getJSONObject("ikDH")),
                signedPreKeyPrivate = priv(archive.getJSONObject("spk")),
                oneTimePrivate = priv(archive.getJSONObject("otp").getJSONObject("1")),
                initiatorIdentityDH = Bytes.unb64(x3dh.getString("ik_dh")),
                initiatorEphemeral = Bytes.unb64(x3dh.getString("ek")),
                pq = pqSeed to ciphertext,
            )
        assertEquals(expect.getString("sk"), hex(sk))

        val ad =
            X3DH.associatedData(
                first = Bytes.unb64(x3dh.getString("ik_sig")),
                second = Bytes.unb64(archive.getJSONObject("ikSig").getString("pub")),
                pqCiphertext = ciphertext,
            )
        assertEquals(expect.getString("ad"), hex(ad))
    }

    /**
     * A downgrade must not produce a key the other side accepts. v1's
     * `DH1‖DH2‖DH3‖DH4` and v2's `DH1‖DH2‖DH3‖SS` are both 128 bytes, so
     * without the info string changing with the version they are
     * indistinguishable KDF inputs (§12.5).
     */
    @Test
    fun `running the same handshake without the KEM derives a different key`() {
        val archive = vector.getJSONObject("bobArchive")
        val x3dh = vector.getJSONObject("env").getJSONObject("x3dh")

        fun priv(path: JSONObject) = X25519.privateKeyFromPkcs8(Bytes.unb64(path.getString("priv")))

        val classical =
            X3DH.respond(
                identityPrivateDH = priv(archive.getJSONObject("ikDH")),
                signedPreKeyPrivate = priv(archive.getJSONObject("spk")),
                oneTimePrivate = priv(archive.getJSONObject("otp").getJSONObject("1")),
                initiatorIdentityDH = Bytes.unb64(x3dh.getString("ik_dh")),
                initiatorEphemeral = Bytes.unb64(x3dh.getString("ek")),
            )
        assertNotEquals(vector.getJSONObject("expect").getString("sk"), hex(classical))
    }

    /**
     * The associated data binds the ciphertext, so a relay that swapped one
     * cannot leave the session looking untouched.
     */
    @Test
    fun `changing one byte of the ciphertext changes the associated data`() {
        val ciphertext = Bytes.unb64(vector.getJSONObject("env").getJSONObject("pq").getString("ct"))
        val tampered = ciphertext.copyOf().also { it[7] = (it[7].toInt() xor 1).toByte() }
        val a = X3DH.associatedData(ByteArray(32), ByteArray(32), ciphertext)
        val b = X3DH.associatedData(ByteArray(32), ByteArray(32), tampered)
        assertNotEquals(hex(a), hex(b))
    }

    /** A v1 session's AD is the two identity keys and nothing else. */
    @Test
    fun `without a ciphertext the associated data is just the two identities`() {
        val a = ByteArray(32) { 1 }
        val b = ByteArray(32) { 2 }
        assertContentEquals(a + b, X3DH.associatedData(a, b, null))
    }
}