// Author:  cirakg & claude ai
// Summary: Key derivation: resumable PBKDF2-HMAC-SHA256 (RFC 8018) with calibration, and HKDF-SHA256 (RFC 5869).
//
// PBKDF2 turns the PIN into a key (README 7.2). It runs as a Jobs.Job, a
// bounded number of iterations per step, because the iteration count is
// calibrated to about one second on the watch, far beyond one event's
// watchdog budget. It uses the native HMAC-SHA256: the object keeps the key
// and digest() resets it, so each iteration is one native call.
//
// HKDF splits one master key into separate encryption and MAC keys
// (README 7.3).
//
// Verified with RFC 7914 section 11 (PBKDF2) and RFC 5869 A.1-A.3 (HKDF),
// plus seeded vectors from Python's hashlib (test/crypto/KdfTest.mc).

import Toybox.Cryptography;
import Toybox.Lang;
import Toybox.System;

module Kdf {

    const SHA256_BYTES = 32;
    // Bounds for the calibrated PBKDF2 iteration count. The lower bound keeps
    // a fast watch from choosing a trivially cheap PIN key; the upper bound
    // keeps the simulator (a PC) from choosing a count a watch can't afford.
    const MIN_ITERATIONS = 2000;
    const MAX_ITERATIONS = 200000;
    // Iterations per job step. The watchdog allows 120,000 bytecodes per
    // event on the Forerunner 255/255S (240,000 elsewhere); 200 iterations of
    // the byte-wise version tripped it there in the simulator (M4). With the
    // word-wise XOR this keeps a wide margin; the UI flow derives PINs on the
    // Forerunner 255S profile in every `make screenshots`.
    const ITERATIONS_PER_STEP = 100;

    // decodeNumber/encodeNumber options for one 32-bit word.
    typedef WordOptions as {:offset as Number, :endianness as Lang.Endian};

    // PBKDF2-HMAC-SHA256 as a resumable job (RFC 8018 section 5.2).
    class Pbkdf2 extends Jobs.Job {

        private var _mac as Cryptography.HashBasedMessageAuthenticationCode;
        private var _salt as ByteArray;
        private var _iterations as Number;
        private var _length as Number;
        private var _blockCount as Number;
        private var _block as Number = 1; // current output block, 1-based
        private var _done as Number = 0; // iterations done for the current block
        private var _u as ByteArray = []b;
        private var _out as ByteArray = []b;
        // T of the current block as eight 32-bit words: XOR-ing words takes
        // far fewer bytecodes than XOR-ing 32 bytes, which matters for the
        // per-event watchdog budget (see ITERATIONS_PER_STEP).
        private var _t0 as Number = 0;
        private var _t1 as Number = 0;
        private var _t2 as Number = 0;
        private var _t3 as Number = 0;
        private var _t4 as Number = 0;
        private var _t5 as Number = 0;
        private var _t6 as Number = 0;
        private var _t7 as Number = 0;
        // decodeNumber options for word 0..7, built once.
        private var _w as Array<WordOptions> = [] as Array<WordOptions>;

        // iterations >= 1; 1 <= length <= 255 * 32 bytes.
        function initialize(password as ByteArray, salt as ByteArray, iterations as Number, length as Number) {
            Job.initialize();
            if (iterations < 1 || length < 1 || length > 255 * SHA256_BYTES) {
                throw new Lang.InvalidValueException("PBKDF2: bad iteration count or key length");
            }
            _mac = new Cryptography.HashBasedMessageAuthenticationCode({
                :algorithm => Cryptography.HASH_SHA256,
                :key => password
            });
            _salt = Bytes.copy(salt);
            _iterations = iterations;
            _length = length;
            _blockCount = (length + SHA256_BYTES - 1) / SHA256_BYTES;
            for (var k = 0; k < 8; k++) {
                _w.add({:offset => k * 4, :endianness => Lang.ENDIAN_BIG});
            }
        }

        // Run up to `budget` iterations. Returns true when the key is complete.
        function step(budget as Number) as Boolean {
            var left = budget;
            var f = Lang.NUMBER_FORMAT_SINT32;
            var w = _w;
            while (left > 0 && _block <= _blockCount) {
                if (_done == 0) {
                    // U1 = HMAC(P, S || INT(i)); T = U1
                    _mac.update(Bytes.concat(_salt, Bytes.uint32BigEndian(_block)));
                    _u = _mac.digest();
                    var u = _u;
                    _t0 = u.decodeNumber(f, w[0]) as Number;
                    _t1 = u.decodeNumber(f, w[1]) as Number;
                    _t2 = u.decodeNumber(f, w[2]) as Number;
                    _t3 = u.decodeNumber(f, w[3]) as Number;
                    _t4 = u.decodeNumber(f, w[4]) as Number;
                    _t5 = u.decodeNumber(f, w[5]) as Number;
                    _t6 = u.decodeNumber(f, w[6]) as Number;
                    _t7 = u.decodeNumber(f, w[7]) as Number;
                } else {
                    // Uj = HMAC(P, Uj-1); T ^= Uj. The previous Uj isn't wiped
                    // every iteration: it is useless without the password, and
                    // a wipe per iteration would cut the iterations the watch
                    // can afford. The final values are wiped below.
                    _mac.update(_u);
                    _u = _mac.digest();
                    var u = _u;
                    _t0 ^= u.decodeNumber(f, w[0]) as Number;
                    _t1 ^= u.decodeNumber(f, w[1]) as Number;
                    _t2 ^= u.decodeNumber(f, w[2]) as Number;
                    _t3 ^= u.decodeNumber(f, w[3]) as Number;
                    _t4 ^= u.decodeNumber(f, w[4]) as Number;
                    _t5 ^= u.decodeNumber(f, w[5]) as Number;
                    _t6 ^= u.decodeNumber(f, w[6]) as Number;
                    _t7 ^= u.decodeNumber(f, w[7]) as Number;
                }
                _done++;
                left--;
                if (_done == _iterations) {
                    finishBlock();
                }
            }
            return _block > _blockCount;
        }

        // Append T to the output and clear it.
        private function finishBlock() as Void {
            var t = new [SHA256_BYTES]b;
            var words = [_t0, _t1, _t2, _t3, _t4, _t5, _t6, _t7];
            for (var k = 0; k < 8; k++) {
                t.encodeNumber(words[k] as Number, Lang.NUMBER_FORMAT_SINT32, _w[k]);
                words[k] = 0;
            }
            _out.addAll(t);
            Bytes.wipe(t);
            Bytes.wipe(_u);
            clearWords();
            _block++;
            _done = 0;
        }

        private function clearWords() as Void {
            _t0 = 0;
            _t1 = 0;
            _t2 = 0;
            _t3 = 0;
            _t4 = 0;
            _t5 = 0;
            _t6 = 0;
            _t7 = 0;
        }

        // The derived key; only valid after step() returned true.
        function result() as ByteArray {
            return _out.slice(0, _length);
        }

        // Wipe every intermediate value.
        function wipe() as Void {
            Bytes.wipe(_u);
            Bytes.wipe(_out);
            Bytes.wipe(_salt);
            clearWords();
        }
    }

    // PBKDF2-HMAC-SHA256 in one call. Only for tests and small iteration counts.
    function pbkdf2(password as ByteArray, salt as ByteArray, iterations as Number, length as Number) as ByteArray {
        var job = new Pbkdf2(password, salt, iterations, length);
        Jobs.runToCompletion(job, ITERATIONS_PER_STEP);
        var key = job.result();
        job.wipe();
        return key;
    }

    // Iteration count for which one derivation takes about targetMs on this
    // device, clamped to [MIN_ITERATIONS, MAX_ITERATIONS]. Measures a short
    // run of `sample` iterations (call it again from a job if a longer sample
    // would trip the watchdog).
    function calibrate(targetMs as Number, sample as Number) as Number {
        var job = new Pbkdf2([0, 0, 0, 0, 0, 0]b, new [16]b, sample, SHA256_BYTES);
        var start = System.getTimer();
        Jobs.runToCompletion(job, ITERATIONS_PER_STEP);
        var elapsed = System.getTimer() - start;
        job.wipe();
        var estimate = elapsed > 0 ? (sample.toLong() * targetMs / elapsed) : MAX_ITERATIONS.toLong();
        if (estimate < MIN_ITERATIONS) {
            return MIN_ITERATIONS;
        }
        return estimate > MAX_ITERATIONS ? MAX_ITERATIONS : estimate.toNumber();
    }

    // HKDF-SHA256 (RFC 5869): extract with salt (empty means 32 zero bytes),
    // then expand to `length` bytes with info. length <= 255 * 32.
    function hkdf(ikm as ByteArray, salt as ByteArray, info as ByteArray, length as Number) as ByteArray {
        if (length < 1 || length > 255 * SHA256_BYTES) {
            throw new Lang.InvalidValueException("HKDF: bad key length");
        }
        var prk = Hmac.nativeSha256(salt.size() > 0 ? salt : new [SHA256_BYTES]b, ikm);
        var okm = new [0]b;
        var block = new [0]b;
        for (var counter = 1; okm.size() < length; counter++) {
            var input = Bytes.concat(block, info);
            input.add(counter);
            Bytes.wipe(block);
            block = Hmac.nativeSha256(prk, input);
            okm.addAll(block);
        }
        Bytes.wipe(prk);
        Bytes.wipe(block);
        var key = okm.slice(0, length);
        Bytes.wipe(okm);
        return key;
    }
}