erldistpy/erldistpy/handshake.py

"""Distribution handshake — v6 (DFLAG_HANDSHAKE_23) only.

Drives a freshly-opened TCP socket through the OTP 23+ handshake dance:

    client                          server
      |  send_name ('N')  ─────────► |
      |  ◄─────── recv_status ('s')  |
      |  ◄─── recv_challenge ('N')   |
      |  send_challenge_reply ('r') ►|
      |  ◄── recv_challenge_ack ('a')|

All frames during the handshake are wrapped in a 2-byte big-endian
length prefix. After ``handshake()`` returns the socket is authenticated
and ready for the post-handshake data phase (4-byte framing, Phase 4).

Cookie discipline: ``cookie`` is a string passed in by the caller and
never logged by this module. Treat it like any other secret.

Spec: https://www.erlang.org/doc/apps/erts/erl_dist_protocol.html#distribution-handshake
"""

from __future__ import annotations

import hashlib
import os
import socket
import struct
from dataclasses import dataclass

from erldistpy.flags import DEFAULT_FLAGS, DFLAG_HANDSHAKE_23

TAG_SEND_NAME = 0x4E         # 'N'
TAG_RECV_STATUS = 0x73       # 's'
TAG_RECV_CHALLENGE = 0x4E    # 'N' (same as send_name; distinguished by direction)
TAG_CHALLENGE_REPLY = 0x72   # 'r'
TAG_CHALLENGE_ACK = 0x61     # 'a'


class HandshakeError(RuntimeError):
    """Handshake protocol or I/O failure."""


@dataclass(frozen=True)
class HandshakeResult:
    peer_name: str
    peer_flags: int
    peer_creation: int
    our_flags: int
    our_creation: int


def handshake(
    sock: socket.socket,
    *,
    our_name: str,
    cookie: str,
    our_flags: int = DEFAULT_FLAGS,
    our_creation: int | None = None,
    timeout: float = 5.0,
    frame_size: int = 2,
) -> HandshakeResult:
    """Drive ``sock`` through a v6 distribution handshake.

    ``our_name`` is the full ``name@host`` we are calling as.
    ``cookie`` is the shared secret with the peer.

    ``frame_size`` is the length-prefix size in bytes — 2 for plain TCP
    dist (``inet_tcp_dist``, the default), 4 for TLS dist
    (``inet_tls_dist``, which sets ``{packet, 4}`` on the SSL socket
    before the handshake). Pass ``4`` whenever the socket is wrapped in
    TLS; :class:`Node` selects this automatically based on ``tls_context``.

    On success the socket is left open and authenticated. On any
    protocol or auth failure raises :class:`HandshakeError`.
    """
    if "@" not in our_name:
        raise ValueError(f"our_name must look like 'name@host', got {our_name!r}")
    if not (our_flags & DFLAG_HANDSHAKE_23):
        raise ValueError("our_flags must include DFLAG_HANDSHAKE_23 (v6 handshake)")
    if frame_size not in (2, 4):
        raise ValueError(f"frame_size must be 2 or 4, got {frame_size}")
    if our_creation is None:
        our_creation = _random_creation()

    sock.settimeout(timeout)

    # Step 1: send_name
    _send_frame(sock, build_send_name(our_name, our_flags, our_creation), frame_size)

    # Step 2: recv_status
    status_frame = _recv_frame(sock, frame_size)
    status = parse_recv_status(status_frame)
    if status not in ("ok", "ok_simultaneous"):
        raise HandshakeError(f"peer rejected handshake: status={status!r}")

    # Step 3: recv_challenge
    challenge_frame = _recv_frame(sock, frame_size)
    peer_flags, peer_challenge, peer_creation, peer_name = parse_recv_challenge(challenge_frame)
    if not (peer_flags & DFLAG_HANDSHAKE_23):
        raise HandshakeError("peer does not speak v6 handshake (no DFLAG_HANDSHAKE_23)")

    # Step 4: send_challenge_reply with our challenge + MD5(cookie || peer_challenge)
    our_challenge = _random_challenge()
    reply_digest = cookie_digest(cookie, peer_challenge)
    _send_frame(sock, build_challenge_reply(our_challenge, reply_digest), frame_size)

    # Step 5: recv_challenge_ack — peer must answer with MD5(cookie || our_challenge)
    ack_frame = _recv_frame(sock, frame_size)
    ack_digest = parse_challenge_ack(ack_frame)
    expected = cookie_digest(cookie, our_challenge)
    if ack_digest != expected:
        raise HandshakeError("cookie digest mismatch — wrong cookie or peer impostor")

    return HandshakeResult(
        peer_name=peer_name,
        peer_flags=peer_flags,
        peer_creation=peer_creation,
        our_flags=our_flags,
        our_creation=our_creation,
    )


# ---------------------------------------------------------------------------
# Frame builders / parsers (pure functions, exposed for tests)
# ---------------------------------------------------------------------------


def build_send_name(name: str, flags: int, creation: int) -> bytes:
    """Build the v6 ``send_name`` frame body (no length prefix)."""
    name_bytes = name.encode("utf-8")
    return (
        bytes([TAG_SEND_NAME])
        + struct.pack(">Q", flags)
        + struct.pack(">I", creation)
        + struct.pack(">H", len(name_bytes))
        + name_bytes
    )


def parse_recv_status(frame: bytes) -> str:
    if not frame or frame[0] != TAG_RECV_STATUS:
        raise HandshakeError(
            f"expected recv_status tag {TAG_RECV_STATUS}, got {frame[:1].hex()}"
        )
    return frame[1:].decode("latin-1")


def parse_recv_challenge(frame: bytes) -> tuple[int, int, int, str]:
    """Returns (flags, challenge, creation, name)."""
    if not frame or frame[0] != TAG_RECV_CHALLENGE:
        raise HandshakeError(
            f"expected recv_challenge tag {TAG_RECV_CHALLENGE}, got {frame[:1].hex()}"
        )
    if len(frame) < 1 + 8 + 4 + 4 + 2:
        raise HandshakeError(f"recv_challenge frame truncated: {len(frame)} bytes")
    flags = struct.unpack(">Q", frame[1:9])[0]
    challenge = struct.unpack(">I", frame[9:13])[0]
    creation = struct.unpack(">I", frame[13:17])[0]
    nlen = struct.unpack(">H", frame[17:19])[0]
    if len(frame) < 19 + nlen:
        raise HandshakeError("recv_challenge name truncated")
    name = frame[19:19 + nlen].decode("utf-8")
    return flags, challenge, creation, name


def build_challenge_reply(challenge: int, digest: bytes) -> bytes:
    if len(digest) != 16:
        raise ValueError(f"digest must be 16 bytes, got {len(digest)}")
    return bytes([TAG_CHALLENGE_REPLY]) + struct.pack(">I", challenge) + digest


def parse_challenge_ack(frame: bytes) -> bytes:
    if not frame or frame[0] != TAG_CHALLENGE_ACK:
        raise HandshakeError(
            f"expected challenge_ack tag {TAG_CHALLENGE_ACK}, got {frame[:1].hex()}"
        )
    if len(frame) != 1 + 16:
        raise HandshakeError(f"challenge_ack frame wrong length: {len(frame)} bytes")
    return frame[1:17]


def cookie_digest(cookie: str, challenge: int) -> bytes:
    """MD5 of cookie concatenated with the base-10 string of the challenge.

    Per the OTP handshake spec: ``md5(Cookie ++ integer_to_list(Challenge))``.
    The challenge is treated as an unsigned 32-bit integer.
    """
    if challenge < 0:
        challenge = challenge & 0xFFFFFFFF
    payload = cookie.encode("latin-1") + str(challenge).encode("ascii")
    return hashlib.md5(payload).digest()


# ---------------------------------------------------------------------------
# Wire helpers (2-byte length-prefixed frames)
# ---------------------------------------------------------------------------


def _send_frame(sock: socket.socket, body: bytes, frame_size: int = 2) -> None:
    if frame_size == 2:
        if len(body) > 0xFFFF:
            raise HandshakeError(f"handshake frame too large for 2-byte length: {len(body)} bytes")
        sock.sendall(struct.pack(">H", len(body)) + body)
    else:
        sock.sendall(struct.pack(">I", len(body)) + body)


def _recv_frame(sock: socket.socket, frame_size: int = 2) -> bytes:
    hdr = _recv_exact(sock, frame_size)
    if frame_size == 2:
        (length,) = struct.unpack(">H", hdr)
    else:
        (length,) = struct.unpack(">I", hdr)
    if length == 0:
        return b""
    return _recv_exact(sock, length)


def _recv_exact(sock: socket.socket, n: int) -> bytes:
    chunks = []
    remaining = n
    while remaining > 0:
        chunk = sock.recv(remaining)
        if not chunk:
            raise HandshakeError(f"peer closed connection after {n - remaining}/{n} bytes")
        chunks.append(chunk)
        remaining -= len(chunk)
    return b"".join(chunks)


def _random_creation() -> int:
    """Non-zero uint32. The spec wants a value that survives node restarts."""
    n = struct.unpack(">I", os.urandom(4))[0]
    return n if n else 1


def _random_challenge() -> int:
    return struct.unpack(">I", os.urandom(4))[0]