#!/usr/bin/env qore
# -*- mode: qore; indent-tabs-mode: nil -*-

%new-style
%enable-all-warnings
%require-types
%strict-args

%requires ../../../../qlib/Util.qm
%requires ../../../../qlib/QUnit.qm

%exec-class CryptoTest

public class CryptoTest inherits QUnit::Test {
    public {
        const Str = "Hello There This is a Test - 1234567890";
        const Key = "1234567812345abcabcdefghijklmnopqrstuvwxyz";
        const Aad = "additional authenticated data";
    }

    constructor() : Test("Crypto test", "1.0") {
        addTestCase("Test crypto functions", \test(), NOTHING);

        # Return for compatibility with test harness that checks return value.
        set_return_value(main());
    }

    test() {
        assertEq(Type::Int, call_builtin_function("type", 1));

        foreach string str in (Str, strmul("x", 99999)) {
            testEncryption("des", str, \des_encrypt_cbc(), \des_decrypt_cbc_to_string());
            testEncryption("des-cfb", str);
            testEncryption("des-ofb", str);
            testEncryption("desede", str, \des_ede_encrypt_cbc(), \des_ede_decrypt_cbc_to_string());
            testEncryption("desede-cfb", str);
            testEncryption("desede-ofb", str);
            testEncryption("desede3", str, \des_ede3_encrypt_cbc(), \des_ede3_decrypt_cbc_to_string());
            testEncryption("desede3-cfb", str);
            testEncryption("desede3-ofb", str);
            testEncryption("desx", str, \desx_encrypt_cbc(), \desx_decrypt_cbc_to_string());
            testEncryption("blowfish", str, \blowfish_encrypt_cbc(), \blowfish_decrypt_cbc_to_string());
            testEncryption("blowfish-cfb", str);
            testEncryption("blowfish-ofb", str);
            testEncryption("rc2", str, \rc2_encrypt_cbc(), \rc2_decrypt_cbc_to_string());
            testEncryption("rc2-cfb", str);
            testEncryption("rc2-ofb", str);
%ifdef HAVE_RC4
            testEncryption("rc4", str, \rc4_encrypt_cbc(), \rc4_decrypt_cbc_to_string());
%endif
            testEncryption("cast5", str, \cast5_encrypt_cbc(), \cast5_decrypt_cbc_to_string());
            testEncryption("cast5-cfb", str);
            testEncryption("cast5-ofb", str);

            binary bkey = des_random_key();
            binary x = des_encrypt_cbc(str, bkey);
            string xstr = des_decrypt_cbc_to_string(x, bkey);
            assertEq(str, xstr, "DES random single key encrypt-decrypt");
        }

        # test AES ciphers with MACs and AAD
        foreach string str in (Str, strmul("x", 99999)) {
            map testGcmEncryption($1, str), ("aes128", "aes192", "aes256");
        }
    }

    testEncryption(string alg, string str = Str, *code enc_func, *code dec_func) {
        binary x1 = encrypt(alg, str, Key);

        on_error printf("alg: %y\n", alg);
        hash<auto> h = get_cipher_info(alg);
        assertEq(Type::Int, h.key_len.type());

        if (enc_func) {
            binary x = enc_func(str, Key);
            string xstr = dec_func(x, Key);
            assertEq(str, xstr, alg + "encrypt+decrypt");
            assertEq(x, x1);
        }

        assertEq(str, binary_to_string(decrypt_to_binary(alg, x1, Key)), alg + " encrypt bin");
        assertEq(str, decrypt_to_string(alg, x1, Key), alg + " encrypt str");

        # decryption must fail with a different key
        try {
            binary x = decrypt_to_binary(alg, x1, "x" + Key);
            assertEq(False, x == x1, alg + " decrypt");
        } catch (hash<ExceptionInfo> ex) {
            assertEq("DECRYPT-ERROR", ex.err, alg + " decrypt");
        }

        # encrypt with an iv
        binary iv = get_random_bytes(12);
        x1 = encrypt(alg, str, Key, iv);

        assertEq(str, binary_to_string(decrypt_to_binary(alg, x1, Key, iv)), alg + " encrypt bin+iv");
        assertEq(str, decrypt_to_string(alg, x1, Key, iv), alg + " encrypt str+iv");

        # decryption must fail with a different key
        try {
            binary x = decrypt_to_binary(alg, x1, Key);
            assertEq(False, x == x1, alg + " decrypt+iv");
        } catch (hash<ExceptionInfo> ex) {
            assertEq("DECRYPT-ERROR", ex.err, alg + " decrypt+iv");
        }
    }

    testGcmEncryption(string alg, string str = Str) {
        binary mac;

        on_error printf("alg: %y\n", alg);
        hash<auto> h = get_cipher_info(alg);
        assertEq("GCM", h.type);

        binary x1 = encrypt(alg, str, Key, NOTHING, \mac, 16);
        assertEq(str, binary_to_string(decrypt_to_binary(alg, x1, Key, NOTHING, mac)), alg + " encrypt bin");
        assertEq(str, decrypt_to_string(alg, x1, Key, NOTHING, mac), alg + " encrypt str");

        # decryption must fail with a different key
        assertThrows("DECRYPT-ERROR", \decrypt_to_binary(), (alg, x1, "x" + Key, NOTHING, mac), alg + " encrypt bin");
        assertThrows("DECRYPT-ERROR", \decrypt_to_string(), (alg, x1, "x" + Key, NOTHING, mac), alg + " encrypt bin");

        # decryption must fail with a different MAC
        mac = get_random_bytes(16);
        assertThrows("DECRYPT-ERROR", \decrypt_to_binary(), (alg, x1, "x" + Key, NOTHING, mac), alg + " encrypt bin");
        assertThrows("DECRYPT-ERROR", \decrypt_to_string(), (alg, x1, "x" + Key, NOTHING, mac), alg + " encrypt bin");

        # test with IVs
        binary iv = get_random_bytes(12);
        x1 = encrypt(alg, str, Key, iv, \mac, 16);
        assertEq(str, binary_to_string(decrypt_to_binary(alg, x1, Key, iv, mac)), alg + " encrypt bin");
        assertEq(str, decrypt_to_string(alg, x1, Key, iv, mac), alg + " encrypt str");

        # decryption must fail with a different IV
        assertThrows("DECRYPT-ERROR", \decrypt_to_binary(), (alg, x1, Key, NOTHING, mac), alg + " encrypt bin");
        assertThrows("DECRYPT-ERROR", \decrypt_to_string(), (alg, x1, Key, NOTHING, mac), alg + " encrypt bin");

        # test with AAD
        x1 = encrypt(alg, str, Key, NOTHING, \mac, 16, Aad);
        assertEq(str, binary_to_string(decrypt_to_binary(alg, x1, Key, NOTHING, mac, Aad)), alg + " encrypt bin");
        assertEq(str, decrypt_to_string(alg, x1, Key, NOTHING, mac, Aad), alg + " encrypt str");

        # decryption must fail with different AAD
        assertThrows("DECRYPT-ERROR", \decrypt_to_binary(), (alg, x1, Key, NOTHING, mac), alg + " encrypt bin");
        assertThrows("DECRYPT-ERROR", \decrypt_to_string(), (alg, x1, Key, NOTHING, mac), alg + " encrypt bin");
    }
}
