package com.dangerousthings.aliro;

import javacard.framework.JCSystem;
import javacard.framework.Util;
import javacard.security.CryptoException;
import javacard.security.HMACKey;
import javacard.security.KeyBuilder;
import javacard.security.MessageDigest;
import javacard.security.Signature;

/**
 * HMAC-SHA-256 with runtime native-or-userland selection.
 *
 * <p>The constructor probes for {@link Signature#ALG_HMAC_SHA_256} plus
 * {@link KeyBuilder#TYPE_HMAC_TRANSIENT_DESELECT} at length
 * {@link KeyBuilder#LENGTH_HMAC_SHA_256_BLOCK_64}. If both succeed (J3R452 and
 * other newer NXP J3R-family cards), {@link #compute} delegates to the native
 * Signature engine, which is dramatically faster than computing HMAC by hand
 * on top of SHA-256 in JC bytecode (measured ~5 ms native vs ~30 ms userland
 * per call on the J3R-family). If the probe throws CryptoException (J3R180
 * lacks any TYPE_HMAC slot at any tested size, and lacks ALG_HMAC_SHA_256
 * entirely), we fall back to the original RFC 2104 implementation on top of
 * {@link MessageDigest#ALG_SHA_256}, which is present on both.
 *
 * <p>HMAC per RFC 2104:
 * <pre>
 *   HMAC(K, M) = SHA-256((K_pad XOR opad) || SHA-256((K_pad XOR ipad) || M))
 * </pre>
 * with {@code ipad = 0x36 × 64} and {@code opad = 0x5C × 64}, key padded
 * with zeros to the 64-byte SHA-256 block size (or pre-hashed if &gt; 64 B).
 * Output is always 32 bytes.
 *
 * <p>Intended to be instantiated once per applet and reused — both the native
 * Signature/HMACKey objects and the userland MessageDigest are expensive to
 * allocate.
 */
final class AliroHmac {

    private static final short BLOCK_LEN = 64;   // SHA-256 block size
    private static final short HASH_LEN  = 32;   // SHA-256 output size

    private static final byte IPAD_BYTE = (byte) 0x36;
    private static final byte OPAD_BYTE = (byte) 0x5C;

    // Scratch layout: 64B pad block || 32B inner digest = 96B total. Reused
    // by both paths -- the native path uses OFF_PAD as the zero-padded key
    // staging buffer (HMACKey requires a fixed 64-byte slot).
    private static final short OFF_PAD   = 0;
    private static final short OFF_INNER = 64;
    private static final short SCRATCH_LEN = 96;

    /** 1 if the constructor successfully allocated native HMAC primitives, 0
     *  if we fell back to the SHA-256-on-top-of-MessageDigest userland path. */
    private final byte usesNative;

    // Native path (J3R452 etc.) -- null when usesNative == 0.
    private Signature hmacSig;
    private HMACKey hmacKey;

    // Always allocated: needed by the userland path AND by the native path's
    // ">64 byte key" pre-hash branch.
    private final MessageDigest sha256;
    private final byte[] scratch;

    AliroHmac() {
        sha256 = MessageDigest.getInstance(MessageDigest.ALG_SHA_256, false);
        scratch = JCSystem.makeTransientByteArray(SCRATCH_LEN, JCSystem.CLEAR_ON_DESELECT);

        // Probe for native HMAC end-to-end. The two-stage check matters:
        // JCAlgTest reports "supported" if buildKey succeeds, but some cards
        // (J3R452) accept buildKey then reject setKey with ILLEGAL_VALUE on
        // every key length we try. So we run a full smoke-test (buildKey ->
        // setKey -> init -> sign with one byte of throwaway data) and only
        // mark native if the whole flow runs without throwing. Anything that
        // misfires falls us back to userland, which always works.
        Signature tmpSig = null;
        HMACKey tmpKey = null;
        byte nativeOk = 0;
        try {
            tmpSig = Signature.getInstance(Signature.ALG_HMAC_SHA_256, false);
            tmpKey = (HMACKey) KeyBuilder.buildKey(
                KeyBuilder.TYPE_HMAC,
                KeyBuilder.LENGTH_HMAC_SHA_256_BLOCK_64, false);
            // Smoke test: zero-filled 64-byte key, sign a 1-byte message,
            // discard the digest. If any of these throw, native is unusable
            // on this card and we fall through to userland.
            Util.arrayFillNonAtomic(scratch, OFF_PAD, BLOCK_LEN, (byte) 0);
            tmpKey.setKey(scratch, OFF_PAD, BLOCK_LEN);
            tmpSig.init(tmpKey, Signature.MODE_SIGN);
            tmpSig.sign(scratch, OFF_PAD, (short) 1, scratch, OFF_INNER);
            nativeOk = 1;
        } catch (CryptoException e) {
            nativeOk = 0;
        } catch (Throwable t) {
            nativeOk = 0;
        }
        if (nativeOk == 1) {
            usesNative = 1;
            hmacSig = tmpSig;
            hmacKey = tmpKey;
        } else {
            usesNative = 0;
            // tmpSig / tmpKey become garbage-collectable; no need to retain.
        }
        // Wipe any smoke-test bytes left in scratch.
        Util.arrayFillNonAtomic(scratch, (short) 0, SCRATCH_LEN, (byte) 0);
    }

    /**
     * HMAC-SHA-256(key, msg). Writes 32 bytes to {@code out[outOff..outOff+32)}
     * and returns 32. The key is staged through internal scratch before any
     * hashing, and the message is consumed before anything is written to
     * {@code out}, so the {@code key}, {@code msg}, and {@code out} buffers
     * may alias each other freely.
     */
    short compute(
            byte[] key, short keyOff, short keyLen,
            byte[] msg, short msgOff, short msgLen,
            byte[] out, short outOff) {
        if (usesNative != 0) {
            return computeNative(key, keyOff, keyLen, msg, msgOff, msgLen, out, outOff);
        }
        return computeUserland(key, keyOff, keyLen, msg, msgOff, msgLen, out, outOff);
    }

    /**
     * Native HMAC path: stage K as a fixed 64-byte block in scratch (zero-pad
     * if shorter, SHA-256 hash if longer per RFC 2104), then setKey with the
     * full block length. The J3R452's HMACKey rejects variable-length setKey
     * with ILLEGAL_VALUE; the runtime wants the slot fully populated. RFC 2104
     * specifies zero-padding shorter keys to block size before the inner/outer
     * XOR, so zero-padding here and passing the full 64 bytes yields the
     * correct HMAC output.
     */
    private short computeNative(
            byte[] key, short keyOff, short keyLen,
            byte[] msg, short msgOff, short msgLen,
            byte[] out, short outOff) {

        // Build the 64-byte staged key in scratch[OFF_PAD..OFF_PAD+64).
        Util.arrayFillNonAtomic(scratch, OFF_PAD, BLOCK_LEN, (byte) 0);
        if (keyLen > BLOCK_LEN) {
            // K' = SHA-256(K), 32 bytes; remaining 32 bytes stay zero.
            sha256.reset();
            sha256.doFinal(key, keyOff, keyLen, scratch, OFF_PAD);
        } else if (keyLen > 0) {
            Util.arrayCopyNonAtomic(key, keyOff, scratch, OFF_PAD, keyLen);
        }
        // else keyLen == 0: scratch[OFF_PAD..OFF_PAD+64) stays all zeros,
        // which is the valid (if degenerate) HMAC key derived from the empty
        // input -- matches our userland path's behavior for empty keys.

        hmacKey.setKey(scratch, OFF_PAD, BLOCK_LEN);
        hmacSig.init(hmacKey, Signature.MODE_SIGN);
        short outLen = hmacSig.sign(msg, msgOff, msgLen, out, outOff);

        // Wipe the staged key -- HMACKey now holds it internally, but the
        // scratch copy would otherwise outlive this call until next compute().
        Util.arrayFillNonAtomic(scratch, OFF_PAD, BLOCK_LEN, (byte) 0);
        return outLen;
    }

    /**
     * Userland HMAC path (J3R180 and any card without native HMAC). Original
     * RFC 2104 implementation on top of MessageDigest.ALG_SHA_256, unchanged
     * from the pre-refactor version.
     */
    private short computeUserland(
            byte[] key, short keyOff, short keyLen,
            byte[] msg, short msgOff, short msgLen,
            byte[] out, short outOff) {

        // --- Build K_pad into scratch[OFF_PAD..OFF_PAD+64) ---
        // Zero the 64-byte pad region first so any unused key-tail bytes stay 0.
        Util.arrayFillNonAtomic(scratch, OFF_PAD, BLOCK_LEN, (byte) 0);

        if (keyLen > BLOCK_LEN) {
            // K' = SHA-256(K), 32 bytes; the remaining 32 bytes stay zero.
            sha256.reset();
            sha256.doFinal(key, keyOff, keyLen, scratch, OFF_PAD);
        } else if (keyLen > 0) {
            Util.arrayCopyNonAtomic(key, keyOff, scratch, OFF_PAD, keyLen);
        }
        // else keyLen == 0: scratch[OFF_PAD..OFF_PAD+64) stays all zeros.

        // --- Inner: SHA-256((K_pad XOR ipad) || M) into scratch[OFF_INNER..) ---
        for (short i = 0; i < BLOCK_LEN; i++) {
            scratch[(short) (OFF_PAD + i)] ^= IPAD_BYTE;
        }
        sha256.reset();
        sha256.update(scratch, OFF_PAD, BLOCK_LEN);
        if (msgLen > 0) {
            sha256.doFinal(msg, msgOff, msgLen, scratch, OFF_INNER);
        } else {
            // doFinal with length 0 -- pass the same buffer as a harmless source.
            sha256.doFinal(scratch, OFF_PAD, (short) 0, scratch, OFF_INNER);
        }

        // --- Recover K_pad from (K_pad XOR ipad), then XOR opad. ---
        // (K_pad XOR ipad) XOR (ipad XOR opad) = K_pad XOR opad, so one pass
        // over the block suffices. ipad XOR opad = 0x36 XOR 0x5C = 0x6A.
        for (short i = 0; i < BLOCK_LEN; i++) {
            scratch[(short) (OFF_PAD + i)] ^= (byte) 0x6A;
        }

        // --- Outer: SHA-256((K_pad XOR opad) || inner) into out[outOff..) ---
        sha256.reset();
        sha256.update(scratch, OFF_PAD, BLOCK_LEN);
        sha256.doFinal(scratch, OFF_INNER, HASH_LEN, out, outOff);

        // Wipe scratch -- it held K_pad (derived from the HMAC key) and the
        // inner digest, both of which would leak key-dependent state to a
        // later HMAC invocation on a different key otherwise. CLEAR_ON_DESELECT
        // cleans up at deselect, but multi-APDU sessions could see stale
        // state until then.
        Util.arrayFillNonAtomic(scratch, (short) 0, SCRATCH_LEN, (byte) 0);

        return HASH_LEN;
    }
}
