import os, sys, tempfile
from pathlib import Path

import pytest
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric.utils import decode_dss_signature

REPO_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(REPO_ROOT / "tools"))
from sign_firmware import sign_firmware, load_private_key


@pytest.fixture()
def keypair(tmp_path):
    key = ec.generate_private_key(ec.SECP256R1())
    pem_path = tmp_path / "key.pem"
    pem_path.write_bytes(key.private_bytes(
        serialization.Encoding.PEM,
        serialization.PrivateFormat.PKCS8,
        serialization.NoEncryption(),
    ))
    return key, pem_path


def test_signature_is_64_bytes(keypair, tmp_path):
    key, key_path = keypair
    firmware = tmp_path / "fw.bin"
    firmware.write_bytes(b"fake firmware data" * 100)
    sig = sign_firmware(firmware, key_path)
    assert len(sig) == 64


def test_signature_verifies(keypair, tmp_path):
    key, key_path = keypair
    data = b"test firmware payload"
    firmware = tmp_path / "fw.bin"
    firmware.write_bytes(data)
    sig_raw = sign_firmware(firmware, key_path)

    # Convert raw r||s back to DER for cryptography lib verify
    r = int.from_bytes(sig_raw[:32], 'big')
    s = int.from_bytes(sig_raw[32:], 'big')
    from cryptography.hazmat.primitives.asymmetric.utils import encode_dss_signature
    sig_der = encode_dss_signature(r, s)

    key.public_key().verify(sig_der, data, ec.ECDSA(hashes.SHA256()))


def test_wrong_key_fails_verification(keypair, tmp_path):
    key, key_path = keypair
    firmware = tmp_path / "fw.bin"
    firmware.write_bytes(b"firmware")
    sig_raw = sign_firmware(firmware, key_path)

    other_key = ec.generate_private_key(ec.SECP256R1())
    r = int.from_bytes(sig_raw[:32], 'big')
    s = int.from_bytes(sig_raw[32:], 'big')
    from cryptography.hazmat.primitives.asymmetric.utils import encode_dss_signature
    sig_der = encode_dss_signature(r, s)

    with pytest.raises(Exception):
        other_key.public_key().verify(sig_der, b"firmware", ec.ECDSA(hashes.SHA256()))