// Author:  cirakg & claude ai
// Summary: The device key and the PIN-wrapped locked master key: creation, PIN derivation jobs, wrap and unwrap.
//
// Keys (README 7.1, 7.2):
//   * Device key "dk": 32 random bytes, created on first use, stored as is.
//     It seals unlocked entries, pending items and the import buffer.
//   * Locked master key (LMK): 32 random bytes, stored only wrapped under
//     the PIN key, in "w".
//   * PIN key = PBKDF2-HMAC-SHA256(PIN as ASCII digits, salt, n, 32 bytes).
//
// Wrap format, version 1 (one ByteArray under "w"):
//     salt (16 bytes) | n (4 bytes, big-endian) | SealedBox(Keys(PIN key), aad, LMK)
//     aad = "OTPW-WRAP-1" || salt || n
// A correct PIN is exactly "the box verifies"; no PIN hash exists. The salt
// and iteration count are in the aad, so changing either fails the MAC.
// Changing the PIN re-wraps the same LMK; entries aren't re-encrypted.
//
// PBKDF2 runs as a PinJob (a Jobs.Job) because one derivation is calibrated
// to about a second, far beyond one event's watchdog budget.

import Toybox.Cryptography;
import Toybox.Lang;

class KeyManager {

    static const DEVICE_KEY = "dk";
    static const WRAP_KEY = "w";
    static const KEY_BYTES = 32;
    static const SALT_BYTES = 16;
    static const WRAP_LABEL = "OTPW-WRAP-1";
    static const MIN_PIN_DIGITS = 4;
    static const MAX_PIN_DIGITS = 10;

    private var _store as Persist.Backend;
    private var _deviceKeys as SealedBox.Keys? = null;

    function initialize(store as Persist.Backend) {
        _store = store;
    }

    // The device key's derived keys, creating the device key on first use.
    function deviceKeys() as SealedBox.Keys {
        var keys = _deviceKeys;
        if (keys != null) {
            return keys;
        }
        var stored = _store.getValue(DEVICE_KEY);
        var raw = null;
        if (stored instanceof ByteArray && stored.size() == KEY_BYTES) {
            raw = stored;
        } else {
            // First start (or a lost key, which makes unlocked records
            // unreadable; they then show as damaged).
            raw = Cryptography.randomBytes(KEY_BYTES);
            _store.setValue(DEVICE_KEY, raw);
        }
        keys = new SealedBox.Keys(raw);
        Bytes.wipe(raw);
        _deviceKeys = keys;
        return keys;
    }

    // True when a PIN has been set.
    function hasPin() as Boolean {
        return _store.getValue(WRAP_KEY) != null;
    }

    // A fresh random master key (when the first PIN is created).
    function newMasterKey() as ByteArray {
        return Cryptography.randomBytes(KEY_BYTES);
    }

    // True for 4 to 10 ASCII digits.
    static function isValidPin(pin as ByteArray) as Boolean {
        if (pin.size() < MIN_PIN_DIGITS || pin.size() > MAX_PIN_DIGITS) {
            return false;
        }
        for (var i = 0; i < pin.size(); i++) {
            if (pin[i] < 0x30 || pin[i] > 0x39) {
                return false;
            }
        }
        return true;
    }

    // The derivation for checking `pin` against the stored wrap. Throws
    // NO_PIN when no PIN is set, DAMAGED when the wrap is malformed.
    function unlockJob(pin as ByteArray) as PinJob {
        var wrap = _store.getValue(WRAP_KEY);
        if (wrap == null) {
            throw new VaultError(VaultError.NO_PIN, "no PIN set");
        }
        if (!(wrap instanceof ByteArray) || wrap.size() < SALT_BYTES + 4 + SealedBox.MIN_SIZE) {
            throw new VaultError(VaultError.DAMAGED, "PIN data damaged");
        }
        var iterations = wrap.decodeNumber(Lang.NUMBER_FORMAT_UINT32, {:offset => SALT_BYTES, :endianness => Lang.ENDIAN_BIG}) as Number;
        if (iterations < Kdf.MIN_ITERATIONS || iterations > Kdf.MAX_ITERATIONS) {
            // Also stops a tampered count from making the derivation endless.
            throw new VaultError(VaultError.DAMAGED, "PIN data damaged");
        }
        return new PinJob(pin, wrap.slice(0, SALT_BYTES), iterations);
    }

    // The derivation for a new PIN, with a fresh salt.
    function newPinJob(pin as ByteArray, iterations as Number) as PinJob {
        if (iterations < Kdf.MIN_ITERATIONS || iterations > Kdf.MAX_ITERATIONS) {
            throw new Lang.InvalidValueException("PIN: iteration count out of range");
        }
        return new PinJob(pin, Cryptography.randomBytes(SALT_BYTES), iterations);
    }

    // After a finished unlockJob: the master key, or null for a wrong PIN.
    // Wipes the PIN key.
    function unwrap(job as PinJob) as ByteArray? {
        var wrap = _store.getValue(WRAP_KEY);
        if (!(wrap instanceof ByteArray) || wrap.size() < SALT_BYTES + 4 + SealedBox.MIN_SIZE) {
            job.wipe();
            return null;
        }
        var keys = job.keys();
        var masterKey = SealedBox.open(keys, wrapAad(job), wrap.slice(SALT_BYTES + 4, wrap.size()));
        keys.wipe();
        job.wipe();
        if (masterKey != null && masterKey.size() != KEY_BYTES) {
            Bytes.wipe(masterKey);
            return null;
        }
        return masterKey;
    }

    // After a finished newPinJob: wrap masterKey under the new PIN and store
    // it, replacing any previous wrap. Wipes the PIN key, not masterKey.
    function storeWrap(job as PinJob, masterKey as ByteArray) as Void {
        var keys = job.keys();
        var wrap = Bytes.copy(job.salt);
        wrap.addAll(Bytes.uint32BigEndian(job.iterations));
        wrap.addAll(SealedBox.seal(keys, wrapAad(job), masterKey));
        keys.wipe();
        job.wipe();
        _store.setValue(WRAP_KEY, wrap);
    }

    // Forget the PIN (reset or erase). The master key is then gone for good.
    function deleteWrap() as Void {
        _store.deleteValue(WRAP_KEY);
    }

    private function wrapAad(job as PinJob) as ByteArray {
        var aad = Bytes.fromUtf8(WRAP_LABEL);
        aad.addAll(job.salt);
        aad.addAll(Bytes.uint32BigEndian(job.iterations));
        return aad;
    }
}

// PBKDF2 of a PIN as a resumable job, remembering its salt and count.
class PinJob extends Jobs.Job {

    var salt as ByteArray;
    var iterations as Number;
    private var _pbkdf2 as Kdf.Pbkdf2;

    function initialize(pin as ByteArray, salt as ByteArray, iterations as Number) {
        Job.initialize();
        if (!KeyManager.isValidPin(pin)) {
            throw new Lang.InvalidValueException("PIN: 4 to 10 digits");
        }
        self.salt = salt;
        self.iterations = iterations;
        _pbkdf2 = new Kdf.Pbkdf2(pin, salt, iterations, KeyManager.KEY_BYTES);
    }

    function step(budget as Number) as Boolean {
        return _pbkdf2.step(budget);
    }

    // The wrap keys derived from the PIN key; only after step() returned true.
    function keys() as SealedBox.Keys {
        var pinKey = _pbkdf2.result();
        var keys = new SealedBox.Keys(pinKey);
        Bytes.wipe(pinKey);
        return keys;
    }

    function wipe() as Void {
        _pbkdf2.wipe();
    }
}