// 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: small enough for the 120,000-bytecode watchdog.
    const ITERATIONS_PER_STEP = 200;

    // 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 _t as ByteArray = []b;
        private var _out as ByteArray = []b;

        // 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;
        }

        // Run up to `budget` iterations. Returns true when the key is complete.
        function step(budget as Number) as Boolean {
            var left = budget;
            while (left > 0 && _block <= _blockCount) {
                if (_done == 0) {
                    // U1 = HMAC(P, S || INT(i))
                    _mac.update(Bytes.concat(_salt, Bytes.uint32BigEndian(_block)));
                    _u = _mac.digest();
                    _t = Bytes.copy(_u);
                } 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 loop per iteration would halve the iterations the
                    // watch can afford in its time budget. The final values
                    // are wiped below.
                    _mac.update(_u);
                    _u = _mac.digest();
                    var t = _t;
                    var u = _u;
                    for (var i = 0; i < SHA256_BYTES; i++) {
                        t[i] = t[i] ^ u[i];
                    }
                }
                _done++;
                left--;
                if (_done == _iterations) {
                    _out.addAll(_t);
                    Bytes.wipe(_t);
                    Bytes.wipe(_u);
                    _block++;
                    _done = 0;
                }
            }
            return _block > _blockCount;
        }

        // 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(_t);
            Bytes.wipe(_out);
            Bytes.wipe(_salt);
        }
    }

    // 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;
    }
}