package dev.privacyllc.period.core.security

import java.security.MessageDigest
import java.security.SecureRandom
import javax.crypto.SecretKeyFactory
import javax.crypto.spec.PBEKeySpec

/**
 * The three outcomes of checking a PIN, and they are three rather than two on
 * purpose.
 *
 * [KeyUnavailable] is the one that matters. If the Keystore entry has gone —
 * a restored device, a cleared keystore, an OEM upgrade that dropped it — then
 * no PIN can ever verify again. Reporting that as [Wrong] would send the user
 * into a backoff that gets longer forever while they type a PIN that was always
 * correct, and under the no-recovery policy they would have no way to learn
 * otherwise. It has to be a different branch with different words on screen.
 */
internal sealed interface PinCheck {

    /**
     * @param rehashed a record at the current cost, when the stored one was
     *   cheaper. Non-null only here, because the PIN exists in memory only at
     *   this instant — see [VerifierRecord].
     */
    data class Correct(val rehashed: VerifierRecord?) : PinCheck

    data object Wrong : PinCheck

    data object KeyUnavailable : PinCheck
}

/**
 * Turns a PIN into something that can be checked without storing it.
 *
 * ```
 * stretched = PBKDF2WithHmacSHA256(pin, salt, iterations, 256 bits)
 * mac       = HMAC-SHA256(keystoreKey, 0x01 || salt || stretched)
 * ```
 *
 * Two layers, because they defend against different things and neither is
 * sufficient alone.
 *
 * **The Keystore MAC is what makes a six-digit PIN safe.** A PIN has at most a
 * million values; any hash an attacker can compute offline falls in seconds, at
 * any iteration count. The key here cannot be exported from the Android
 * Keystore, so an attacker holding a copy of the app's files cannot try even one
 * candidate without the physical device.
 *
 * **PBKDF2 underneath is for the day that assumption breaks.** If the key ever
 * leaks, the stretch is the only thing between the record and the PIN.
 *
 * `0x01` is a domain-separation tag. The same Keystore key also MACs the
 * lockout counter under `0x02`, and without a tag a value from one context could
 * be replayed as the other.
 */
internal class PinVerifier(
    private val macs: MacProvider,
    private val random: SecureRandom = SecureRandom(),
) {

    /** Creates the key if absent, and returns the record to persist. */
    fun enroll(pin: CharArray, iterations: Int = TARGET_ITERATIONS): VerifierRecord {
        macs.ensureKey()
        val salt = ByteArray(SALT_BYTES).also(random::nextBytes)
        return VerifierRecord(iterations, salt, macFor(pin, salt, iterations))
    }

    fun verify(pin: CharArray, record: VerifierRecord): PinCheck {
        if (!macs.hasKey()) return PinCheck.KeyUnavailable

        val candidate = macFor(pin, record.salt, record.iterations)
        // MessageDigest.isEqual, not contentEquals: this is the one comparison
        // in the app where an early return leaks how much of a guess was right.
        // Hex-string comparison would leak the same way and is the usual mistake.
        if (!MessageDigest.isEqual(candidate, record.mac)) return PinCheck.Wrong

        // The PIN is in memory exactly here and nowhere else, so this is the
        // only moment a cheaper record can be brought up to the current cost.
        val rehashed = if (record.iterations < TARGET_ITERATIONS) enroll(pin) else null
        return PinCheck.Correct(rehashed)
    }

    private fun macFor(pin: CharArray, salt: ByteArray, iterations: Int): ByteArray {
        var stretched: ByteArray? = null
        try {
            stretched = stretch(pin, salt, iterations)
            return macs.mac(byteArrayOf(TAG_VERIFIER) + salt + stretched)
        } finally {
            // Best-effort, and worth doing anyway: these are ByteArrays we own,
            // unlike the String a text field hands us, which is immutable and
            // cannot be cleared at all. The KDoc on the lock screen explains
            // why a custom keypad is still not worth the accessibility cost.
            stretched?.fill(0)
        }
    }

    private fun stretch(pin: CharArray, salt: ByteArray, iterations: Int): ByteArray {
        val spec = PBEKeySpec(pin, salt, iterations, STRETCHED_BITS)
        try {
            return SecretKeyFactory.getInstance(PBKDF2).generateSecret(spec).encoded
        } finally {
            spec.clearPassword()
        }
    }

    companion object {
        /**
         * `PBKDF2WithHmacSHA256` arrived in API 26, which is this project's
         * minSdk exactly — so this is the one algorithm choice with no margin,
         * and it is checked on the `PeriodMinSdk26` emulator rather than assumed.
         */
        const val PBKDF2 = "PBKDF2WithHmacSHA256"

        const val SALT_BYTES = 16
        const val STRETCHED_BITS = 256

        /**
         * Chosen for an unlock a person waits through, not for a password
         * database. The usual published figures assume the hash can be attacked
         * offline; here it cannot, because the MAC key never leaves the
         * Keystore — so this is a second line rather than the only one, and
         * buying it at a second of latency on an old phone would be paying in
         * the wrong currency. Raise it by bumping this: a cheaper stored record
         * is re-derived on the next successful unlock.
         */
        const val TARGET_ITERATIONS = 210_000

        /** Domain-separation tags. The lockout counter uses [TAG_BACKOFF]. */
        const val TAG_VERIFIER: Byte = 0x01
        const val TAG_BACKOFF: Byte = 0x02
    }
}
