#include "ecies.hpp"
#include "mock_credential_storage.hpp"
#include "openssl_security_module.hpp"
#include <gtest/gtest.h>
#include <openssl/bn.h>
#include <openssl/ec.h>
#include <openssl/ecdh.h>
#include <openssl/obj_mac.h>

using vanetza::ByteBuffer;
using namespace vanetza::pki;

/**
 * Test vectors from IEEE 1609.2-2017 (Annex D)
 */

TEST(Ecies, kdf2_vector1)
{
    const ByteBuffer shared {{
        0x96, 0xC0, 0x56, 0x19, 0xD5, 0x6C, 0x32, 0x8A,
        0xB9, 0x5F, 0xE8, 0x4B, 0x18, 0x26, 0x4B, 0x08,
        0x72, 0x5B, 0x85, 0xE3, 0x3F, 0xD3, 0x4F, 0x08
    }};
    const ByteBuffer kdp;
    const ByteBuffer expected {{
        0x44, 0x30, 0x24, 0xc3, 0xda, 0xe6, 0x6b, 0x95,
        0xe6, 0xf5, 0x67, 0x06, 0x01, 0x55, 0x8f, 0x71
    }};

    OpenSslSecurityModule security(std::make_shared<MockCredentialStorage>());
    EXPECT_EQ(expected, calculate_kdf2(security, shared, kdp, expected.size()));
}

TEST(Ecies, kdf2_vector2)
{
    const ByteBuffer shared {{
        0x96, 0xF6, 0x00, 0xB7, 0x3A, 0xD6, 0xAC, 0x56,
        0x29, 0x57, 0x7E, 0xCE, 0xD5, 0x17, 0x43, 0xDD,
        0x2C, 0x24, 0xC2, 0x1B, 0x1A, 0xC8, 0x3E, 0xE4
    }};
    const ByteBuffer kdp;
    const ByteBuffer expected {{
        0xb6, 0x29, 0x51, 0x62, 0xa7, 0x80, 0x4f, 0x56,
        0x67, 0xba, 0x90, 0x70, 0xf8, 0x2f, 0xa5, 0x22
    }};

    OpenSslSecurityModule security(std::make_shared<MockCredentialStorage>());
    EXPECT_EQ(expected, calculate_kdf2(security, shared, kdp, expected.size()));
}

TEST(Ecies, kdf2_vector3)
{
    const ByteBuffer shared {{
        0x22, 0x51, 0x8B, 0x10, 0xE7, 0x0F, 0x2A, 0x3F,
        0x24, 0x38, 0x10, 0xAE, 0x32, 0x54, 0x13, 0x9E,
        0xFB, 0xEE, 0x04, 0xAA, 0x57, 0xC7, 0xAF, 0x7D
    }};
    const ByteBuffer kdp {{
        0x75, 0xEE, 0xF8, 0x1A, 0xA3, 0x04, 0x1E, 0x33,
        0xB8, 0x09, 0x71, 0x20, 0x3D, 0x2C, 0x0C, 0x52
    }};
    const ByteBuffer expected {{
        0xc4, 0x98, 0xaf, 0x77, 0x16, 0x1c, 0xc5, 0x9f,
        0x29, 0x62, 0xb9, 0xa7, 0x13, 0xe2, 0xb2, 0x15,
        0x15, 0x2d, 0x13, 0x97, 0x66, 0xce, 0x34, 0xa7,
        0x76, 0xdf, 0x11, 0x86, 0x6a, 0x69, 0xbf, 0x2e,
        0x52, 0xa1, 0x3d, 0x9c, 0x7c, 0x6f, 0xc8, 0x78,
        0xc5, 0x0c, 0x5e, 0xa0, 0xbc, 0x7b, 0x00, 0xe0,
        0xda, 0x24, 0x47, 0xcf, 0xd8, 0x74, 0xf6, 0xcf,
        0x92, 0xf3, 0x0d, 0x00, 0x97, 0x11, 0x14, 0x85,
        0x50, 0x0c, 0x90, 0xc3, 0xaf, 0x8b, 0x48, 0x78,
        0x72, 0xd0, 0x46, 0x85, 0xd1, 0x4c, 0x8d, 0x1d,
        0xc8, 0xd7, 0xfa, 0x08, 0xbe, 0xb0, 0xce, 0x0a,
        0xba, 0xbc, 0x11, 0xf0, 0xbd, 0x49, 0x62, 0x69,
        0x14, 0x2d, 0x43, 0x52, 0x5a, 0x78, 0xe5, 0xbc,
        0x79, 0xa1, 0x7f, 0x59, 0x67, 0x6a, 0x57, 0x06,
        0xdc, 0x54, 0xd5, 0x4d, 0x4d, 0x1f, 0x0b, 0xd7,
        0xe3, 0x86, 0x12, 0x8e, 0xc2, 0x6a, 0xfc, 0x21
    }};

    OpenSslSecurityModule security(std::make_shared<MockCredentialStorage>());
    EXPECT_EQ(expected, calculate_kdf2(security, shared, kdp, expected.size()));
}

TEST(Ecies, kdf2_vector4)
{
    const ByteBuffer shared {{
        0x7E, 0x33, 0x5A, 0xFA, 0x4B, 0x31, 0xD7, 0x72,
        0xC0, 0x63, 0x5C, 0x7B, 0x0E, 0x06, 0xF2, 0x6F,
        0xCD, 0x78, 0x1D, 0xF9, 0x47, 0xD2, 0x99, 0x0A
    }};
    const ByteBuffer kdp {{
        0xD6, 0x5A, 0x48, 0x12, 0x73, 0x3F, 0x8C, 0xDB,
        0xCD, 0xFB, 0x4B, 0x2F, 0x4C, 0x19, 0x1D, 0x87
    }};
    const ByteBuffer expected {{
        0xc0, 0xbd, 0x9e, 0x38, 0xa8, 0xf9, 0xde, 0x14,
        0xc2, 0xac, 0xd3, 0x5b, 0x2f, 0x34, 0x10, 0xc6,
        0x98, 0x8c, 0xf0, 0x24, 0x00, 0x54, 0x36, 0x31,
        0xe0, 0xd6, 0xa4, 0xc1, 0xd0, 0x30, 0x36, 0x5a,
        0xcb, 0xf3, 0x98, 0x11, 0x5e, 0x51, 0xaa, 0xdd,
        0xeb, 0xdc, 0x95, 0x90, 0x66, 0x42, 0x10, 0xf9,
        0xaa, 0x9f, 0xed, 0x77, 0x0d, 0x4c, 0x57, 0xed,
        0xea, 0xfa, 0x0b, 0x8c, 0x14, 0xf9, 0x33, 0x00,
        0x86, 0x52, 0x51, 0x21, 0x8c, 0x26, 0x2d, 0x63,
        0xda, 0xdc, 0x47, 0xdf, 0xa0, 0xe0, 0x28, 0x48,
        0x26, 0x79, 0x39, 0x85, 0x13, 0x7e, 0x0a, 0x54,
        0x4e, 0xc8, 0x0a, 0xbf, 0x2f, 0xdf, 0x5a, 0xb9,
        0x0b, 0xda, 0xea, 0x66, 0x20, 0x40, 0x12, 0xef,
        0xe3, 0x49, 0x71, 0xdc, 0x43, 0x1d, 0x62, 0x5c,
        0xd9, 0xa3, 0x29, 0xb8, 0x21, 0x7c, 0xc8, 0xfd,
        0x0d, 0x9f, 0x02, 0xb1, 0x3f, 0x2f, 0x6b, 0x0b
    }};

    OpenSslSecurityModule security(std::make_shared<MockCredentialStorage>());
    EXPECT_EQ(expected, calculate_kdf2(security, shared, kdp, expected.size()));
}

// Regression test: previously, the 0_EU-EA_L0 certificate's encryption key is
// stored in compressed-y-1 form and the ECIES context must still derive a
// working shared secret from it.
TEST(Ecies, shared_secret_with_compressed_y1_key)
{
    // x coordinate of the 0_EU-EA_L0 EA certificate encryption key
    const ByteBuffer eu_ea_x {{
        0x53, 0x8A, 0x8D, 0x36, 0x0D, 0x00, 0xD8, 0x48,
        0x0F, 0x5B, 0x29, 0x54, 0xFA, 0xB2, 0x42, 0xFB,
        0x98, 0xB3, 0x38, 0x70, 0xF7, 0xA0, 0xC3, 0x3C,
        0xF6, 0x52, 0xCB, 0x46, 0xA1, 0x74, 0xE2, 0x48
    }};

    PublicKey compressed;
    compressed.type = KeyType::NistP256;
    compressed.compression = KeyCompression::Y1;
    compressed.x = eu_ea_x;

    OpenSslSecurityModule security(std::make_shared<MockCredentialStorage>());
    Sha256Hash info {};
    // Must not throw and must produce a non-empty 32-byte shared secret
    auto ecies = security.create_ecies_context(compressed, info);
    EXPECT_EQ(32u, ecies->shared_secret().size());
    EXPECT_FALSE(ecies->shared_secret() == ByteBuffer(32, 0));
}

TEST(Ecies, encryption_roundtrip_compressed_receiver)
{
    // Receiver was created as uncompressed but is then compressed before handing it
    // to the ECIES context. Decryption must still succeed -- otherwise we have a
    // bug in how compressed keys are mapped back to curve points.
    OpenSslSecurityModule security(std::make_shared<MockCredentialStorage>());
    PublicKey receiver = security.create_key(KeyType::NistP256);
    ASSERT_NE(KeyCompression::NoCompression, receiver.compression);

    Sha256Hash info { 0x12, 0x34, 0x56, 0x78, 0x90, 0xab, 0xcd, 0xef };
    auto ecies = security.create_ecies_context(receiver, info);

    const ByteBuffer plaintext {{
        0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77,
        0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff
    }};
    const ByteBuffer ciphertext = ecies->encrypt(plaintext);
    const ByteBuffer decrypted = ecies->decrypt(ciphertext.data(), ciphertext.size());
    EXPECT_NE(plaintext, ciphertext);
    EXPECT_EQ(plaintext, decrypted);
}

TEST(Ecies, encryption_roundtrip)
{
    OpenSslSecurityModule security(std::make_shared<MockCredentialStorage>());
    PublicKey receiver = security.create_key(KeyType::NistP256);
    Sha256Hash info { 0x12, 0x34, 0x56, 0x78, 0x90, 0xab, 0xcd, 0xef };
    auto ecies = security.create_ecies_context(receiver, info);

    const ByteBuffer plaintext {{
        0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77,
        0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff
    }};
    const ByteBuffer ciphertext = ecies->encrypt(plaintext);
    const ByteBuffer decrypted = ecies->decrypt(ciphertext.data(), ciphertext.size());
    EXPECT_NE(plaintext, ciphertext);
    EXPECT_EQ(plaintext, decrypted);
}

namespace
{

// Mock that returns fixed test-vector values, so encrypt_key() can be tested
// against IEEE 1609.2 Annex D.6.2 known-answer data.
class FixedEciesContext : public SecurityModule::EciesContext
{
public:
    FixedEciesContext(ByteBuffer shared) : m_shared_secret(std::move(shared))
    {
    }

    PublicKey ephemeral_public_key() const override
    {
        return {};
    }

    PublicKey recipient_public_key() const override
    {
        return {};
    }

    ByteBuffer shared_secret() const override
    {
        return m_shared_secret;
    }

    ByteBuffer encrypted_key() const override
    {
        return {};
    }

    ByteBuffer authentication_tag() const override
    {
        return {};
    }

    ByteBuffer nonce() const override
    {
        return {};
    }

    void nonce(const ByteBuffer&) override
    {
    }

    ByteBuffer encrypt(const ByteBuffer&) override
    {
        return {};
    }

    ByteBuffer decrypt(const std::uint8_t*, std::size_t) override
    {
        return {};
    }

private:
    ByteBuffer m_shared_secret;
};

} // anonymous namespace

// IEEE 1609.2-2016 Annex D.6.2 ECIES1 test vector
// End-to-end: runs ECDH with the vector's ephemeral private key and recipient
// public key through OpenSSL, then pipes the resulting shared secret through
// encrypt_key(). This catches any bug in the ECDH input handling that a
// standalone KDF2 test would miss.
TEST(Ecies, ecies_IEEE_1609_2_Annex_D_6_2_ECIES1_full)
{
    const ByteBuffer sender_eph_priv {{
        0x13, 0x84, 0xC3, 0x1D, 0x69, 0x82, 0xD5, 0x2B,
        0xCA, 0x3B, 0xED, 0x8A, 0x7E, 0x60, 0xF5, 0x2F,
        0xEC, 0xDA, 0xB4, 0x4E, 0x5C, 0x0E, 0xA1, 0x66,
        0x81, 0x5A, 0x81, 0x59, 0xE0, 0x9F, 0xFB, 0x42
    }};
    const ByteBuffer recipient_pub_x {{
        0x8C, 0x5E, 0x20, 0xFE, 0x31, 0x93, 0x5F, 0x6F,
        0xA6, 0x82, 0xA1, 0xF6, 0xD4, 0x6E, 0x44, 0x68,
        0x53, 0x4F, 0xFE, 0xA1, 0xA6, 0x98, 0xB1, 0x4B,
        0x0B, 0x12, 0x51, 0x3E, 0xED, 0x8D, 0xEB, 0x11
    }};
    const ByteBuffer recipient_pub_y {{
        0x12, 0x70, 0xFE, 0xC2, 0x42, 0x7E, 0x6A, 0x15,
        0x4D, 0xFC, 0xAE, 0x33, 0x68, 0x58, 0x43, 0x96,
        0xC8, 0x25, 0x1A, 0x04, 0xE2, 0xAE, 0x7D, 0x87,
        0xB0, 0x16, 0xFF, 0x65, 0xD2, 0x2D, 0x6F, 0x9E
    }};
    const ByteBuffer aes_key {{
        0x91, 0x69, 0x15, 0x5B, 0x08, 0xB0, 0x76, 0x74,
        0xCB, 0xAD, 0xF7, 0x5F, 0xB4, 0x6A, 0x7B, 0x0D
    }};
    const Sha256Hash recipient_info {{
        0xA6, 0xB7, 0xB5, 0x25, 0x54, 0xB4, 0x20, 0x3F,
        0x7E, 0x3A, 0xCF, 0xDB, 0x3A, 0x3E, 0xD8, 0x67,
        0x4E, 0xE0, 0x86, 0xCE, 0x59, 0x06, 0xA7, 0xCA,
        0xC2, 0xF8, 0xA3, 0x98, 0x30, 0x6D, 0x3B, 0xE9
    }};
    const ByteBuffer expected_wrapped {{
        0xA6, 0x34, 0x20, 0x13, 0xD6, 0x23, 0xAD, 0x6C,
        0x5F, 0x68, 0x82, 0x46, 0x96, 0x73, 0xAE, 0x33
    }};
    const ByteBuffer expected_tag {{
        0x80, 0xE1, 0xD8, 0x5D, 0x30, 0xF1, 0xBA, 0xE4,
        0xEC, 0xF1, 0xA5, 0x34, 0xA8, 0x9A, 0x07, 0x86
    }};

    // Set up sender ephemeral EC_KEY with the test vector's private scalar
    EC_KEY* sender = EC_KEY_new_by_curve_name(NID_X9_62_prime256v1);
    BIGNUM* priv_bn = BN_bin2bn(sender_eph_priv.data(), sender_eph_priv.size(), nullptr);
    ASSERT_EQ(1, EC_KEY_set_private_key(sender, priv_bn));

    // Reconstruct recipient public point from (x, y)
    const EC_GROUP* group = EC_KEY_get0_group(sender);
    EC_POINT* recipient_point = EC_POINT_new(group);
    BIGNUM* rx = BN_bin2bn(recipient_pub_x.data(), recipient_pub_x.size(), nullptr);
    BIGNUM* ry = BN_bin2bn(recipient_pub_y.data(), recipient_pub_y.size(), nullptr);
    ASSERT_EQ(1, EC_POINT_set_affine_coordinates(group, recipient_point, rx, ry, nullptr));

    // Compute shared secret
    ByteBuffer shared(32);
    int got = ECDH_compute_key(shared.data(), shared.size(), recipient_point, sender, nullptr);
    ASSERT_EQ(32, got);

    BN_free(priv_bn);
    BN_free(rx);
    BN_free(ry);
    EC_POINT_free(recipient_point);
    EC_KEY_free(sender);

    // Now feed the shared secret through encrypt_key
    OpenSslSecurityModule security(std::make_shared<MockCredentialStorage>());
    FixedEciesContext ecies(shared);
    auto result = encrypt_key(security, ecies, aes_key, recipient_info);

    EXPECT_EQ(expected_wrapped, result.wrapped_key);
    EXPECT_EQ(expected_tag, result.authentication_tag);
}

// IEEE 1609.2-2016 Annex D.6.2 ECIES1 test vector
// Verifies the KDF2 + key wrapping + HMAC-SHA256 tag chain used by encrypt_key().
TEST(Ecies, encrypt_key_IEEE_1609_2_Annex_D_6_2_ECIES1)
{
    // Shared secret z = x-coord of [k_E] * Q_recipient (precomputed)
    const ByteBuffer shared_secret {{
        0xd4, 0x43, 0x08, 0x02, 0x3f, 0xbb, 0xd3, 0xe9,
        0x06, 0xb9, 0xf3, 0xb4, 0x0e, 0x5e, 0x0b, 0x62,
        0x54, 0x11, 0x70, 0x3c, 0x4a, 0x99, 0x24, 0x9d,
        0x35, 0xa1, 0x39, 0x06, 0x38, 0x6a, 0x3e, 0xe3
    }};
    const ByteBuffer aes_key {{
        0x91, 0x69, 0x15, 0x5B, 0x08, 0xB0, 0x76, 0x74,
        0xCB, 0xAD, 0xF7, 0x5F, 0xB4, 0x6A, 0x7B, 0x0D
    }};
    const Sha256Hash recipient_info {{
        0xA6, 0xB7, 0xB5, 0x25, 0x54, 0xB4, 0x20, 0x3F,
        0x7E, 0x3A, 0xCF, 0xDB, 0x3A, 0x3E, 0xD8, 0x67,
        0x4E, 0xE0, 0x86, 0xCE, 0x59, 0x06, 0xA7, 0xCA,
        0xC2, 0xF8, 0xA3, 0x98, 0x30, 0x6D, 0x3B, 0xE9
    }};
    const ByteBuffer expected_wrapped {{
        0xA6, 0x34, 0x20, 0x13, 0xD6, 0x23, 0xAD, 0x6C,
        0x5F, 0x68, 0x82, 0x46, 0x96, 0x73, 0xAE, 0x33
    }};
    const ByteBuffer expected_tag {{
        0x80, 0xE1, 0xD8, 0x5D, 0x30, 0xF1, 0xBA, 0xE4,
        0xEC, 0xF1, 0xA5, 0x34, 0xA8, 0x9A, 0x07, 0x86
    }};

    OpenSslSecurityModule security(std::make_shared<MockCredentialStorage>());
    FixedEciesContext ecies(shared_secret);
    auto result = encrypt_key(security, ecies, aes_key, recipient_info);

    EXPECT_EQ(expected_wrapped, result.wrapped_key);
    EXPECT_EQ(expected_tag, result.authentication_tag);
}