erldistpy/tests/test_handshake.py

"""Distribution handshake tests.

Unit tests cover the frame builders/parsers and the cookie digest
formula (cross-checked against ``erlang:md5/1`` output). Live tests
spawn an Erlang node with a known cookie and drive a full handshake
end-to-end.
"""

from __future__ import annotations

import shutil
import socket
import struct
import subprocess
import time

import pytest

from erldistpy.epmd import lookup
from erldistpy.flags import DEFAULT_FLAGS, DFLAG_HANDSHAKE_23
from erldistpy.handshake import (
    HandshakeError,
    build_challenge_reply,
    build_send_name,
    cookie_digest,
    handshake,
    parse_challenge_ack,
    parse_recv_challenge,
    parse_recv_status,
)

# Reference value computed by `erlang:md5("TESTCOOKIE" ++ integer_to_list(1234567890))`
COOKIE_DIGEST_REF = bytes.fromhex("1f856146c1dd930d59ef38a9eeef21f6")


# --------------------------------------------------------------------------
# Cookie digest
# --------------------------------------------------------------------------


def test_cookie_digest_matches_erlang_reference():
    assert cookie_digest("TESTCOOKIE", 1234567890) == COOKIE_DIGEST_REF


def test_cookie_digest_treats_challenge_as_unsigned():
    # The wire challenge is a uint32. -1 should hash like 0xFFFFFFFF.
    a = cookie_digest("X", -1)
    b = cookie_digest("X", 0xFFFFFFFF)
    assert a == b


def test_cookie_digest_length_is_16():
    assert len(cookie_digest("anything", 0)) == 16


# --------------------------------------------------------------------------
# Frame builders / parsers
# --------------------------------------------------------------------------


def test_build_send_name_shape():
    frame = build_send_name("erldistpy@localhost", DEFAULT_FLAGS, creation=42)
    assert frame[0] == 0x4E
    flags = struct.unpack(">Q", frame[1:9])[0]
    creation = struct.unpack(">I", frame[9:13])[0]
    nlen = struct.unpack(">H", frame[13:15])[0]
    name = frame[15:15 + nlen]
    assert flags == DEFAULT_FLAGS
    assert creation == 42
    assert name == b"erldistpy@localhost"


def test_parse_recv_status_ok():
    assert parse_recv_status(b"sok") == "ok"


def test_parse_recv_status_nok():
    assert parse_recv_status(b"snot_allowed") == "not_allowed"


def test_parse_recv_status_wrong_tag():
    with pytest.raises(HandshakeError, match="recv_status tag"):
        parse_recv_status(b"Xok")


def test_parse_recv_challenge_roundtrip():
    name = b"peer@host"
    flags = DEFAULT_FLAGS
    challenge = 0xDEADBEEF
    creation = 7
    frame = (
        bytes([0x4E])
        + struct.pack(">Q", flags)
        + struct.pack(">I", challenge)
        + struct.pack(">I", creation)
        + struct.pack(">H", len(name))
        + name
    )
    f, c, cr, n = parse_recv_challenge(frame)
    assert (f, c, cr, n) == (flags, challenge, creation, "peer@host")


def test_parse_recv_challenge_truncated():
    with pytest.raises(HandshakeError, match="truncated"):
        parse_recv_challenge(bytes([0x4E]) + b"\x00" * 5)


def test_challenge_reply_roundtrip():
    digest = b"\xAA" * 16
    frame = build_challenge_reply(0x12345678, digest)
    assert frame[0] == 0x72
    assert struct.unpack(">I", frame[1:5])[0] == 0x12345678
    assert frame[5:] == digest


def test_challenge_reply_rejects_wrong_digest_length():
    with pytest.raises(ValueError):
        build_challenge_reply(0, b"\x00" * 15)


def test_parse_challenge_ack_ok():
    digest = b"\xBB" * 16
    assert parse_challenge_ack(bytes([0x61]) + digest) == digest


def test_parse_challenge_ack_wrong_tag():
    with pytest.raises(HandshakeError, match="challenge_ack tag"):
        parse_challenge_ack(bytes([0x60]) + b"\x00" * 16)


def test_parse_challenge_ack_wrong_length():
    with pytest.raises(HandshakeError, match="wrong length"):
        parse_challenge_ack(bytes([0x61]) + b"\x00" * 15)


# --------------------------------------------------------------------------
# Top-level handshake() validation
# --------------------------------------------------------------------------


def test_handshake_requires_qualified_name():
    sock = socket.socket()
    with pytest.raises(ValueError, match="name@host"):
        handshake(sock, our_name="justname", cookie="X")


def test_handshake_requires_v6_flag():
    sock = socket.socket()
    with pytest.raises(ValueError, match="DFLAG_HANDSHAKE_23"):
        handshake(sock, our_name="x@h", cookie="X", our_flags=0)


# --------------------------------------------------------------------------
# Live interop — spawn an erl node with a known cookie and shake hands
# --------------------------------------------------------------------------


def _epmd_running() -> bool:
    try:
        with socket.create_connection(("localhost", 4369), timeout=0.5):
            return True
    except OSError:
        return False


@pytest.fixture(scope="module")
def live_erlang_peer():
    """Spawn `erl -sname erldistpy_hs -setcookie ERLDISTPY_HS_COOKIE`.

    Returns (sname, cookie). Skips the suite if erl/EPMD are absent.
    """
    if not shutil.which("erl"):
        pytest.skip("erl not installed")
    if not _epmd_running():
        pytest.skip("EPMD not running on localhost")

    sname = "erldistpy_hs"
    cookie = "ERLDISTPY_HS_COOKIE"
    proc = subprocess.Popen(
        [
            "erl",
            "-sname", sname,
            "-setcookie", cookie,
            "-noshell",
            "-eval", "timer:sleep(infinity).",
        ],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    for _ in range(40):
        time.sleep(0.1)
        if lookup(sname, timeout=0.5) is not None:
            break
    else:
        proc.terminate()
        proc.wait(timeout=2)
        pytest.skip("erl node failed to register with EPMD")

    yield sname, cookie

    proc.terminate()
    try:
        proc.wait(timeout=2)
    except subprocess.TimeoutExpired:
        proc.kill()


def _connect_to(sname: str) -> socket.socket:
    info = lookup(sname)
    assert info is not None, f"node {sname} not in EPMD"
    return socket.create_connection(("localhost", info.port), timeout=2)


def test_live_handshake_succeeds(live_erlang_peer):
    sname, cookie = live_erlang_peer
    sock = _connect_to(sname)
    try:
        # Use a distinct sname for our client side
        result = handshake(
            sock,
            our_name=f"erldistpy_client@{socket.gethostname()}",
            cookie=cookie,
        )
        assert result.peer_name.startswith(sname + "@")
        assert result.peer_flags & DFLAG_HANDSHAKE_23
        assert result.peer_creation > 0
        assert result.our_creation > 0
    finally:
        sock.close()


def test_live_handshake_bad_cookie_fails(live_erlang_peer):
    sname, _ = live_erlang_peer
    sock = _connect_to(sname)
    try:
        with pytest.raises(HandshakeError, match="digest mismatch|peer closed"):
            handshake(
                sock,
                our_name=f"erldistpy_client_bad@{socket.gethostname()}",
                cookie="WRONG_COOKIE",
            )
    finally:
        sock.close()