from cryptography.hazmat.primitives.asymmetric import ec

from aliro_harness.reader.crypto import (
    derive_kdh,
    ecdh_shared_x,
    hkdf_sha256,
)


# RFC 5869 Test Case 1
RFC5869_T1_IKM = bytes.fromhex("0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b")
RFC5869_T1_SALT = bytes.fromhex("000102030405060708090a0b0c")
RFC5869_T1_INFO = bytes.fromhex("f0f1f2f3f4f5f6f7f8f9")
RFC5869_T1_OKM_42 = bytes.fromhex(
    "3cb25f25faacd57a90434f64d0362f2a"
    "2d2d0a90cf1a5a4c5db02d56ecc4c5bf"
    "34007208d5b887185865"
)


def test_hkdf_sha256_matches_rfc5869_test_case_1():
    assert hkdf_sha256(RFC5869_T1_IKM, RFC5869_T1_SALT, RFC5869_T1_INFO, 42) == RFC5869_T1_OKM_42


def test_ecdh_commutative():
    a = ec.generate_private_key(ec.SECP256R1())
    b = ec.generate_private_key(ec.SECP256R1())
    a_pub_uncomp = bytes([0x04]) + a.public_key().public_numbers().x.to_bytes(32, "big") + a.public_key().public_numbers().y.to_bytes(32, "big")
    b_pub_uncomp = bytes([0x04]) + b.public_key().public_numbers().x.to_bytes(32, "big") + b.public_key().public_numbers().y.to_bytes(32, "big")
    assert ecdh_shared_x(a, b_pub_uncomp) == ecdh_shared_x(b, a_pub_uncomp)


def test_derive_kdh_agrees_on_both_sides():
    a = ec.generate_private_key(ec.SECP256R1())
    b = ec.generate_private_key(ec.SECP256R1())
    a_pub = bytes([0x04]) + a.public_key().public_numbers().x.to_bytes(32, "big") + a.public_key().public_numbers().y.to_bytes(32, "big")
    b_pub = bytes([0x04]) + b.public_key().public_numbers().x.to_bytes(32, "big") + b.public_key().public_numbers().y.to_bytes(32, "big")
    txn = b"\x01" * 16
    assert derive_kdh(a, b_pub, txn) == derive_kdh(b, a_pub, txn)