"""Minimal RustDesk rendezvous protocol encoder/probe (stdlib only, zero deps).

hbbs speaks `RendezvousMessage` protobuf over UDP :21116 (one message per datagram, no
framing). The wire format is hand-encoded here from the upstream proto
(rustdesk/hbb_common protos/rendezvous.proto — field numbers are IDENTICAL between the
1.1.12-vendored and the current shared-proto era, verified against both tags):

    message RendezvousMessage {
      oneof union {
        RegisterPeer        register_peer         = 6;
        RegisterPeerResponse register_peer_response = 7;
        PunchHoleRequest    punch_hole_request    = 8;
        PunchHoleResponse   punch_hole_response   = 11;
        ...
      }
    }
    message RegisterPeer     { string id = 1; int32 serial = 2; }
    message PunchHoleRequest { string id = 1; ... }

This is the surface a real RustDesk client uses to register and request connections; it is
the app's own interface (STYLE.md §1), exercised from the cc-ci host against the
host-published :21116 (the stock compose publishes 21115-21119 mode:host).
"""

from __future__ import annotations

import socket
import time

RENDEZVOUS_PORT = 21116


def varint(n: int) -> bytes:
    out = bytearray()
    while n > 0x7F:
        out.append((n & 0x7F) | 0x80)
        n >>= 7
    out.append(n)
    return bytes(out)


def field_str(field: int, s: str) -> bytes:
    b = s.encode()
    return varint((field << 3) | 2) + varint(len(b)) + b


def field_varint(field: int, n: int) -> bytes:
    return varint((field << 3) | 0) + varint(n)


def wrap(union_field: int, inner: bytes) -> bytes:
    """Wrap an inner message as the given oneof field of RendezvousMessage."""
    return varint((union_field << 3) | 2) + varint(len(inner)) + inner


def register_peer(peer_id: str) -> bytes:
    """RendezvousMessage{ register_peer { id, serial=0 } } — a plain peer registration."""
    return wrap(6, field_str(1, peer_id) + field_varint(2, 0))


def punch_hole_request(target_id: str) -> bytes:
    """RendezvousMessage{ punch_hole_request { id } } — an UNAUTHENTICATED connection request
    for `target_id` (no licence_key/token — exactly the spoofable request shape from the
    CVE-2026-30784 / GH #670 amplification report)."""
    return wrap(8, field_str(1, target_id))


def parse_union_field(data: bytes) -> int | None:
    """Which RendezvousMessage oneof field the server answered with (None if unparseable)."""
    if not data:
        return None
    key = data[0]
    field, wire = key >> 3, key & 0x7
    return field if wire == 2 else field  # field number regardless of wire type


def udp_probe(
    msg: bytes, host: str = "127.0.0.1", port: int = RENDEZVOUS_PORT, wait: float = 2.0
) -> bytes | None:
    """Send one datagram, return the server's reply (or None after `wait` seconds of silence)."""
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.settimeout(wait)
    try:
        s.sendto(msg, (host, port))
        try:
            data, _ = s.recvfrom(65535)
            return data
        except TimeoutError:
            return None
    finally:
        s.close()


def retry_register_peer(peer_id: str, attempts: int = 6, interval: float = 5.0) -> bytes | None:
    """Poll register_peer until the server answers (post-deploy settle), or give up.

    Returns the raw register_peer_response datagram."""
    deadline = time.time() + attempts * interval
    last: bytes | None = None
    while time.time() < deadline:
        last = udp_probe(register_peer(peer_id))
        if last is not None:
            return last
        time.sleep(interval)
    return last