// 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.
//
// A new PIN has exactly PIN_DIGITS (4) digits and is stored with the
// marker "wl" = PIN_DIGITS next to the wrap (the wrap format is unchanged).
// A wrap without the marker is a legacy PIN of 4 to 10 digits from an
// earlier version: it still unlocks, and the PIN screen then asks for a new
// 4-digit PIN (isLegacyPin).
//
// 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";
    // Every new PIN has exactly this many digits.
    static const PIN_DIGITS = 4;
    // Legacy PINs (set before the 4-digit rule) have 4 to 10 digits.
    static const MIN_PIN_DIGITS = 4;
    static const MAX_PIN_DIGITS = 10;
    // The marker of a PIN set with PIN_DIGITS digits (see the header).
    static const PIN_LENGTH_KEY = "wl";

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

    static const FINGERPRINT_BYTES = 8;

    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);
        _fingerprintKey = Hkdf.derive(raw, []b, Bytes.fromUtf8("OTPW fp v1"), KEY_BYTES);
        Bytes.wipe(raw);
        _deviceKeys = keys;
        return keys;
    }

    // An 8-byte keyed fingerprint of an entry's kind, type and secret (or
    // payload): equal for duplicates, meaningless without the device key.
    // Stored in the readable header so imports can find duplicates of locked
    // entries while the watch is locked (docs/SECURITY.md).
    function fingerprint(entry as Entry) as ByteArray {
        deviceKeys();
        var input = Bytes.fromUtf8("OTPW-FP-1");
        input.add(entry.kind);
        input.add(entry.type);
        var material = entry.kind == Entry.KIND_OTP ? entry.secret : entry.payload;
        if (material != null) {
            input.addAll(material);
        }
        var mac = Hmac.nativeSha256(_fingerprintKey as ByteArray, input);
        Bytes.wipe(input);
        var out = mac.slice(0, FINGERPRINT_BYTES);
        Bytes.wipe(mac);
        return out;
    }

    // A keyed hash (16 hex characters) of a settings form's content, for
    // the applied-forms log (README 10.3). Keyed like the fingerprints, with
    // its own label, because the content includes secrets. The parts are
    // fed to the native HMAC one by one, each preceded by `separator`, so a
    // large import form costs no byte loop here (the watchdog).
    function formHash(parts as Array<ByteArray>, separator as Number) as String {
        deviceKeys();
        var mac = new Cryptography.HashBasedMessageAuthenticationCode({
            :algorithm => Cryptography.HASH_SHA256,
            :key => _fingerprintKey as ByteArray
        });
        mac.update(Bytes.fromUtf8("OTPW-FH-1"));
        for (var i = 0; i < parts.size(); i++) {
            mac.update([separator]b);
            if (parts[i].size() > 0) {
                mac.update(parts[i]);
            }
        }
        var digest = mac.digest();
        var out = Bytes.toHex(digest.slice(0, FINGERPRINT_BYTES));
        Bytes.wipe(digest);
        return out;
    }

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

    // True when the stored PIN predates the 4-digit rule (no marker), so it
    // may have 4 to 10 digits and should be replaced by a 4-digit one.
    function isLegacyPin() as Boolean {
        var marker = _store.getValue(PIN_LENGTH_KEY);
        return hasPin() && !(marker instanceof Number && marker == PIN_DIGITS);
    }

    // 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: any PIN that can unlock, legacy
    // lengths included.
    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;
    }

    // True for exactly PIN_DIGITS ASCII digits: what a new PIN must be.
    static function isValidNewPin(pin as ByteArray) as Boolean {
        return pin.size() == PIN_DIGITS && isValidPin(pin);
    }

    // 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. Throws VaultError
    // INVALID unless the PIN has exactly PIN_DIGITS digits.
    function newPinJob(pin as ByteArray, iterations as Number) as PinJob {
        if (!isValidNewPin(pin)) {
            throw new VaultError(VaultError.INVALID, "PIN: " + PIN_DIGITS + " digits");
        }
        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, and mark the PIN as a 4-digit one.
    // 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);
        // After the wrap: a crash in between leaves the PIN unmarked, which
        // only means being asked for a new PIN once more.
        _store.setValue(PIN_LENGTH_KEY, PIN_DIGITS);
    }

    // Forget the PIN (reset or erase). The master key is then gone for good.
    function deleteWrap() as Void {
        _store.deleteValue(WRAP_KEY);
        _store.deleteValue(PIN_LENGTH_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: " + KeyManager.MIN_PIN_DIGITS + " to " + KeyManager.MAX_PIN_DIGITS + " 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();
    }
}