package com.peachpasswords.vault

import com.github.luben.zstd.Zstd
import org.bouncycastle.crypto.digests.Blake3Digest
import org.bouncycastle.crypto.generators.Argon2BytesGenerator
import org.bouncycastle.crypto.params.Argon2Parameters
import org.junit.Assert.assertEquals
import org.junit.Assert.assertThrows
import org.junit.Assert.assertTrue
import org.junit.Test
import javax.crypto.Cipher
import javax.crypto.spec.GCMParameterSpec
import javax.crypto.spec.SecretKeySpec

/**
 * M-02 regression tests: PeachScript parser must reject attacker-controlled
 * lengths/counts before any unbounded allocation. Each test crafts a payload
 * with a malicious header field and asserts that [PeachPayloadCodec.parsePayload]
 * refuses it with the expected error code.
 */
class PeachScriptHardeningTest {

    private val passphrase = "test-passphrase"

    @Test fun oversized_envelope_rejected() {
        // First guard fires before any allocation, so this just needs to be
        // above MAX_ENVELOPE_BYTES.
        val env = ByteArray(50 * 1024 * 1024 + 1)
        val ex = assertThrows(PeachPayloadCodec.PeachScriptException::class.java) {
            PeachPayloadCodec.parsePayload(env, passphrase)
        }
        assertEquals("PAYLOAD_TOO_LARGE", ex.message)
    }

    @Test fun declared_length_exceeds_actual_remaining_rejected() {
        // v3 plaintext whose declared innerLen (1 MiB) is well under the cap
        // but vastly larger than the actual plaintext — must trip the
        // remaining-bytes check without overflowing Int arithmetic.
        val plaintext = ByteArray(64).apply {
            this[0] = 0x03 // FORMAT_V3
            writeU32LE(this, 1, 1024L * 1024) // declared innerLen
        }
        val env = encryptPlaintext(plaintext, passphrase)
        val ex = assertThrows(PeachPayloadCodec.PeachScriptException::class.java) {
            PeachPayloadCodec.parsePayload(env, passphrase)
        }
        assertTrue(
            "expected PAYLOAD_TOO_SHORT or INNER_TOO_LARGE, got ${ex.message}",
            ex.message == "PAYLOAD_TOO_SHORT" || ex.message == "INNER_TOO_LARGE",
        )
    }

    @Test fun decompression_bomb_rejected() {
        // Section header declares a 4 GiB uncompressed length but the compressed
        // payload is tiny — must be rejected before Zstd allocates the bomb size.
        val body = ByteArray(8 + 9 + 16).apply {
            writeU32BE(this, 0, 0x50434853L) // MAGIC
            this[4] = 0x01 // VERSION
            this[5] = 0x01 // sectionCount
            this[8] = 0x01 // SECTION_ENTRIES
            writeU32LE(this, 9, 0xFFFFFFFFL) // rawLen = 4 GiB (bomb)
            writeU32LE(this, 13, 16L) // compLen = 16
        }
        val env = encryptPlaintext(v3Plaintext(body), passphrase)
        val ex = assertThrows(PeachPayloadCodec.PeachScriptException::class.java) {
            PeachPayloadCodec.parsePayload(env, passphrase)
        }
        assertEquals("SECTION_TOO_LARGE", ex.message)
    }

    @Test fun entry_count_overflow_rejected() {
        // Entries section declares a 4 GiB count — must be rejected before the
        // ArrayList pre-allocation runs.
        val entriesRaw = ByteArray(4).apply { writeU32LE(this, 0, 0xFFFFFFFFL) }
        val compressed = Zstd.compress(entriesRaw, 19)
        val body = ByteArray(8 + 9 + compressed.size).apply {
            writeU32BE(this, 0, 0x50434853L)
            this[4] = 0x01
            this[5] = 0x01
            this[8] = 0x01 // SECTION_ENTRIES
            writeU32LE(this, 9, entriesRaw.size.toLong())
            writeU32LE(this, 13, compressed.size.toLong())
            System.arraycopy(compressed, 0, this, 17, compressed.size)
        }
        val env = encryptPlaintext(v3Plaintext(body), passphrase)
        val ex = assertThrows(PeachPayloadCodec.PeachScriptException::class.java) {
            PeachPayloadCodec.parsePayload(env, passphrase)
        }
        assertEquals("ENTRY_COUNT_OVERFLOW", ex.message)
    }

    @Test fun oversize_passkey_blob_rejected() {
        // One passkey entry whose blobLen is 4 GiB — must be rejected before
        // r.read() allocates the declared size.
        // Entry header layout: type(1) | id(16) | flags(1) | flags2(1) |
        // createdSec(4) | updatedDelta(4) | deletedSec(4) | titleIdx(2) | blobLen(4)
        val entry = ByteArray(33 + 4).apply {
            this[0] = 0x04 // ENTRY_TYPE_PASSKEY
            writeU32LE(this, 33, 0xFFFFFFFFL) // blobLen = 4 GiB
        }
        val entriesRaw = ByteArray(6 + entry.size).apply {
            writeU32LE(this, 0, 1L) // count = 1
            writeU16LE(this, 4, 0) // empty string table
            System.arraycopy(entry, 0, this, 6, entry.size)
        }
        val compressed = Zstd.compress(entriesRaw, 19)
        val body = ByteArray(8 + 9 + compressed.size).apply {
            writeU32BE(this, 0, 0x50434853L)
            this[4] = 0x01
            this[5] = 0x01
            this[8] = 0x01
            writeU32LE(this, 9, entriesRaw.size.toLong())
            writeU32LE(this, 13, compressed.size.toLong())
            System.arraycopy(compressed, 0, this, 17, compressed.size)
        }
        val env = encryptPlaintext(v3Plaintext(body), passphrase)
        val ex = assertThrows(PeachPayloadCodec.PeachScriptException::class.java) {
            PeachPayloadCodec.parsePayload(env, passphrase)
        }
        assertEquals("PASSKEY_BLOB_TOO_LARGE", ex.message)
    }

    // ---- crypto / framing helpers (mirror production wire format) ---------

    private fun v3Plaintext(body: ByteArray): ByteArray {
        val checksum = blake3Truncated(body)
        val inner = body + checksum
        val out = ByteArray(5 + inner.size)
        out[0] = 0x03 // FORMAT_V3
        writeU32LE(out, 1, inner.size.toLong())
        System.arraycopy(inner, 0, out, 5, inner.size)
        return out
    }

    private fun encryptPlaintext(plaintext: ByteArray, passphrase: String): ByteArray {
        val salt = ByteArray(16).also { java.security.SecureRandom().nextBytes(it) }
        val nonce = ByteArray(12).also { java.security.SecureRandom().nextBytes(it) }
        val key = deriveKey(passphrase, salt)
        val ctTag = try {
            val cipher = Cipher.getInstance("AES/GCM/NoPadding")
            cipher.init(Cipher.ENCRYPT_MODE, SecretKeySpec(key, "AES"), GCMParameterSpec(128, nonce))
            cipher.doFinal(plaintext)
        } finally { key.fill(0) }
        return salt + nonce + ctTag
    }

    private fun deriveKey(passphrase: String, salt: ByteArray): ByteArray {
        val params = Argon2Parameters.Builder(Argon2Parameters.ARGON2_id)
            .withSalt(salt).withParallelism(4).withMemoryAsKB(65536).withIterations(3).build()
        val gen = Argon2BytesGenerator().apply { init(params) }
        val out = ByteArray(32)
        gen.generateBytes(passphrase.toByteArray(Charsets.UTF_8), out)
        return out
    }

    private fun blake3Truncated(data: ByteArray): ByteArray {
        val digest = Blake3Digest(256)
        digest.update(data, 0, data.size)
        val full = ByteArray(32)
        digest.doFinal(full, 0)
        return full.copyOfRange(0, 4)
    }

    private fun writeU32LE(buf: ByteArray, off: Int, v: Long) {
        buf[off] = (v and 0xff).toByte()
        buf[off + 1] = ((v ushr 8) and 0xff).toByte()
        buf[off + 2] = ((v ushr 16) and 0xff).toByte()
        buf[off + 3] = ((v ushr 24) and 0xff).toByte()
    }

    private fun writeU32BE(buf: ByteArray, off: Int, v: Long) {
        buf[off] = ((v ushr 24) and 0xff).toByte()
        buf[off + 1] = ((v ushr 16) and 0xff).toByte()
        buf[off + 2] = ((v ushr 8) and 0xff).toByte()
        buf[off + 3] = (v and 0xff).toByte()
    }

    private fun writeU16LE(buf: ByteArray, off: Int, v: Int) {
        buf[off] = (v and 0xff).toByte()
        buf[off + 1] = ((v ushr 8) and 0xff).toByte()
    }
}