flamenet-messenger-android/obfs4/src/test/kotlin/io/flamenet/obfs4/LadderTest.kt

package io.flamenet.obfs4

import org.json.JSONObject
import kotlin.test.Test
import kotlin.test.assertEquals
import kotlin.test.assertNotNull
import kotlin.test.assertTrue

/**
 * Cofactor-dirty scalar base multiplication, against lyrebird.
 *
 * 🔴 **The wrong answer sits right next to the right one.** The clean X25519
 * public key — what every library computes — is a perfectly good Curve25519
 * key. It agrees on shared secrets. It round-trips. It is simply confined to
 * the prime-order subgroup and therefore *distinguishable*, which defeats the
 * entire purpose. Clamping the private key before the call is all it takes to
 * get one, so this checks structure as well as bytes.
 */
class LadderTest {

    private fun vectors(): JSONObject {
        val stream = javaClass.classLoader.getResourceAsStream("elligator-vectors.json")
        assertNotNull(stream, "the vectors are missing from test resources")
        return JSONObject(stream.readBytes().decodeToString())
    }

    private fun hex(s: String) = ByteArray(s.length / 2) {
        s.substring(it * 2, it * 2 + 2).toInt(16).toByte()
    }
    private fun hex(b: ByteArray) = b.joinToString("") { "%02x".format(it) }

    /**
     * ⚠️ The `private` field is the 32 bytes lyrebird actually fed to
     * `ScalarBaseMult` — already past its SHA-512 pass, and *not* clamped.
     * Clamping before this call would clear the cofactor and quietly turn this
     * into the clean function.
     */
    @Test
    fun `the dirty public key matches lyrebird`() {
        val keygen = vectors().getJSONArray("keygen")
        assertTrue(keygen.length() > 100, "too few vectors to mean anything")

        val mismatches = mutableListOf<String>()
        for (i in 0 until keygen.length()) {
            val v = keygen.getJSONObject(i)
            val got = Ladder.dirtyPublicKey(hex(v.getString("private")))
            if (hex(got) != v.getString("public")) {
                mismatches += "private ${v.getString("private")}\n" +
                    "  want ${v.getString("public")}\n  got  ${hex(got)}"
            }
        }
        assertTrue(
            mismatches.isEmpty(),
            "${mismatches.size} of ${keygen.length()} private keys produced the wrong " +
                "public key.\n\n${mismatches.take(2).joinToString("\n")}",
        )
    }

    /**
     * 🔴 The structural check. If the cofactor were being cleared this would be
     * the clean key, so it must differ from the clean key exactly when the low
     * three bits of the private key are non-zero, and match it exactly when
     * they are not.
     */
    @Test
    fun `it differs from the clean key exactly when the cofactor is set`() {
        val keygen = vectors().getJSONArray("keygen")
        val wrong = mutableListOf<String>()
        var withCofactor = 0
        var without = 0

        for (i in 0 until keygen.length()) {
            val priv = hex(keygen.getJSONObject(i).getString("private"))
            val dirty = Ladder.dirtyPublicKey(priv)
            val clean = Ladder.publicKey(priv)
            if ((priv[0].toInt() and 7) != 0) {
                withCofactor++
                if (dirty.contentEquals(clean)) wrong += "cofactor set but the key is clean"
            } else {
                without++
                if (!dirty.contentEquals(clean)) wrong += "no cofactor but the keys differ"
            }
        }
        assertTrue(wrong.isEmpty(), "${wrong.size} keys had the wrong relationship to the clean key")
        // ⚠️ Both arms must actually occur, or the assertion above is vacuous
        // for one of them.
        assertTrue(withCofactor > 10, "no keys with a cofactor component")
        assertTrue(without > 10, "no keys without a cofactor component")
    }

    @Test
    fun `the whole keypair matches lyrebird end to end`() {
        val keygen = vectors().getJSONArray("keygen")
        var mismatches = 0
        for (i in 0 until keygen.length()) {
            val v = keygen.getJSONObject(i)
            val pub = Ladder.dirtyPublicKey(hex(v.getString("private")))
            val want = hex(v.getString("representative"))
                .also { it[31] = (it[31].toInt() and 0x3f).toByte() }
            val matched = listOf(0, 1).any { tweak ->
                Elligator2.representativeForPublicKey(pub, tweak)
                    ?.also { it[31] = (it[31].toInt() and 0x3f).toByte() }
                    ?.contentEquals(want) == true
            }
            if (!matched || hex(pub) != v.getString("public")) mismatches++
        }
        assertEquals(0, mismatches, "$mismatches keypairs did not match lyrebird end to end")
    }

    @Test
    fun `trimScalar clamps the standard bits`() {
        val all = ByteArray(32) { 0xff.toByte() }
        val t = Ladder.trimScalar(all)
        assertEquals(0xf8, t[0].toInt() and 0xff, "the low three bits must be cleared")
        assertEquals(0x7f, t[31].toInt() and 0xff, "bit 255 cleared and bit 254 set")
        assertEquals(0x40, Ladder.trimScalar(ByteArray(32))[31].toInt() and 0xff)
    }

    @Test
    fun `adding the group order is driven by the low three bits`() {
        val s = ByteArray(32) { 0x11 }
        assertTrue(Ladder.addGroupOrderTimes(0xf8, s).contentEquals(s),
                   "low three bits clear must leave the scalar alone")
        for (x in 1..7) {
            assertTrue(!Ladder.addGroupOrderTimes(x, s).contentEquals(s), "x=$x must change it")
        }
    }
}