// Author:  cirakg & claude ai
// Summary: PBKDF2, HKDF and AES-CBC tests (RFC 7914, RFC 5869, NIST SP 800-38A) plus padding and job edge cases.

import Toybox.Lang;
import Toybox.Test;

(:test)
function pbkdf2MatchesVectors(logger as Test.Logger) as Boolean {
    var ok = true;
    for (var i = 0; i < TestVectors.PBKDF2.size(); i++) {
        var r = TestSupport.row(TestVectors.PBKDF2, i);
        var key = Kdf.pbkdf2(TestSupport.bytes(r, 0), TestSupport.bytes(r, 1), TestSupport.num(r, 2), TestSupport.num(r, 3));
        ok = TestSupport.sameBytes(logger, "PBKDF2 vector " + i, key, TestSupport.bytes(r, 4)) && ok;
    }
    return ok;
}

// The result must not depend on how the work is split into steps.
(:test)
function pbkdf2StepSizeDoesNotMatter(logger as Test.Logger) as Boolean {
    var password = Bytes.fromUtf8("123456");
    var salt = new [16]b;
    var oneByOne = new Kdf.Pbkdf2(password, salt, 37, 40);
    var steps = 0;
    while (!oneByOne.step(1)) {
        steps++;
    }
    // 2 output blocks x 37 iterations, the last step returns true.
    return Bytes.equal(oneByOne.result(), Kdf.pbkdf2(password, salt, 37, 40)) && steps == 73;
}

(:test)
function pbkdf2RejectsBadParameters(logger as Test.Logger) as Boolean {
    var rejected = 0;
    var cases = [[0, 32], [1, 0], [1, 255 * 32 + 1]];
    for (var i = 0; i < cases.size(); i++) {
        var c = cases[i] as Array<Number>;
        try {
            new Kdf.Pbkdf2([1]b, [2]b, c[0], c[1]);
        } catch (e instanceof Lang.InvalidValueException) {
            rejected++;
        }
    }
    return rejected == cases.size();
}

(:test)
function pbkdf2CalibrationStaysInBounds(logger as Test.Logger) as Boolean {
    var iterations = Kdf.calibrate(1000, 500);
    logger.debug("calibrated iterations for 1 s in this simulator: " + iterations);
    return iterations >= Kdf.MIN_ITERATIONS && iterations <= Kdf.MAX_ITERATIONS;
}

(:test)
function hkdfMatchesRfcVectors(logger as Test.Logger) as Boolean {
    var ok = true;
    for (var i = 0; i < TestVectors.HKDF.size(); i++) {
        var r = TestSupport.row(TestVectors.HKDF, i);
        var okm = Kdf.hkdf(TestSupport.bytes(r, 0), TestSupport.bytes(r, 1), TestSupport.bytes(r, 2), TestSupport.num(r, 3));
        ok = TestSupport.sameBytes(logger, "HKDF vector " + i, okm, TestSupport.bytes(r, 4)) && ok;
    }
    return ok;
}

(:test)
function aesCbcMatchesNistVectors(logger as Test.Logger) as Boolean {
    var ok = true;
    for (var i = 0; i < TestVectors.AES_CBC.size(); i++) {
        var r = TestSupport.row(TestVectors.AES_CBC, i);
        var key = TestSupport.bytes(r, 0);
        var iv = TestSupport.bytes(r, 1);
        var plain = TestSupport.bytes(r, 2);
        var cipher = TestSupport.bytes(r, 3);
        ok = TestSupport.sameBytes(logger, "AES encrypt " + i, Aes.cbcRaw(key, iv, plain, true), cipher) && ok;
        ok = TestSupport.sameBytes(logger, "AES decrypt " + i, Aes.cbcRaw(key, iv, cipher, false), plain) && ok;
    }
    return ok;
}

// PKCS#7 round trips for every length from 0 to 40 bytes, both key sizes.
(:test)
function aesPaddingRoundTrips(logger as Test.Logger) as Boolean {
    var iv = new [16]b;
    var keys = [new [16]b, new [32]b];
    for (var k = 0; k < keys.size(); k++) {
        for (var length = 0; length <= 40; length++) {
            var plain = new [length]b;
            for (var i = 0; i < length; i++) {
                plain[i] = (i * 7 + length) & 0xFF;
            }
            var cipher = Aes.cbcEncrypt(keys[k] as ByteArray, iv, plain);
            var back = Aes.cbcDecrypt(keys[k] as ByteArray, iv, cipher);
            if (cipher.size() != (length / 16 + 1) * 16 || back == null || !Bytes.equal(back, plain)) {
                logger.error("padding round trip failed at length " + length);
                return false;
            }
        }
    }
    return true;
}

(:test)
function aesRejectsBadPaddingAndLength(logger as Test.Logger) as Boolean {
    var key = new [32]b;
    var iv = new [16]b;
    // A block whose decryption ends in 0x00, which is never valid padding.
    var badPadding = Aes.cbcRaw(key, iv, new [16]b, true);
    var badKey = false;
    try {
        Aes.cbcEncrypt(new [24]b, iv, [1]b);
    } catch (e instanceof Lang.InvalidValueException) {
        badKey = true;
    }
    return Aes.cbcDecrypt(key, iv, badPadding) == null && Aes.cbcDecrypt(key, iv, new [15]b) == null
        && Aes.cbcDecrypt(key, iv, []b) == null && badKey;
}