package invalid.lena.rsend

import android.security.keystore.KeyGenParameterSpec
import android.security.keystore.KeyProperties
import java.security.KeyStore
import javax.crypto.Cipher
import javax.crypto.KeyGenerator
import javax.crypto.SecretKey
import javax.crypto.spec.GCMParameterSpec

// KeyVault encrypts the SSH private key at rest with an AES-GCM key in the
// Android Keystore. The Keystore key cannot be exported, so the blob
// on disk is useless off this device. Blob layout is IV || ciphertext+tag.
object KeyVault {

    private const val ALIAS = "rsend-key"
    private const val TRANSFORM = "AES/GCM/NoPadding"
    private const val IV_LEN = 12
    private const val TAG_BITS = 128

    private fun existing(): SecretKey? {
        val ks = KeyStore.getInstance("AndroidKeyStore").apply { load(null) }
        return (ks.getEntry(ALIAS, null) as? KeyStore.SecretKeyEntry)?.secretKey
    }

    private fun create(): SecretKey {
        val gen = KeyGenerator.getInstance(KeyProperties.KEY_ALGORITHM_AES, "AndroidKeyStore")
        gen.init(
            KeyGenParameterSpec.Builder(
                ALIAS,
                KeyProperties.PURPOSE_ENCRYPT or KeyProperties.PURPOSE_DECRYPT,
            )
                .setBlockModes(KeyProperties.BLOCK_MODE_GCM)
                .setEncryptionPaddings(KeyProperties.ENCRYPTION_PADDING_NONE)
                .build(),
        )
        return gen.generateKey()
    }

    fun encrypt(plain: ByteArray): ByteArray {
        val c = Cipher.getInstance(TRANSFORM)
        c.init(Cipher.ENCRYPT_MODE, existing() ?: create())
        return c.iv + c.doFinal(plain)
    }

    // Decrypting must never mint a key: doing so would replace the one that
    // wrapped the stored blob, making it undecryptable forever while looking
    // like an ordinary cipher failure. Say what happened instead, so the user
    // is told to import or generate the identity key again.
    fun decrypt(blob: ByteArray): ByteArray {
        require(blob.size >= IV_LEN + TAG_BITS / 8) { "encrypted private key is truncated" }
        val key = existing() ?: throw IllegalStateException(
            "the device keystore entry for rsend is gone; generate or import the key again"
        )
        val c = Cipher.getInstance(TRANSFORM)
        c.init(Cipher.DECRYPT_MODE, key, GCMParameterSpec(TAG_BITS, blob, 0, IV_LEN))
        return c.doFinal(blob, IV_LEN, blob.size - IV_LEN)
    }
}
