package io.flamenet.e2e

import kotlin.test.Test
import kotlin.test.assertContentEquals
import kotlin.test.assertFalse
import kotlin.test.assertTrue

/**
 * X3DH, both directions.
 */
class X3DHTest {
    private fun raw(kp: java.security.KeyPair) = X25519.rawFromPublicKey(kp.public)

    @Test
    fun `both sides derive the same secret, with a one-time prekey`() {
        val aliceIdentity = X25519.generate()
        val bobIdentity = X25519.generate()
        val bobSigned = X25519.generate()
        val bobOneTime = X25519.generate()

        val alice =
            X3DH.initiate(
                identityPrivateDH = aliceIdentity.private,
                bundle =
                    X3DH.Bundle(
                        identityKeyDH = raw(bobIdentity),
                        signedPreKey = raw(bobSigned),
                        oneTimePreKey = raw(bobOneTime),
                    ),
            )
        val bob =
            X3DH.respond(
                identityPrivateDH = bobIdentity.private,
                signedPreKeyPrivate = bobSigned.private,
                oneTimePrivate = bobOneTime.private,
                initiatorIdentityDH = raw(aliceIdentity),
                initiatorEphemeral = alice.ephemeralPublic,
            )
        assertContentEquals(alice.sharedSecret, bob, "the two sides disagree")
        assertTrue(alice.sharedSecret.size == 32)
    }

    @Test
    fun `both sides agree when the one-time prekey pool is drained`() {
        // A degraded but valid state: initial forward secrecy is lost, and the
        // session must still work.
        val aliceIdentity = X25519.generate()
        val bobIdentity = X25519.generate()
        val bobSigned = X25519.generate()

        val alice =
            X3DH.initiate(
                identityPrivateDH = aliceIdentity.private,
                bundle = X3DH.Bundle(raw(bobIdentity), raw(bobSigned), null),
            )
        val bob =
            X3DH.respond(
                identityPrivateDH = bobIdentity.private,
                signedPreKeyPrivate = bobSigned.private,
                oneTimePrivate = null,
                initiatorIdentityDH = raw(aliceIdentity),
                initiatorEphemeral = alice.ephemeralPublic,
            )
        assertContentEquals(alice.sharedSecret, bob)
    }

    @Test
    fun `using a one-time prekey changes the result`() {
        // Otherwise DH4 is contributing nothing and its absence would be
        // undetectable, which would make a drained pool silently equivalent to
        // a healthy one.
        val aliceIdentity = X25519.generate()
        val bobIdentity = X25519.generate()
        val bobSigned = X25519.generate()
        val bobOneTime = X25519.generate()
        val ephemeral = X25519.generate()

        val withOpk =
            X3DH.initiate(
                aliceIdentity.private,
                X3DH.Bundle(raw(bobIdentity), raw(bobSigned), raw(bobOneTime)),
                ephemeral,
            )
        val withoutOpk =
            X3DH.initiate(
                aliceIdentity.private,
                X3DH.Bundle(raw(bobIdentity), raw(bobSigned), null),
                ephemeral,
            )
        assertFalse(withOpk.sharedSecret.contentEquals(withoutOpk.sharedSecret))
    }

    @Test
    fun `a different responder derives a different secret`() {
        val aliceIdentity = X25519.generate()
        val bobIdentity = X25519.generate()
        val bobSigned = X25519.generate()
        val malloryIdentity = X25519.generate()
        val mallorySigned = X25519.generate()

        val toBob =
            X3DH.initiate(aliceIdentity.private, X3DH.Bundle(raw(bobIdentity), raw(bobSigned), null))
        val malloryTries =
            X3DH.respond(
                identityPrivateDH = malloryIdentity.private,
                signedPreKeyPrivate = mallorySigned.private,
                oneTimePrivate = null,
                initiatorIdentityDH = raw(aliceIdentity),
                initiatorEphemeral = toBob.ephemeralPublic,
            )
        assertFalse(toBob.sharedSecret.contentEquals(malloryTries))
    }
}