// trunk-ignore-all(gitleaks): These are dummy values. Not real secrets.
#include "CryptoEngine.h"

#include "TestUtil.h"
#include "aes-ccm.h"
#include <XEdDSA.h>
#include <cassert>
#include <unity.h>

void HexToBytes(uint8_t *result, const std::string hex, size_t len = 0)
{
    if (len) {
        memset(result, 0, len);
    }
    for (unsigned int i = 0; i < hex.length(); i += 2) {
        std::string byteString = hex.substr(i, 2);
        result[i / 2] = (uint8_t)strtol(byteString.c_str(), NULL, 16);
    }
    return;
}

void setUp(void)
{
    // set stuff up here
}

void tearDown(void)
{
    // clean stuff up here
}

void test_SHA256(void)
{
    uint8_t expected[32];
    uint8_t hash[32] = {0};

    HexToBytes(expected, "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855");
    crypto->hash(hash, 0);
    TEST_ASSERT_EQUAL_MEMORY(hash, expected, 32);

    HexToBytes(hash, "d3", 32);
    HexToBytes(expected, "28969cdfa74a12c82f3bad960b0b000aca2ac329deea5c2328ebc6f2ba9802c1");
    crypto->hash(hash, 1);
    TEST_ASSERT_EQUAL_MEMORY(hash, expected, 32);

    HexToBytes(hash, "11af", 32);
    HexToBytes(expected, "5ca7133fa735326081558ac312c620eeca9970d1e70a4b95533d956f072d1f98");
    crypto->hash(hash, 2);
    TEST_ASSERT_EQUAL_MEMORY(hash, expected, 32);
}

void test_SHA256_large_input(void)
{
    uint8_t hash[300] = {0};
    uint8_t expected[32];

    HexToBytes(expected, "d13d4a8b3b8add19b5970157f09d00c12cbda4fed4d74d8493156523f7069b66");
    crypto->hash(hash, sizeof(hash));
    TEST_ASSERT_EQUAL_MEMORY(hash, expected, sizeof(expected));
}

void test_ECB_AES256(void)
{
    // https://csrc.nist.gov/CSRC/media/Projects/Cryptographic-Standards-and-Guidelines/documents/examples/AES_ECB.pdf

    uint8_t key[32] = {0};
    uint8_t plain[16] = {0};
    uint8_t result[16] = {0};
    uint8_t expected[16] = {0};

    HexToBytes(key, "603DEB1015CA71BE2B73AEF0857D77811F352C073B6108D72D9810A30914DFF4");

    HexToBytes(plain, "6BC1BEE22E409F96E93D7E117393172A");
    HexToBytes(expected, "F3EED1BDB5D2A03C064B5A7E3DB181F8");
    crypto->aesSetKey(key, 32);
    crypto->aesEncrypt(plain, result); // Does 16 bytes at a time
    TEST_ASSERT_EQUAL_MEMORY(expected, result, 16);

    HexToBytes(plain, "AE2D8A571E03AC9C9EB76FAC45AF8E51");
    HexToBytes(expected, "591CCB10D410ED26DC5BA74A31362870");
    crypto->aesSetKey(key, 32);
    crypto->aesEncrypt(plain, result); // Does 16 bytes at a time
    TEST_ASSERT_EQUAL_MEMORY(expected, result, 16);

    HexToBytes(plain, "30C81C46A35CE411E5FBC1191A0A52EF");
    HexToBytes(expected, "B6ED21B99CA6F4F9F153E7B1BEAFED1D");
    crypto->aesSetKey(key, 32);
    crypto->aesEncrypt(plain, result); // Does 16 bytes at a time
    TEST_ASSERT_EQUAL_MEMORY(expected, result, 16);
}
void test_ECB_AES128(void)
{
    // https://csrc.nist.gov/CSRC/media/Projects/Cryptographic-Standards-and-Guidelines/documents/examples/AES_ECB.pdf
    uint8_t key[16] = {0};
    uint8_t plain[16] = {0};
    uint8_t result[16] = {0};
    uint8_t expected[16] = {0};

    HexToBytes(key, "2B7E151628AED2A6ABF7158809CF4F3C");

    HexToBytes(plain, "6BC1BEE22E409F96E93D7E117393172A");
    HexToBytes(expected, "3AD77BB40D7A3660A89ECAF32466EF97");
    crypto->aesSetKey(key, 16);
    crypto->aesEncrypt(plain, result);
    TEST_ASSERT_EQUAL_MEMORY(expected, result, 16);

    HexToBytes(plain, "AE2D8A571E03AC9C9EB76FAC45AF8E51");
    HexToBytes(expected, "F5D3D58503B9699DE785895A96FDBAAF");
    crypto->aesSetKey(key, 16);
    crypto->aesEncrypt(plain, result);
    TEST_ASSERT_EQUAL_MEMORY(expected, result, 16);
}

void test_DH25519(void)
{
    // test vectors from wycheproof x25519
    // https://github.com/C2SP/wycheproof/blob/master/testvectors/x25519_test.json
    uint8_t private_key[32];
    uint8_t public_key[32];
    uint8_t expected_shared[32];

    HexToBytes(public_key, "504a36999f489cd2fdbc08baff3d88fa00569ba986cba22548ffde80f9806829");
    HexToBytes(private_key, "c8a9d5a91091ad851c668b0736c1c9a02936c0d3ad62670858088047ba057475");
    HexToBytes(expected_shared, "436a2c040cf45fea9b29a0cb81b1f41458f863d0d61b453d0a982720d6d61320");
    crypto->setDHPrivateKey(private_key);
    TEST_ASSERT(crypto->setDHPublicKey(public_key));
    TEST_ASSERT_EQUAL_MEMORY(expected_shared, crypto->shared_key, 32);

    HexToBytes(public_key, "63aa40c6e38346c5caf23a6df0a5e6c80889a08647e551b3563449befcfc9733");
    HexToBytes(private_key, "d85d8c061a50804ac488ad774ac716c3f5ba714b2712e048491379a500211958");
    HexToBytes(expected_shared, "279df67a7c4611db4708a0e8282b195e5ac0ed6f4b2f292c6fbd0acac30d1332");
    crypto->setDHPrivateKey(private_key);
    TEST_ASSERT(crypto->setDHPublicKey(public_key));
    TEST_ASSERT_EQUAL_MEMORY(expected_shared, crypto->shared_key, 32);

    HexToBytes(public_key, "ecffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff7f");
    HexToBytes(private_key, "18630f93598637c35da623a74559cf944374a559114c7937811041fc8605564a");
    crypto->setDHPrivateKey(private_key);
    TEST_ASSERT(!crypto->setDHPublicKey(public_key)); // Weak public key results in 0 shared key

    HexToBytes(public_key, "f7e13a1a067d2f4e1061bf9936fde5be6b0c2494a8f809cbac7f290ef719e91c");
    HexToBytes(private_key, "10300724f3bea134eb1575245ef26ff9b8ccd59849cd98ce1a59002fe1d5986c");
    HexToBytes(expected_shared, "24becd5dfed9e9289ba2e15b82b0d54f8e9aacb72f5e4248c58d8d74b451ce76");
    crypto->setDHPrivateKey(private_key);
    TEST_ASSERT(crypto->setDHPublicKey(public_key));
    crypto->hash(crypto->shared_key, 32);
    TEST_ASSERT_EQUAL_MEMORY(expected_shared, crypto->shared_key, 32);
}

void test_PKC(void)
{
    uint8_t private_key[32];
    meshtastic_NodeInfoLite_public_key_t public_key;
    uint8_t expected_shared[32];
    uint8_t expected_decrypted[32];
    uint8_t radioBytes[128] __attribute__((__aligned__));
    uint8_t decrypted[128] __attribute__((__aligned__));
    uint8_t expected_nonce[16];

    uint32_t fromNode = 0x0929;
    uint64_t packetNum = 0x13b2d662;
    HexToBytes(public_key.bytes, "db18fc50eea47f00251cb784819a3cf5fc361882597f589f0d7ff820e8064457");
    public_key.size = 32;
    HexToBytes(private_key, "a00330633e63522f8a4d81ec6d9d1e6617f6c8ffd3a4c698229537d44e522277");
    HexToBytes(expected_shared, "777b1545c9d6f9a2");
    HexToBytes(expected_decrypted, "08011204746573744800");
    HexToBytes(radioBytes, "8c646d7a2909000062d6b2136b00000040df24abfcc30a17a3d9046726099e796a1c036a792b");
    HexToBytes(expected_nonce, "62d6b213036a792b2909000000");
    crypto->setDHPrivateKey(private_key);

    TEST_ASSERT(crypto->decryptCurve25519(fromNode, public_key, packetNum, 22, radioBytes + 16, decrypted));
    TEST_ASSERT_EQUAL_MEMORY(expected_shared, crypto->shared_key, 8);
    TEST_ASSERT_EQUAL_MEMORY(expected_nonce, crypto->nonce, 13);
    TEST_ASSERT_EQUAL_MEMORY(expected_decrypted, decrypted, 10);

    uint32_t toNode = 0; // Only impacts logging
    uint8_t encrypted[128] __attribute__((__aligned__));
    TEST_ASSERT(crypto->encryptCurve25519(toNode, fromNode, public_key, packetNum, 10, decrypted, encrypted));
    TEST_ASSERT_EQUAL_MEMORY(expected_shared, crypto->shared_key, 8);
    // The extraNonce is random, so skip checking the nonce and encrypted output here

    // Copy the nonce to check it after encryption
    memcpy(expected_nonce, crypto->nonce, 16);

    // Decrypt the re-encrypted bytes and check they are the same as what we expect
    TEST_ASSERT(crypto->decryptCurve25519(fromNode, public_key, packetNum, 22, encrypted, decrypted));
    TEST_ASSERT_EQUAL_MEMORY(expected_shared, crypto->shared_key, 8);
    TEST_ASSERT_EQUAL_MEMORY(expected_nonce, crypto->nonce, 13);
    TEST_ASSERT_EQUAL_MEMORY(expected_decrypted, decrypted, 10);
}

void test_XEdDSA(void)
{
    uint8_t private_key[32];
    uint8_t x_public_key[32];
    uint8_t ed_private_key[32];
    uint8_t ed_public_key[32];
    uint8_t ed_public_key2[32];
    uint8_t message[] = "This is a test!";
    uint8_t message2[] = "This is a test.";
    uint8_t signature[64];
    uint32_t fromNode = 0x1234;
    uint32_t packetId = 0xDEADBEEF;
    uint32_t portnum = 1;
    for (int times = 0; times < 10; times++) {
        printf("Start of time %u\n", times);
        crypto->generateKeyPair(x_public_key, private_key);
        XEdDSA::priv_curve_to_ed_keys(private_key, ed_private_key, ed_public_key);
        crypto->curve_to_ed_pub(x_public_key, ed_public_key2);
        TEST_ASSERT_EQUAL_MEMORY(ed_public_key, ed_public_key2, 32);

        // Sign and verify with metadata
        TEST_ASSERT(crypto->xeddsa_sign(fromNode, packetId, portnum, message, sizeof(message), signature));
        TEST_ASSERT(crypto->xeddsa_verify(x_public_key, fromNode, packetId, portnum, message, sizeof(message), signature));

        // Different payload fails
        TEST_ASSERT_FALSE(
            crypto->xeddsa_verify(x_public_key, fromNode, packetId, portnum, message2, sizeof(message2), signature));

        // Different fromNode fails
        TEST_ASSERT_FALSE(
            crypto->xeddsa_verify(x_public_key, fromNode + 1, packetId, portnum, message, sizeof(message), signature));

        // Different packetId fails
        TEST_ASSERT_FALSE(
            crypto->xeddsa_verify(x_public_key, fromNode, packetId + 1, portnum, message, sizeof(message), signature));

        // Different portnum fails
        TEST_ASSERT_FALSE(
            crypto->xeddsa_verify(x_public_key, fromNode, packetId, portnum + 1, message, sizeof(message), signature));
    }
}

// A signature only verifies under the signer's own key; a different key (or an all-zero key) fails.
void test_XEdDSA_cross_key_reject(void)
{
    uint8_t pubA[32], privA[32];
    uint8_t pubB[32], privB[32];
    uint8_t signature[64];
    uint8_t message[] = "cross-key check";
    uint32_t fromNode = 0x4242, packetId = 0xABCD1234, portnum = 7;

    crypto->generateKeyPair(pubA, privA); // engine now holds key A
    TEST_ASSERT(crypto->xeddsa_sign(fromNode, packetId, portnum, message, sizeof(message), signature));

    crypto->generateKeyPair(pubB, privB); // unrelated key pair

    TEST_ASSERT_TRUE(crypto->xeddsa_verify(pubA, fromNode, packetId, portnum, message, sizeof(message), signature));
    TEST_ASSERT_FALSE(crypto->xeddsa_verify(pubB, fromNode, packetId, portnum, message, sizeof(message), signature));

    uint8_t zeroKey[32] = {0};
    TEST_ASSERT_FALSE(crypto->xeddsa_verify(zeroKey, fromNode, packetId, portnum, message, sizeof(message), signature));
}

// Signing with an unset (all-zero) private key must fail rather than emit a bogus signature.
void test_XEdDSA_empty_key_sign_fails(void)
{
    CryptoEngine fresh; // freshly constructed: xeddsa_private_key is all zero
    uint8_t signature[64];
    uint8_t message[] = "no key";
    TEST_ASSERT_FALSE(fresh.xeddsa_sign(0x1, 0x2, 0x3, message, sizeof(message), signature));
}

// curve_to_ed_pub caches the last converted key; verifying A, then B, then A must stay correct.
void test_XEdDSA_curve_to_ed_cache(void)
{
    uint8_t pubA[32], privA[32], sigA[64];
    uint8_t pubB[32], privB[32], sigB[64];
    uint8_t message[] = "cache check";
    uint32_t fromNode = 0x11, packetId = 0x22, portnum = 3;

    crypto->generateKeyPair(pubA, privA);
    TEST_ASSERT(crypto->xeddsa_sign(fromNode, packetId, portnum, message, sizeof(message), sigA));
    crypto->generateKeyPair(pubB, privB);
    TEST_ASSERT(crypto->xeddsa_sign(fromNode, packetId, portnum, message, sizeof(message), sigB));

    // Interleave keys to exercise both cache hits and cache invalidation.
    TEST_ASSERT_TRUE(crypto->xeddsa_verify(pubA, fromNode, packetId, portnum, message, sizeof(message), sigA));
    TEST_ASSERT_TRUE(crypto->xeddsa_verify(pubB, fromNode, packetId, portnum, message, sizeof(message), sigB));
    TEST_ASSERT_TRUE(crypto->xeddsa_verify(pubA, fromNode, packetId, portnum, message, sizeof(message), sigA));
    TEST_ASSERT_FALSE(crypto->xeddsa_verify(pubA, fromNode, packetId, portnum, message, sizeof(message), sigB));
}

// A payload at the maximum signable size (DATA_PAYLOAD_LEN - signature) round-trips and detects tampering.
void test_XEdDSA_max_payload(void)
{
    const size_t len = meshtastic_Constants_DATA_PAYLOAD_LEN - XEDDSA_SIGNATURE_SIZE;
    uint8_t payload[meshtastic_Constants_DATA_PAYLOAD_LEN];
    for (size_t i = 0; i < len; i++)
        payload[i] = (uint8_t)(i * 7 + 1);

    uint8_t pub[32], priv[32], signature[64];
    crypto->generateKeyPair(pub, priv);
    uint32_t fromNode = 0xFEED, packetId = 0xC0DE, portnum = 1;

    TEST_ASSERT(crypto->xeddsa_sign(fromNode, packetId, portnum, payload, len, signature));
    TEST_ASSERT(crypto->xeddsa_verify(pub, fromNode, packetId, portnum, payload, len, signature));
    payload[0] ^= 0x01;
    TEST_ASSERT_FALSE(crypto->xeddsa_verify(pub, fromNode, packetId, portnum, payload, len, signature));
}

// XEdDSA is a randomized (hedged) scheme: the nonce mixes in Z, caller-supplied randomness
// (Signal spec; meshtastic/Crypto#3). CryptoEngine::xeddsa_sign seeds Z from the hardware RNG, so
// signing the same message twice yields *different* signatures that both verify. This pins that
// the randomization is actually wired through end to end - if signing regresses to deterministic
// (Z dropped by the library, or xeddsa_sign stops seeding entropy), the inequality assertion fails.
void test_XEdDSA_repeated_sign_is_randomized(void)
{
    uint8_t pub[32], priv[32], sig1[64], sig2[64];
    uint8_t message[] = "same message";
    uint32_t fromNode = 0x9, packetId = 0x9, portnum = 9;

    crypto->generateKeyPair(pub, priv);
    TEST_ASSERT(crypto->xeddsa_sign(fromNode, packetId, portnum, message, sizeof(message), sig1));
    TEST_ASSERT(crypto->xeddsa_sign(fromNode, packetId, portnum, message, sizeof(message), sig2));

    TEST_ASSERT_TRUE_MESSAGE(memcmp(sig1, sig2, sizeof(sig1)) != 0,
                             "signatures must differ - XEdDSA Z randomization is not wired through");
    TEST_ASSERT_TRUE(crypto->xeddsa_verify(pub, fromNode, packetId, portnum, message, sizeof(message), sig1));
    TEST_ASSERT_TRUE(crypto->xeddsa_verify(pub, fromNode, packetId, portnum, message, sizeof(message), sig2));
}

void test_AES_CTR(void)
{
    uint8_t expected[32];
    uint8_t plain[32];
    uint8_t nonce[32];
    CryptoKey k;

    // vectors from https://www.rfc-editor.org/rfc/rfc3686#section-6
    k.length = 32;
    HexToBytes(k.bytes, "776BEFF2851DB06F4C8A0542C8696F6C6A81AF1EEC96B4D37FC1D689E6C1C104");
    HexToBytes(nonce, "00000060DB5672C97AA8F0B200000001");
    HexToBytes(expected, "145AD01DBF824EC7560863DC71E3E0C0");
    memcpy(plain, "Single block msg", 16);

    crypto->encryptAESCtr(k, nonce, 16, plain);
    TEST_ASSERT_EQUAL_MEMORY(expected, plain, 16);

    k.length = 16;
    memcpy(plain, "Single block msg", 16);
    HexToBytes(k.bytes, "AE6852F8121067CC4BF7A5765577F39E");
    HexToBytes(nonce, "00000030000000000000000000000001");
    HexToBytes(expected, "E4095D4FB7A7B3792D6175A3261311B8");
    crypto->encryptAESCtr(k, nonce, 16, plain);
    TEST_ASSERT_EQUAL_MEMORY(expected, plain, 16);
}

void test_AES_CCM_partial_block_bounds(void)
{
    // aes_ccm_encr() used to write a whole 16-byte AES block at the output before XOR-ing,
    // so a trailing partial block scribbled up to 15 bytes past what the caller allocated.
    const uint8_t guard = 0xA5;
    const size_t guardLen = 16;
    const size_t lengths[] = {5, 20}; // pure partial block, and one full block plus a partial one
    uint8_t key[32];
    uint8_t nonce[13];
    uint8_t auth[8];

    HexToBytes(key, "603DEB1015CA71BE2B73AEF0857D77811F352C073B6108D72D9810A30914DFF4");
    HexToBytes(nonce, "000102030405060708090A0B0C");

    for (size_t n = 0; n < sizeof(lengths) / sizeof(lengths[0]); n++) {
        const size_t len = lengths[n];
        uint8_t plain[32];
        uint8_t crypt[32 + guardLen];
        uint8_t decrypted[32 + guardLen];

        for (size_t i = 0; i < len; i++)
            plain[i] = (uint8_t)i;
        memset(crypt + len, guard, guardLen);
        memset(decrypted + len, guard, guardLen);

        TEST_ASSERT_EQUAL(0, aes_ccm_ae(key, sizeof(key), nonce, sizeof(auth), plain, len, nullptr, 0, crypt, auth));
        for (size_t i = 0; i < guardLen; i++)
            TEST_ASSERT_EQUAL_UINT8(guard, crypt[len + i]);

        TEST_ASSERT_TRUE(aes_ccm_ad(key, sizeof(key), nonce, sizeof(auth), crypt, len, nullptr, 0, auth, decrypted));
        for (size_t i = 0; i < guardLen; i++)
            TEST_ASSERT_EQUAL_UINT8(guard, decrypted[len + i]);
        TEST_ASSERT_EQUAL_MEMORY(plain, decrypted, len);
    }
}

void test_AES_CCM_rfc3610(void)
{
    // Known-answer vectors from RFC 3610 section 8. They all use L=2, which is what
    // aes_ccm_ae()/aes_ccm_ad() hardcode, and each ends in a partial block.
    struct CcmVector {
        const char *key;
        const char *nonce;
        const char *aad;
        const char *plain;
        const char *crypt;
        const char *tag;
    };
    const CcmVector vectors[] = {
        // Packet Vector #1, M=8
        {"C0C1C2C3C4C5C6C7C8C9CACBCCCDCECF", "00000003020100A0A1A2A3A4A5", "0001020304050607",
         "08090A0B0C0D0E0F101112131415161718191A1B1C1D1E", "588C979A61C663D2F066D0C2C0F989806D5F6B61DAC384", "17E8D12CFDF926E0"},
        // Packet Vector #2, M=8
        {"C0C1C2C3C4C5C6C7C8C9CACBCCCDCECF", "00000004030201A0A1A2A3A4A5", "0001020304050607",
         "08090A0B0C0D0E0F101112131415161718191A1B1C1D1E1F", "72C91A36E135F8CF291CA894085C87E3CC15C439C9E43A3B",
         "A091D56E10400916"},
        // Packet Vector #7, M=10
        {"C0C1C2C3C4C5C6C7C8C9CACBCCCDCECF", "00000009080706A0A1A2A3A4A5", "0001020304050607",
         "08090A0B0C0D0E0F101112131415161718191A1B1C1D1E", "0135D1B2C95F41D5D1D4FEC185D166B8094E999DFED96C",
         "048C56602C97ACBB7490"},
    };

    for (size_t v = 0; v < sizeof(vectors) / sizeof(vectors[0]); v++) {
        const CcmVector &vec = vectors[v];
        const size_t plainLen = strlen(vec.plain) / 2;
        const size_t aadLen = strlen(vec.aad) / 2;
        const size_t tagLen = strlen(vec.tag) / 2;

        uint8_t key[16], nonce[13], aad[8];
        uint8_t plain[32], expectedCrypt[32], expectedTag[16];
        uint8_t crypt[32], tag[16], decrypted[32];

        HexToBytes(key, vec.key);
        HexToBytes(nonce, vec.nonce);
        HexToBytes(aad, vec.aad);
        HexToBytes(plain, vec.plain);
        HexToBytes(expectedCrypt, vec.crypt);
        HexToBytes(expectedTag, vec.tag);

        TEST_ASSERT_EQUAL(0, aes_ccm_ae(key, sizeof(key), nonce, tagLen, plain, plainLen, aad, aadLen, crypt, tag));
        TEST_ASSERT_EQUAL_MEMORY(expectedCrypt, crypt, plainLen);
        TEST_ASSERT_EQUAL_MEMORY(expectedTag, tag, tagLen);

        TEST_ASSERT_TRUE(aes_ccm_ad(key, sizeof(key), nonce, tagLen, crypt, plainLen, aad, aadLen, tag, decrypted));
        TEST_ASSERT_EQUAL_MEMORY(plain, decrypted, plainLen);

        // The AAD is authenticated but not encrypted: corrupting it must fail the tag check
        aad[0] ^= 0x01;
        TEST_ASSERT_FALSE(aes_ccm_ad(key, sizeof(key), nonce, tagLen, crypt, plainLen, aad, aadLen, tag, decrypted));
    }
}

// Helper to create a zero-initialized CryptoKey (matching Channels::getKey() behavior)
static CryptoKey makePsk(const std::string &hex)
{
    CryptoKey k;
    assert(hex.length() / 2 <= sizeof(k.bytes));
    memset(k.bytes, 0, sizeof(k.bytes));
    k.length = hex.length() / 2;
    HexToBytes(k.bytes, hex);
    return k;
}

void test_AES_CCM_AEAD_smoke(void)
{
    // Smoke test - encryption changes the payload and produces a tag
    // (the known-answer coverage lives in test_AES_CCM_rfc3610)
    CryptoKey psk = makePsk("d4f1bb3a20290759f0bcffabcf4e6901");

    uint32_t fromNode = 0x12345678;
    uint32_t toNode = 0x0000AAAA;
    uint64_t packetId = 0xAABBCCDD;

    uint8_t plaintext[10];
    HexToBytes(plaintext, "08011204746573744800");

    uint8_t ciphertextWithTag[10 + CryptoEngine::AEAD_TAG_SIZE];
    memset(ciphertextWithTag, 0, sizeof(ciphertextWithTag));

    TEST_ASSERT_TRUE(crypto->encryptPacketCCM(psk, fromNode, toNode, packetId, 10, plaintext, ciphertextWithTag));

    // Ciphertext should differ from plaintext
    TEST_ASSERT_FALSE(memcmp(plaintext, ciphertextWithTag, 10) == 0);

    // Tag bytes (last 12) should not all be zero
    bool tagAllZero = true;
    for (size_t i = 0; i < CryptoEngine::AEAD_TAG_SIZE; i++) {
        if (ciphertextWithTag[10 + i] != 0) {
            tagAllZero = false;
            break;
        }
    }
    TEST_ASSERT_FALSE(tagAllZero);
}

void test_AES_CCM_AEAD_roundtrip_aes256(void)
{
    // Round-trip encrypt → decrypt → compare (AES-256)
    CryptoKey psk = makePsk("603DEB1015CA71BE2B73AEF0857D77811F352C073B6108D72D9810A30914DFF4");

    uint32_t fromNode = 0xDEADBEEF;
    uint32_t toNode = 0xFFFFFFFF;
    uint64_t packetId = 0x0102030405060708;

    const char *msg = "Hello Meshtastic AEAD!";
    size_t msgLen = strlen(msg);

    uint8_t ciphertextWithTag[64];
    memset(ciphertextWithTag, 0, sizeof(ciphertextWithTag));
    TEST_ASSERT_TRUE(crypto->encryptPacketCCM(psk, fromNode, toNode, packetId, msgLen, (const uint8_t *)msg, ciphertextWithTag));

    uint8_t decrypted[64];
    memset(decrypted, 0, sizeof(decrypted));
    size_t totalBytes = msgLen + CryptoEngine::AEAD_TAG_SIZE;
    TEST_ASSERT_TRUE(crypto->decryptPacketCCM(psk, fromNode, toNode, packetId, totalBytes, ciphertextWithTag, decrypted));

    TEST_ASSERT_EQUAL_MEMORY(msg, decrypted, msgLen);
}

void test_AES_CCM_AEAD_rejects_tampering(void)
{
    // Tampered ciphertext - flip a bit, verify rejection
    {
        CryptoKey psk = makePsk("d4f1bb3a20290759f0bcffabcf4e6901");

        uint32_t fromNode = 0xABCD1234;
        uint32_t toNode = 0x00000001;
        uint64_t packetId = 0x11223344;

        uint8_t plaintext[8] = {0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08};
        uint8_t ciphertextWithTag[8 + CryptoEngine::AEAD_TAG_SIZE];

        TEST_ASSERT_TRUE(crypto->encryptPacketCCM(psk, fromNode, toNode, packetId, 8, plaintext, ciphertextWithTag));

        // Flip a bit in the ciphertext portion
        ciphertextWithTag[3] ^= 0x01;

        uint8_t decrypted[8];
        TEST_ASSERT_FALSE(crypto->decryptPacketCCM(psk, fromNode, toNode, packetId, 8 + CryptoEngine::AEAD_TAG_SIZE,
                                                   ciphertextWithTag, decrypted));
    }

    // Tampered auth tag - modify tag, verify rejection
    {
        CryptoKey psk = makePsk("d4f1bb3a20290759f0bcffabcf4e6901");

        uint32_t fromNode = 0xABCD1234;
        uint32_t toNode = 0x87654321;
        uint64_t packetId = 0x55667788;

        uint8_t plaintext[16] = {0};
        for (int i = 0; i < 16; i++)
            plaintext[i] = (uint8_t)i;

        uint8_t ciphertextWithTag[16 + CryptoEngine::AEAD_TAG_SIZE];
        TEST_ASSERT_TRUE(crypto->encryptPacketCCM(psk, fromNode, toNode, packetId, 16, plaintext, ciphertextWithTag));

        // Corrupt the auth tag (last byte)
        ciphertextWithTag[16 + CryptoEngine::AEAD_TAG_SIZE - 1] ^= 0xFF;

        uint8_t decrypted[16];
        TEST_ASSERT_FALSE(crypto->decryptPacketCCM(psk, fromNode, toNode, packetId, 16 + CryptoEngine::AEAD_TAG_SIZE,
                                                   ciphertextWithTag, decrypted));
    }
}

void test_AES_CCM_AEAD_rejects_undersized(void)
{
    // Packet too small for AEAD - totalBytes <= AEAD_TAG_SIZE
    CryptoKey psk = makePsk("d4f1bb3a20290759f0bcffabcf4e6901");

    uint8_t dummy[CryptoEngine::AEAD_TAG_SIZE] = {0};
    // Sized for the whole input so a regressed length guard fails the assertion below
    // instead of corrupting the stack on its way out.
    uint8_t out[CryptoEngine::AEAD_TAG_SIZE];

    TEST_ASSERT_FALSE(crypto->decryptPacketCCM(psk, 0x1234, 0x4321, 0x5678, CryptoEngine::AEAD_TAG_SIZE, dummy, out));
    TEST_ASSERT_FALSE(crypto->decryptPacketCCM(psk, 0x1234, 0x4321, 0x5678, 0, dummy, out));
}

void test_AES_CCM_AEAD_rejects_wrong_psk(void)
{
    // Wrong PSK - decrypt with different key, verify rejection
    CryptoKey pskA = makePsk("d4f1bb3a20290759f0bcffabcf4e6901");
    CryptoKey pskB = makePsk("00112233445566778899aabbccddeeff");

    uint32_t fromNode = 0x99887766;
    uint32_t toNode = 0x13579BDF;
    uint64_t packetId = 0xDEADFACE;

    uint8_t plaintext[12] = "Hello World";
    uint8_t ciphertextWithTag[12 + CryptoEngine::AEAD_TAG_SIZE];

    TEST_ASSERT_TRUE(crypto->encryptPacketCCM(pskA, fromNode, toNode, packetId, 12, plaintext, ciphertextWithTag));

    // Attempt decryption with wrong key
    uint8_t decrypted[12];
    TEST_ASSERT_FALSE(crypto->decryptPacketCCM(pskB, fromNode, toNode, packetId, 12 + CryptoEngine::AEAD_TAG_SIZE,
                                               ciphertextWithTag, decrypted));
}

void test_AES_CCM_AEAD_roundtrip_aes128(void)
{
    // Round-trip with AES-128 PSK (16-byte key, true AES-128-CCM)
    CryptoKey psk = makePsk("d4f1bb3a20290759f0bcffabcf4e6901");

    uint32_t fromNode = 0x42424242;
    uint32_t toNode = 0x2468ACE0;
    uint64_t packetId = 0xBEEF1234;

    uint8_t plaintext[20] = "AES128 round trip!";
    uint8_t ciphertextWithTag[20 + CryptoEngine::AEAD_TAG_SIZE];

    TEST_ASSERT_TRUE(crypto->encryptPacketCCM(psk, fromNode, toNode, packetId, 20, plaintext, ciphertextWithTag));

    uint8_t decrypted[20];
    TEST_ASSERT_TRUE(crypto->decryptPacketCCM(psk, fromNode, toNode, packetId, 20 + CryptoEngine::AEAD_TAG_SIZE,
                                              ciphertextWithTag, decrypted));
    TEST_ASSERT_EQUAL_MEMORY(plaintext, decrypted, 20);
}

void test_AES_CCM_AEAD_tamper_sweep(void)
{
    // AES-256-CCM round-trip + per-byte tamper detection
    CryptoKey psk = makePsk("603DEB1015CA71BE2B73AEF0857D77811F352C073B6108D72D9810A30914DFF4");

    uint32_t fromNode = 0x01020304;
    uint32_t toNode = 0x0BADCAFE;
    uint64_t packetId = 0x0A0B0C0D0E0F1011;

    uint8_t plaintext[32];
    for (int i = 0; i < 32; i++)
        plaintext[i] = (uint8_t)(i * 7 + 3);

    uint8_t ciphertextWithTag[32 + CryptoEngine::AEAD_TAG_SIZE];
    TEST_ASSERT_TRUE(crypto->encryptPacketCCM(psk, fromNode, toNode, packetId, 32, plaintext, ciphertextWithTag));

    // Valid decrypt
    uint8_t decrypted[32];
    TEST_ASSERT_TRUE(crypto->decryptPacketCCM(psk, fromNode, toNode, packetId, 32 + CryptoEngine::AEAD_TAG_SIZE,
                                              ciphertextWithTag, decrypted));
    TEST_ASSERT_EQUAL_MEMORY(plaintext, decrypted, 32);

    // Flip a bit in every byte in turn, tag included, and verify each one is rejected
    for (size_t i = 0; i < 32 + CryptoEngine::AEAD_TAG_SIZE; i++) {
        uint8_t tampered[32 + CryptoEngine::AEAD_TAG_SIZE];
        memcpy(tampered, ciphertextWithTag, sizeof(tampered));
        tampered[i] ^= 0x80;
        TEST_ASSERT_FALSE(
            crypto->decryptPacketCCM(psk, fromNode, toNode, packetId, 32 + CryptoEngine::AEAD_TAG_SIZE, tampered, decrypted));
    }
}

void test_AES_CCM_AEAD_is_deterministic(void)
{
    // Deterministic - same inputs produce same output
    CryptoKey psk = makePsk("d4f1bb3a20290759f0bcffabcf4e6901");

    uint32_t fromNode = 0xCAFEBABE;
    uint32_t toNode = 0x5A5A5A5A;
    uint64_t packetId = 0xFEEDFACE;

    uint8_t plaintext[5] = {0xDE, 0xAD, 0xBE, 0xEF, 0x42};
    uint8_t ct1[5 + CryptoEngine::AEAD_TAG_SIZE];
    uint8_t ct2[5 + CryptoEngine::AEAD_TAG_SIZE];

    TEST_ASSERT_TRUE(crypto->encryptPacketCCM(psk, fromNode, toNode, packetId, 5, plaintext, ct1));
    TEST_ASSERT_TRUE(crypto->encryptPacketCCM(psk, fromNode, toNode, packetId, 5, plaintext, ct2));

    TEST_ASSERT_EQUAL_MEMORY(ct1, ct2, 5 + CryptoEngine::AEAD_TAG_SIZE);
}

void test_AES_CCM_AEAD_binds_nonce_inputs(void)
{
    // Wrong nonce input - the nonce derives from both fromNode and packetId,
    // so each one on its own must be enough to make the tag check fail
    CryptoKey psk = makePsk("d4f1bb3a20290759f0bcffabcf4e6901");

    uint32_t fromNodeA = 0x11111111;
    uint32_t fromNodeB = 0x22222222;
    uint32_t toNode = 0x77777777;
    uint64_t packetIdA = 0xAAAABBBB;
    uint64_t packetIdB = 0xCCCCDDDD;

    uint8_t plaintext[6] = {0x01, 0x02, 0x03, 0x04, 0x05, 0x06};
    uint8_t ciphertextWithTag[6 + CryptoEngine::AEAD_TAG_SIZE];

    TEST_ASSERT_TRUE(crypto->encryptPacketCCM(psk, fromNodeA, toNode, packetIdA, 6, plaintext, ciphertextWithTag));

    uint8_t decrypted[6];
    // Wrong fromNode, right packetId
    TEST_ASSERT_FALSE(crypto->decryptPacketCCM(psk, fromNodeB, toNode, packetIdA, 6 + CryptoEngine::AEAD_TAG_SIZE,
                                               ciphertextWithTag, decrypted));
    // Right fromNode, wrong packetId
    TEST_ASSERT_FALSE(crypto->decryptPacketCCM(psk, fromNodeA, toNode, packetIdB, 6 + CryptoEngine::AEAD_TAG_SIZE,
                                               ciphertextWithTag, decrypted));
    // Both wrong
    TEST_ASSERT_FALSE(crypto->decryptPacketCCM(psk, fromNodeB, toNode, packetIdB, 6 + CryptoEngine::AEAD_TAG_SIZE,
                                               ciphertextWithTag, decrypted));
    // Both right still succeeds, so the assertions above are not passing for free
    TEST_ASSERT_TRUE(crypto->decryptPacketCCM(psk, fromNodeA, toNode, packetIdA, 6 + CryptoEngine::AEAD_TAG_SIZE,
                                              ciphertextWithTag, decrypted));
    TEST_ASSERT_EQUAL_MEMORY(plaintext, decrypted, 6);
}

void test_AES_CCM_AEAD_rejects_invalid_psk(void)
{
    // Empty PSK - must return false, not crash
    CryptoKey emptyPsk;
    memset(&emptyPsk, 0, sizeof(emptyPsk));
    emptyPsk.length = 0;

    uint32_t fromNode = 0xDEADBEEF;
    uint32_t toNode = 0x0000BEEF;
    uint64_t packetId = 0x12345678;

    uint8_t plaintext[8] = {0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08};
    uint8_t ciphertextWithTag[8 + CryptoEngine::AEAD_TAG_SIZE];
    uint8_t decrypted[8];

    // Encrypt with empty PSK must fail gracefully
    TEST_ASSERT_FALSE(crypto->encryptPacketCCM(emptyPsk, fromNode, toNode, packetId, 8, plaintext, ciphertextWithTag));

    // Decrypt with empty PSK must fail gracefully
    // (use dummy ciphertext since encrypt failed)
    memset(ciphertextWithTag, 0xAA, sizeof(ciphertextWithTag));
    TEST_ASSERT_FALSE(crypto->decryptPacketCCM(emptyPsk, fromNode, toNode, packetId, 8 + CryptoEngine::AEAD_TAG_SIZE,
                                               ciphertextWithTag, decrypted));

    // CryptoKey uses -1 as its "invalid key - do not use" sentinel, and it would widen
    // into a huge unsigned length rather than be rejected. Both directions must refuse it.
    CryptoKey invalidPsk;
    memset(&invalidPsk, 0, sizeof(invalidPsk));
    invalidPsk.length = -1;

    TEST_ASSERT_FALSE(crypto->encryptPacketCCM(invalidPsk, fromNode, toNode, packetId, 8, plaintext, ciphertextWithTag));
    TEST_ASSERT_FALSE(crypto->decryptPacketCCM(invalidPsk, fromNode, toNode, packetId, 8 + CryptoEngine::AEAD_TAG_SIZE,
                                               ciphertextWithTag, decrypted));
}

void test_AES_CCM_AEAD_key_size_distinction(void)
{
    // AES-128 vs AES-256 produce different ciphertexts
    // Verifies that 16-byte keys use true AES-128, not AES-256 with padding.
    // Same 16 bytes of key material, but one is AES-128 (16 bytes)
    // and the other is AES-256 (32 bytes, zero-padded).
    CryptoKey psk128 = makePsk("d4f1bb3a20290759f0bcffabcf4e6901");
    CryptoKey psk256;
    memset(psk256.bytes, 0, sizeof(psk256.bytes));
    HexToBytes(psk256.bytes, "d4f1bb3a20290759f0bcffabcf4e6901");
    psk256.length = 32; // same first 16 bytes, but treated as AES-256

    uint32_t fromNode = 0x55AA55AA;
    uint32_t toNode = 0x33333333;
    uint64_t packetId = 0x1234ABCD;

    uint8_t plaintext[8] = {0x10, 0x20, 0x30, 0x40, 0x50, 0x60, 0x70, 0x80};
    uint8_t ct128[8 + CryptoEngine::AEAD_TAG_SIZE];
    uint8_t ct256[8 + CryptoEngine::AEAD_TAG_SIZE];

    TEST_ASSERT_TRUE(crypto->encryptPacketCCM(psk128, fromNode, toNode, packetId, 8, plaintext, ct128));
    TEST_ASSERT_TRUE(crypto->encryptPacketCCM(psk256, fromNode, toNode, packetId, 8, plaintext, ct256));

    // AES-128 and AES-256 with the same key material must produce different output
    TEST_ASSERT_FALSE(memcmp(ct128, ct256, 8 + CryptoEngine::AEAD_TAG_SIZE) == 0);

    // Both must still round-trip correctly
    uint8_t dec128[8], dec256[8];
    TEST_ASSERT_TRUE(
        crypto->decryptPacketCCM(psk128, fromNode, toNode, packetId, 8 + CryptoEngine::AEAD_TAG_SIZE, ct128, dec128));
    TEST_ASSERT_EQUAL_MEMORY(plaintext, dec128, 8);
    TEST_ASSERT_TRUE(
        crypto->decryptPacketCCM(psk256, fromNode, toNode, packetId, 8 + CryptoEngine::AEAD_TAG_SIZE, ct256, dec256));
    TEST_ASSERT_EQUAL_MEMORY(plaintext, dec256, 8);

    // Cross-key decryption must fail
    TEST_ASSERT_FALSE(
        crypto->decryptPacketCCM(psk256, fromNode, toNode, packetId, 8 + CryptoEngine::AEAD_TAG_SIZE, ct128, dec128));
    TEST_ASSERT_FALSE(
        crypto->decryptPacketCCM(psk128, fromNode, toNode, packetId, 8 + CryptoEngine::AEAD_TAG_SIZE, ct256, dec256));
}

void test_AES_CCM_AEAD_binds_destination(void)
{
    // Rewritten destination - `to` is authenticated as associated data, so changing
    // it in flight must fail the tag check even though the nonce is unaffected
    CryptoKey psk = makePsk("d4f1bb3a20290759f0bcffabcf4e6901");

    uint32_t fromNode = 0x0A0B0C0D;
    uint32_t toNode = 0x00000042;
    uint32_t otherNode = 0x00000043;
    uint32_t broadcast = 0xFFFFFFFF;
    uint64_t packetId = 0x99887766;

    uint8_t plaintext[9] = {'t', 'o', '-', 'i', 's', '-', 'a', 'a', 'd'};
    uint8_t ciphertextWithTag[9 + CryptoEngine::AEAD_TAG_SIZE];
    uint8_t decrypted[9];

    TEST_ASSERT_TRUE(crypto->encryptPacketCCM(psk, fromNode, toNode, packetId, 9, plaintext, ciphertextWithTag));

    // Redirecting the packet to another node must be rejected
    TEST_ASSERT_FALSE(crypto->decryptPacketCCM(psk, fromNode, otherNode, packetId, 9 + CryptoEngine::AEAD_TAG_SIZE,
                                               ciphertextWithTag, decrypted));

    // Promoting a unicast to a broadcast must be rejected too
    TEST_ASSERT_FALSE(crypto->decryptPacketCCM(psk, fromNode, broadcast, packetId, 9 + CryptoEngine::AEAD_TAG_SIZE,
                                               ciphertextWithTag, decrypted));

    // The unmodified destination still round-trips, so the rejections above are not vacuous
    TEST_ASSERT_TRUE(
        crypto->decryptPacketCCM(psk, fromNode, toNode, packetId, 9 + CryptoEngine::AEAD_TAG_SIZE, ciphertextWithTag, decrypted));
    TEST_ASSERT_EQUAL_MEMORY(plaintext, decrypted, 9);

    // A different destination must also change the tag, not just be rejected on decrypt
    uint8_t otherCiphertextWithTag[9 + CryptoEngine::AEAD_TAG_SIZE];
    TEST_ASSERT_TRUE(crypto->encryptPacketCCM(psk, fromNode, otherNode, packetId, 9, plaintext, otherCiphertextWithTag));
    TEST_ASSERT_FALSE(memcmp(ciphertextWithTag + 9, otherCiphertextWithTag + 9, CryptoEngine::AEAD_TAG_SIZE) == 0);
}

void setup()
{
    // NOTE!!! Wait for >2 secs
    // if board doesn't support software reset via Serial.DTR/RTS
    delay(10);
    delay(2000);

    initializeTestEnvironment();
    UNITY_BEGIN(); // IMPORTANT LINE!
    RUN_TEST(test_SHA256);
    RUN_TEST(test_SHA256_large_input);
    RUN_TEST(test_ECB_AES128);
    RUN_TEST(test_ECB_AES256);
    RUN_TEST(test_DH25519);
    RUN_TEST(test_AES_CTR);
    RUN_TEST(test_AES_CCM_partial_block_bounds);
    RUN_TEST(test_AES_CCM_rfc3610);
    RUN_TEST(test_PKC);
    RUN_TEST(test_XEdDSA);
    RUN_TEST(test_XEdDSA_cross_key_reject);
    RUN_TEST(test_XEdDSA_empty_key_sign_fails);
    RUN_TEST(test_XEdDSA_curve_to_ed_cache);
    RUN_TEST(test_XEdDSA_max_payload);
    RUN_TEST(test_XEdDSA_repeated_sign_is_randomized);
    RUN_TEST(test_AES_CCM_AEAD_smoke);
    RUN_TEST(test_AES_CCM_AEAD_roundtrip_aes256);
    RUN_TEST(test_AES_CCM_AEAD_rejects_tampering);
    RUN_TEST(test_AES_CCM_AEAD_rejects_undersized);
    RUN_TEST(test_AES_CCM_AEAD_rejects_wrong_psk);
    RUN_TEST(test_AES_CCM_AEAD_roundtrip_aes128);
    RUN_TEST(test_AES_CCM_AEAD_tamper_sweep);
    RUN_TEST(test_AES_CCM_AEAD_is_deterministic);
    RUN_TEST(test_AES_CCM_AEAD_binds_nonce_inputs);
    RUN_TEST(test_AES_CCM_AEAD_rejects_invalid_psk);
    RUN_TEST(test_AES_CCM_AEAD_key_size_distinction);
    RUN_TEST(test_AES_CCM_AEAD_binds_destination);
    exit(UNITY_END()); // stop unit testing
}

void loop() {}