package io.flamenet.e2e

import java.security.GeneralSecurityException
import javax.crypto.Cipher
import javax.crypto.spec.GCMParameterSpec
import javax.crypto.spec.SecretKeySpec

/**
 * AES-256-GCM, in the layout the protocol uses: `nonce || ciphertext || tag`.
 *
 * That layout is not incidental. CryptoKit's `AES.GCM.SealedBox.combined` emits
 * exactly this, the JavaScript engine assembles it by hand to match, and a box
 * built any other way decrypts to garbage on the other two engines rather than
 * failing loudly. So it is pinned here and tested, not left to a convention.
 *
 * 12-byte nonces and 16-byte tags throughout — the only sizes the other engines
 * produce, and the sizes GCM is actually safe at.
 */
object Aead {

    const val NONCE_LEN = 12
    const val TAG_LEN = 16

    class DecryptionFailed(message: String) : GeneralSecurityException(message)

    /**
     * Seal [plaintext], returning `nonce || ciphertext || tag`.
     *
     * The nonce is supplied rather than generated so callers cannot accidentally
     * reuse one and so this is testable against fixed vectors. **A nonce must
     * never repeat under one key** — GCM does not degrade gracefully there, it
     * leaks the authentication subkey and with it the ability to forge.
     */
    fun seal(key: ByteArray, nonce: ByteArray, plaintext: ByteArray,
             aad: ByteArray = ByteArray(0)): ByteArray {
        require(key.size == 32) { "AES-256-GCM needs a 32-byte key, got ${key.size}" }
        require(nonce.size == NONCE_LEN) { "GCM nonce must be $NONCE_LEN bytes, got ${nonce.size}" }
        val cipher = Cipher.getInstance("AES/GCM/NoPadding")
        cipher.init(Cipher.ENCRYPT_MODE, SecretKeySpec(key, "AES"),
                    GCMParameterSpec(TAG_LEN * 8, nonce))
        if (aad.isNotEmpty()) cipher.updateAAD(aad)
        return nonce + cipher.doFinal(plaintext)
    }

    /** Open a `nonce || ciphertext || tag` box. */
    fun open(key: ByteArray, box: ByteArray, aad: ByteArray = ByteArray(0)): ByteArray {
        require(key.size == 32) { "AES-256-GCM needs a 32-byte key, got ${key.size}" }
        // Refuse a box too short to contain what it claims, rather than letting
        // the array slicing below throw something unrelated.
        if (box.size < NONCE_LEN + TAG_LEN) {
            throw DecryptionFailed("box is ${box.size} bytes, too short to hold a nonce and a tag")
        }
        val cipher = Cipher.getInstance("AES/GCM/NoPadding")
        cipher.init(Cipher.DECRYPT_MODE, SecretKeySpec(key, "AES"),
                    GCMParameterSpec(TAG_LEN * 8, box, 0, NONCE_LEN))
        if (aad.isNotEmpty()) cipher.updateAAD(aad)
        return try {
            cipher.doFinal(box, NONCE_LEN, box.size - NONCE_LEN)
        } catch (e: GeneralSecurityException) {
            // Deliberately one message for every failure. Distinguishing "bad
            // tag" from "wrong key" hands an attacker a decision oracle, and
            // there is nothing a caller can usefully do differently.
            throw DecryptionFailed("could not open the box")
        }
    }
}