package io.flamenet.e2e

import org.json.JSONObject
import javax.crypto.KeyGenerator
import javax.crypto.SecretKey
import kotlin.test.Test
import kotlin.test.assertEquals
import kotlin.test.assertFailsWith
import kotlin.test.assertFalse
import kotlin.test.assertNotEquals
import kotlin.test.assertNull
import kotlin.test.assertTrue

class VaultTest {
    /** Stands in for a directory on a phone, and lets a test see the raw bytes. */
    private class MapStorage : Vault.Storage {
        val records = LinkedHashMap<String, ByteArray>()

        override fun read(name: String) = records[name]

        override fun write(
            name: String,
            bytes: ByteArray,
        ) {
            records[name] = bytes
        }

        override fun delete(name: String) {
            records.remove(name)
        }

        override fun names() = records.keys.toList()
    }

    private fun key(): SecretKey = KeyGenerator.getInstance("AES").apply { init(256) }.generateKey()

    private fun bundleFor(identity: DeviceIdentity): PeerBundle {
        val oneTime = identity.generateOneTimePreKeys(1).first()
        return PeerBundle.fromWire(
            identity.registrationWire().put(
                "opk",
                JSONObject().put("id", oneTime.id).put("pub", Bytes.b64(oneTime.publicKey)),
            ),
        )
    }

    @Test
    fun `an identity survives being stored and read back`() {
        val storage = MapStorage()
        val secret = key()
        val identity = DeviceIdentity.create()
        identity.generateOneTimePreKeys(2)

        Vault(storage, secret).saveIdentity(identity)
        val restored = Vault(storage, secret).loadIdentity()!!

        assertEquals(identity.deviceId, restored.deviceId)
        assertEquals(2, restored.oneTimePreKeyCount)

        // The private halves came back too, which only a real handshake proves.
        val peer = DeviceIdentity.create()
        val session = E2ESession.initiate(peer, bundleFor(restored))
        val envelope = session.encrypt("after a cold start")
        assertEquals("after a cold start", E2ESession.respond(restored, envelope).decrypt(envelope))
    }

    @Test
    fun `a conversation survives being stored and read back`() {
        val storage = MapStorage()
        val secret = key()
        val alice = DeviceIdentity.create()
        val bob = DeviceIdentity.create()

        val a = E2ESession.initiate(alice, bundleFor(bob))
        val first = a.encrypt("hello")
        val b = E2ESession.respond(bob, first)
        b.decrypt(first)

        Vault(storage, secret).saveSession(alice.deviceId, b)
        val restored = Vault(storage, secret).loadSession(alice.deviceId)!!
        assertEquals("still talking", restored.decrypt(a.encrypt("still talking")))
        assertEquals(listOf(alice.deviceId), Vault(storage, secret).knownPeers())
    }

    /**
     * The chain a re-establishment replaces must survive it.
     */
    @Test
    fun `a retired session is kept and still opens its chain`() {
        val storage = MapStorage()
        val secret = key()
        val alice = DeviceIdentity.create()
        val bob = DeviceIdentity.create()

        // Bob starts his own chain before hearing from Alice.
        val bobsOwn = E2ESession.initiate(bob, bundleFor(alice))
        val bobsFirst = bobsOwn.encrypt("from bob")
        Vault(storage, secret).saveSession(bob.deviceId, bobsOwn)

        // Alice's message arrives on a different chain, so Alice's is adopted
        // and Bob's own is retired rather than dropped.
        val alices = E2ESession.initiate(alice, bundleFor(bob))
        val alicesFirst = alices.encrypt("from alice")
        val adopted = E2ESession.respond(bob, alicesFirst)
        assertEquals("from alice", adopted.decrypt(alicesFirst))

        val vault = Vault(storage, secret)
        vault.saveRetiredSession(bob.deviceId, vault.loadSession(bob.deviceId)!!)
        vault.saveSession(bob.deviceId, adopted)

        // Alice is still on the chain Bob started, she answers on it.
        val onTheRetiredChain = E2ESession.respond(alice, bobsFirst)
        assertEquals("from bob", onTheRetiredChain.decrypt(bobsFirst))
        val reply = onTheRetiredChain.encrypt("bob, this is the old chain")

        val fresh = Vault(storage, secret)
        assertNull(
            runCatching { fresh.loadSession(bob.deviceId)!!.decrypt(reply) }.getOrNull(),
            "the adopted session must NOT open the retired chain, otherwise this " +
                "test would pass without the retired copy existing at all",
        )
        assertEquals(
            "bob, this is the old chain",
            fresh.loadRetiredSession(bob.deviceId)!!.decrypt(reply),
        )
    }

    /**
     * A failed decrypt damages the session object it was handed, so the thing
     * that gets retired must be re-read from the vault.
     */
    @Test
    fun `the session retired after a failed decrypt must be re-read, not reused`() {
        val storage = MapStorage()
        val secret = key()
        val alice = DeviceIdentity.create()
        val bob = DeviceIdentity.create()

        // Bob's own chain, stored, with a message Alice will later answer on
        // it.
        val bobsOwn = E2ESession.initiate(bob, bundleFor(alice))
        val bobsFirst = bobsOwn.encrypt("from bob")
        Vault(storage, secret).saveSession(bob.deviceId, bobsOwn)

        val alicesReplyOnBobsChain =
            E2ESession.respond(alice, bobsFirst).let {
                it.decrypt(bobsFirst)
                it.encrypt("answering on the chain bob started")
            }

        // The receive path: Bob's stored session is handed a foreign envelope
        // and fails.
        val foreign = E2ESession.initiate(alice, bundleFor(bob)).encrypt("different chain")
        val damaged = Vault(storage, secret).loadSession(bob.deviceId)!!
        assertNull(
            runCatching { damaged.decrypt(foreign) }.getOrNull(),
            "the foreign envelope must NOT open, or this test proves nothing",
        )

        // Retiring the damaged instance loses the chain.
        val wrong = MapStorage().also { storage.records.forEach(it.records::put) }
        Vault(wrong, secret).saveRetiredSession(bob.deviceId, damaged)
        assertNull(
            runCatching {
                Vault(wrong, secret)
                    .loadRetiredSession(bob.deviceId)!!
                    .decrypt(alicesReplyOnBobsChain)
            }.getOrNull(),
            "a session damaged by a failed decrypt should NOT still open its chain, " +
                "if this ever starts passing, the hazard is gone and this test can go",
        )

        // Re-reading the vault copy keeps it.
        val pristine = Vault(storage, secret).loadSession(bob.deviceId)!!
        Vault(storage, secret).saveRetiredSession(bob.deviceId, pristine)
        assertEquals(
            "answering on the chain bob started",
            Vault(storage, secret)
                .loadRetiredSession(bob.deviceId)!!
                .decrypt(alicesReplyOnBobsChain),
        )
    }

    /** Signing out must take the retired chains too, they hold chain keys. */
    @Test
    fun `forgetting sessions forgets the retired ones as well`() {
        val storage = MapStorage()
        val secret = key()
        val alice = DeviceIdentity.create()
        val bob = DeviceIdentity.create()

        val a = E2ESession.initiate(alice, bundleFor(bob))
        Vault(storage, secret).saveSession(bob.deviceId, a)
        Vault(storage, secret).saveRetiredSession(bob.deviceId, a)

        Vault(storage, secret).forgetAllSessions()

        assertNull(Vault(storage, secret).loadSession(bob.deviceId))
        assertNull(Vault(storage, secret).loadRetiredSession(bob.deviceId))
        assertTrue(
            storage.records.keys.none { it.contains("session") },
            "a record still names a session: ${storage.records.keys}",
        )
    }

    /** What is on disk must not be readable, or none of this was worth doing. */
    @Test
    fun `nothing recognisable is written in the clear`() {
        val storage = MapStorage()
        val identity = DeviceIdentity.create()
        Vault(storage, key()).saveIdentity(identity)

        val onDisk = String(storage.records.values.first(), Charsets.ISO_8859_1)
        assertFalse(onDisk.contains(identity.deviceId))
        assertFalse(onDisk.contains("ikDH"))
        assertFalse(onDisk.contains(Bytes.b64(identity.identityDHPublic)))
    }

    /**
     * The failure this catches: a stored blob moved from one slot to another.
     */
    @Test
    fun `a record moved to another slot does not open`() {
        val storage = MapStorage()
        val secret = key()
        val vault = Vault(storage, secret)
        val alice = DeviceIdentity.create()
        val bob = DeviceIdentity.create()
        val session = E2ESession.initiate(alice, bundleFor(bob))
        vault.saveSession("peer-a", session)

        val blob = storage.records.getValue("session/peer-a")
        storage.records["session/peer-b"] = blob

        assertFailsWith<Vault.Corrupt> { vault.loadSession("peer-b") }
        // And the record in its own slot is untouched by the attempt.
        assertTrue(vault.loadSession("peer-a") != null)
    }

    @Test
    fun `a different key does not open the vault`() {
        val storage = MapStorage()
        Vault(storage, key()).saveIdentity(DeviceIdentity.create())
        assertFailsWith<Vault.Corrupt> { Vault(storage, key()).loadIdentity() }
    }

    @Test
    fun `a flipped bit is refused rather than half-read`() {
        val storage = MapStorage()
        val secret = key()
        Vault(storage, secret).saveIdentity(DeviceIdentity.create())
        val blob = storage.records.getValue("identity")
        blob[blob.size / 2] = (blob[blob.size / 2].toInt() xor 1).toByte()
        assertFailsWith<Vault.Corrupt> { Vault(storage, secret).loadIdentity() }
    }

    /**
     * A corrupt identity must NOT read as "no identity".
     */
    @Test
    fun `an unreadable identity is not mistaken for a missing one`() {
        val storage = MapStorage()
        val secret = key()
        assertNull(Vault(storage, secret).loadIdentity())
        assertFalse(Vault(storage, secret).hasIdentity())

        Vault(storage, secret).saveIdentity(DeviceIdentity.create())
        storage.records["identity"] = ByteArray(64)
        assertTrue(Vault(storage, secret).hasIdentity())
        assertFailsWith<Vault.Corrupt> { Vault(storage, secret).loadIdentity() }
    }

    @Test
    fun `the inbox watermark round-trips and starts at zero`() {
        val storage = MapStorage()
        val secret = key()
        assertEquals(0, Vault(storage, secret).inboxWatermark)
        Vault(storage, secret).inboxWatermark = 4207
        assertEquals(4207, Vault(storage, secret).inboxWatermark)
    }

    @Test
    fun `forgetting removes rather than marks`() {
        val storage = MapStorage()
        val secret = key()
        val vault = Vault(storage, secret)
        val alice = DeviceIdentity.create()
        vault.saveIdentity(alice)
        vault.saveSession("peer-a", E2ESession.initiate(alice, bundleFor(DeviceIdentity.create())))
        vault.saveSession("peer-b", E2ESession.initiate(alice, bundleFor(DeviceIdentity.create())))

        vault.forget("peer-a")
        assertFalse(storage.records.keys.any { it.contains("peer-a") })
        assertEquals(listOf("peer-b"), vault.knownPeers())

        vault.forgetAllSessions()
        assertEquals(emptyList(), vault.knownPeers())
        assertTrue(vault.hasIdentity(), "signing out of a chat is not signing out of the app")

        // Sign-out takes the identity too: the next person to use this phone
        // must not inherit a device the relay still has messages for.
        vault.forgetEverything()
        assertFalse(vault.hasIdentity())
        assertEquals(emptyList(), storage.names())
    }

    @Test
    fun `two writes of the same value do not produce the same bytes`() {
        // A fresh nonce per write.
        val storage = MapStorage()
        val secret = key()
        val identity = DeviceIdentity.create()
        Vault(storage, secret).saveIdentity(identity)
        val first = storage.records.getValue("identity").copyOf()
        Vault(storage, secret).saveIdentity(identity)
        assertNotEquals(Bytes.b64(first), Bytes.b64(storage.records.getValue("identity")))
    }

    // ---- Trust records and the platform's JSON ---------------------------.

    /**
     * A contact with no known address must come back with NO address.
     */
    @Test
    fun `a trust record with no address writes no address at all`() {
        val storage = MapStorage()
        val secret = key()
        Vault(storage, secret).saveTrust(7, Vault.Trust(listOf("aaa"), verified = false, fnid = null))

        assertNull(
            Vault(storage, secret).loadTrust(7)?.fnid,
            "an absent address came back as a value",
        )
    }

    /**
     * This is the test that reproduces the device, and the one above is not.
     */

    @Test
    fun `a record already holding the literal string null reads as no address`() {
        // Records written the old way exist on phones already.
        val storage = MapStorage()
        val secret = key()
        val vault = Vault(storage, secret)
        vault.saveTrust(7, Vault.Trust(listOf("aaa"), verified = false, fnid = "null"))

        assertNull(
            Vault(storage, secret).loadTrust(7)?.fnid,
            "the literal string \"null\" was accepted as somebody's global address",
        )
    }

    @Test
    fun `a real address still round-trips`() {
        val storage = MapStorage()
        val secret = key()
        Vault(storage, secret).saveTrust(7, Vault.Trust(listOf("aaa"), verified = true, fnid = "fnid:abc123"))
        val back = Vault(storage, secret).loadTrust(7)
        assertEquals("fnid:abc123", back?.fnid)
        assertTrue(back!!.verified)
    }
}