// Author:  cirakg & claude ai
// Summary: Authenticated encryption of records: AES-256-CBC plus HMAC-SHA256, encrypt-then-MAC (README 7.3).
//
// Format of a sealed box (one ByteArray):
//     version (1 byte, = 1) | IV (16 bytes) | AES-256-CBC ciphertext | tag (32 bytes)
//     tag = HMAC-SHA256(macKey, "OTPW-REC-1" || len(aad) as 4 bytes || aad || version || IV || ciphertext)
//
// * Two keys come from one master key with HKDF (labels "OTPW enc v1" and
//   "OTPW mac v1"), so encryption and authentication never share a key.
// * aad ("additional authenticated data") is authenticated but not encrypted;
//   storage puts the record ID and header there (M3), so a record can't be
//   swapped for another one or have its header changed unnoticed. Its length
//   is part of the MAC input, so aad and ciphertext can't be shifted against
//   each other.
// * open() checks the tag in constant time BEFORE decrypting, and returns
//   null for any problem (wrong key, tampering, truncation, unknown version)
//   without saying which.
// * Every seal uses a fresh random IV from Cryptography.randomBytes.

import Toybox.Cryptography;
import Toybox.Lang;

module SealedBox {

    const VERSION = 1;
    const IV_BYTES = 16;
    const TAG_BYTES = 32;
    const KEY_BYTES = 32;
    const MIN_SIZE = 1 + IV_BYTES + 16 + TAG_BYTES; // ciphertext is at least one block

    // The encryption and MAC keys derived from one master key.
    class Keys {
        var encKey as ByteArray;
        var macKey as ByteArray;

        function initialize(masterKey as ByteArray) {
            encKey = Kdf.hkdf(masterKey, []b, Bytes.fromUtf8("OTPW enc v1"), KEY_BYTES);
            macKey = Kdf.hkdf(masterKey, []b, Bytes.fromUtf8("OTPW mac v1"), KEY_BYTES);
        }

        // Wipe both keys (when the session locks).
        function wipe() as Void {
            Bytes.wipe(encKey);
            Bytes.wipe(macKey);
        }
    }

    // Seal plaintext with a fresh random IV.
    function seal(keys as Keys, aad as ByteArray, plaintext as ByteArray) as ByteArray {
        return sealWithIv(keys, aad, plaintext, Cryptography.randomBytes(IV_BYTES));
    }

    // Seal with a given IV. Only tests call this directly (fixed IVs make the
    // output reproducible); reusing an IV in real use would be a mistake.
    function sealWithIv(keys as Keys, aad as ByteArray, plaintext as ByteArray, iv as ByteArray) as ByteArray {
        var ciphertext = Aes.cbcEncrypt(keys.encKey, iv, plaintext);
        var box = [VERSION]b;
        box.addAll(iv);
        box.addAll(ciphertext);
        box.addAll(tag(keys, aad, box));
        return box;
    }

    // The plaintext, or null if the box is malformed, was sealed with another
    // key or aad, or was changed in any way.
    function open(keys as Keys, aad as ByteArray, box as ByteArray) as ByteArray? {
        var size = box.size();
        if (size < MIN_SIZE || (size - 1 - IV_BYTES - TAG_BYTES) % Aes.BLOCK_BYTES != 0 || box[0] != VERSION) {
            return null;
        }
        var body = box.slice(0, size - TAG_BYTES);
        var expected = tag(keys, aad, body);
        var valid = Bytes.equal(expected, box.slice(size - TAG_BYTES, size));
        Bytes.wipe(expected);
        if (!valid) {
            return null;
        }
        return Aes.cbcDecrypt(keys.encKey, box.slice(1, 1 + IV_BYTES), box.slice(1 + IV_BYTES, size - TAG_BYTES));
    }

    // HMAC over the label, the aad (with its length) and the box body.
    function tag(keys as Keys, aad as ByteArray, body as ByteArray) as ByteArray {
        var input = Bytes.fromUtf8("OTPW-REC-1");
        input.addAll(Bytes.uint32BigEndian(aad.size()));
        input.addAll(aad);
        input.addAll(body);
        return Hmac.nativeSha256(keys.macKey, input);
    }
}