// 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 <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_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 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_AES256);
    RUN_TEST(test_DH25519);
    RUN_TEST(test_AES_CTR);
    RUN_TEST(test_AES_CCM_partial_block_bounds);
    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);
    exit(UNITY_END()); // stop unit testing
}

void loop() {}
