// Author:  cirakg & claude ai
// Summary: AES-CBC encryption with PKCS#7 padding on top of the native Connect IQ cipher.
//
// The native Cipher offers AES-128 and AES-256 in ECB and CBC mode, without
// padding (CBC input must be a multiple of 16 bytes). This module adds PKCS#7
// padding. AES-256 is used for records (decided in M2: it works in the
// simulator on all four test devices and costs nothing extra; the PIN, not
// the cipher, limits the security of locked entries).
//
// Security note: CBC alone does not detect tampering. Records are always
// sealed with an HMAC as well (SealedBox.mc), and the MAC is checked before
// anything is decrypted, so padding errors never reach an attacker.
// Verified with NIST SP 800-38A F.2.1 and F.2.5 (test/crypto/AesTest.mc).

import Toybox.Cryptography;
import Toybox.Lang;

module Aes {

    const BLOCK_BYTES = 16;

    // A CBC cipher for a 16- or 32-byte key and a 16-byte IV.
    function cipher(key as ByteArray, iv as ByteArray) as Cryptography.Cipher {
        if ((key.size() != 16 && key.size() != 32) || iv.size() != BLOCK_BYTES) {
            throw new Lang.InvalidValueException("AES: key must be 16 or 32 bytes, IV 16 bytes");
        }
        return new Cryptography.Cipher({
            :algorithm => key.size() == 32 ? Cryptography.CIPHER_AES256 : Cryptography.CIPHER_AES128,
            :mode => Cryptography.MODE_CBC,
            :key => key,
            :iv => iv
        });
    }

    // CBC without padding; data must be a multiple of 16 bytes.
    function cbcRaw(key as ByteArray, iv as ByteArray, data as ByteArray, encrypt as Boolean) as ByteArray {
        var c = cipher(key, iv);
        return encrypt ? c.encrypt(data) : c.decrypt(data);
    }

    // Encrypt with PKCS#7 padding: 1 to 16 bytes, each holding the pad length.
    function cbcEncrypt(key as ByteArray, iv as ByteArray, plaintext as ByteArray) as ByteArray {
        var padLength = BLOCK_BYTES - (plaintext.size() % BLOCK_BYTES);
        var padded = Bytes.copy(plaintext);
        for (var i = 0; i < padLength; i++) {
            padded.add(padLength);
        }
        var out = cbcRaw(key, iv, padded, true);
        Bytes.wipe(padded);
        return out;
    }

    // Decrypt and remove PKCS#7 padding; null when the length or the padding is
    // wrong.
    function cbcDecrypt(key as ByteArray, iv as ByteArray, ciphertext as ByteArray) as ByteArray? {
        var size = ciphertext.size();
        if (size == 0 || size % BLOCK_BYTES != 0) {
            return null;
        }
        var padded = cbcRaw(key, iv, ciphertext, false);
        var padLength = padded[size - 1];
        var valid = padLength >= 1 && padLength <= BLOCK_BYTES;
        if (valid) {
            for (var i = size - padLength; i < size; i++) {
                valid = valid && padded[i] == padLength;
            }
        }
        var plaintext = valid ? padded.slice(0, size - padLength) : null;
        Bytes.wipe(padded);
        return plaintext;
    }
}