import pytest
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey
from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat

from p2p.crypto import NoiseInitiator, NoiseResponder, frame_encrypt, frame_decrypt, frame_decrypt_length
from p2p.errors import CryptoError


def _handshake():
    si, sr = X25519PrivateKey.generate(), X25519PrivateKey.generate()
    i, r = NoiseInitiator(si), NoiseResponder(sr)
    m1 = i.write_message_1()
    m2 = r.read_message_1(m1)
    m3 = i.read_message_2(m2)
    r.read_message_3(m3)
    return i.finalize(), r.finalize()


class TestNoise:
    def test_handshake(self):
        _handshake()

    def test_keys(self):
        si, sr = X25519PrivateKey.generate(), X25519PrivateKey.generate()
        i, r = NoiseInitiator(si), NoiseResponder(sr)
        m2 = r.read_message_1(i.write_message_1())
        m3 = i.read_message_2(m2)
        r.read_message_3(m3)
        iss, rss = i.finalize(), r.finalize()
        assert iss.remote_static_pub == sr.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)
        assert rss.remote_static_pub == si.public_key().public_bytes(Encoding.Raw, PublicFormat.Raw)

    def test_encrypt_decrypt(self):
        iss, rss = _handshake()
        assert rss.decrypt(iss.encrypt(b"hello")) == b"hello"
        assert iss.decrypt(rss.encrypt(b"world")) == b"world"

    def test_multi(self):
        iss, rss = _handshake()
        for n in range(20):
            m = f"msg{n}".encode()
            assert rss.decrypt(iss.encrypt(m)) == m

    def test_tamper(self):
        iss, rss = _handshake()
        ct = bytearray(iss.encrypt(b"x"))
        ct[-1] ^= 0xFF
        with pytest.raises(CryptoError):
            rss.decrypt(bytes(ct))

    def test_framing(self):
        iss, rss = _handshake()
        f = frame_encrypt(iss, b"framed")
        l = frame_decrypt_length(f[:4])
        assert frame_decrypt(rss, f[4:4+l]) == b"framed"

    def test_nonce_overflow_raises(self):
        iss, rss = _handshake()
        iss._send.n = 2**63 - 2
        iss.encrypt(b"ok")  # n = 2**63 - 2, should work
        with pytest.raises(CryptoError, match="Nonce exhausted"):
            iss.encrypt(b"overflow")  # n = 2**63 - 1, should raise

    def test_nonce_overflow_decrypt_raises(self):
        iss, rss = _handshake()
        rss._recv.n = 2**63 - 1
        with pytest.raises(CryptoError, match="Nonce exhausted"):
            rss.decrypt(b"\x00" * 32)
