progpib/progpib/discovery.py

"""Netfinder: the Prologix GPIB-ETHERNET controller's UDP discovery and
network-configuration protocol (broadcast, port 3040).

Struct layouts below are ported from `vendortools/nfutil.py`, which are
correct, but that file has a real Python 3 bug: `MkIdentifyReply` and
`MkAssignment` build single-byte struct fields with `chr(x)`, which produces
a `str`, when `struct.pack`'s `c` format needs a length-1 `bytes` object
(`bytes([x])`). `sampledata.txt` captures this exact call raising
`struct.error: char format requires a bytes object of length 1` against
real hardware, and the working fix used there. Only the pack/unpack pairs a
discovery *client* needs (Identify, Assignment, Test) are implemented here;
the rest (block write/verify/reboot/etc.) are firmware/bootloader tooling,
out of scope.
"""

import contextlib
import dataclasses
import enum
import random
import socket
import struct
import sys
import time
from typing import Callable, Iterator

from bits import Bytes

from .exceptions import DiscoveryError

NETFINDER_PORT = 3040
NF_MAGIC = 0x5A


class NfCommand(enum.IntEnum):
    IDENTIFY = 0
    IDENTIFY_REPLY = 1
    ASSIGNMENT = 2
    ASSIGNMENT_REPLY = 3
    FLASH_ERASE = 4
    FLASH_ERASE_REPLY = 5
    BLOCK_SIZE = 6
    BLOCK_SIZE_REPLY = 7
    BLOCK_WRITE = 8
    BLOCK_WRITE_REPLY = 9
    VERIFY = 10
    VERIFY_REPLY = 11
    REBOOT = 12
    SET_ETHERNET_ADDRESS = 13
    SET_ETHERNET_ADDRESS_REPLY = 14
    TEST = 15
    TEST_REPLY = 16


class NfResult(enum.IntEnum):
    SUCCESS = 0
    CRC_MISMATCH = 1
    INVALID_MEMORY_TYPE = 2
    INVALID_SIZE = 3
    INVALID_IP_TYPE = 4


class IpType(enum.IntEnum):
    DYNAMIC = 0
    STATIC = 1


_HEADER_FMT = "!2cH6s2x"
_IDENTIFY_REPLY_FMT = "!H6c4s4s4s4s4s4s32s"
_ASSIGNMENT_FMT = "!3xc4s4s4s32x"
_ASSIGNMENT_REPLY_FMT = "!c3x"
_TEST_REPLY_FMT = "!32s"

_HEADER_LEN = struct.calcsize(_HEADER_FMT)


@dataclasses.dataclass(frozen=True)
class DeviceInfo:
    """A Prologix GPIB-ETHERNET controller found (or targeted) via netfinder."""

    mac: bytes
    mode: int
    alert: int
    ip_type: IpType
    ip_address: str
    netmask: str
    gateway: str
    app_version: str
    boot_version: str
    hw_version: str
    name: str
    uptime_days: int
    uptime_hours: int
    uptime_minutes: int
    uptime_seconds: int
    host_ip: str | None = None  # which local interface discovered it

    @property
    def mac_str(self) -> str:
        return Bytes(self.mac).hex(sep=":", bytes_per_sep=1)


def _pack_header(command: NfCommand, seq: int, mac: bytes) -> bytes:
    return struct.pack(_HEADER_FMT, bytes([NF_MAGIC]), bytes([int(command)]), seq, mac)


def _unpack_header(msg: bytes) -> dict:
    magic, command, seq, mac = struct.unpack(_HEADER_FMT, msg)
    return {"magic": ord(magic), "command": ord(command), "seq": seq, "mac": mac}


def _pack_identify(seq: int) -> bytes:
    return _pack_header(NfCommand.IDENTIFY, seq, b"\xff" * 6)


def _unpack_identify_reply(msg: bytes) -> DeviceInfo:
    header = _unpack_header(msg[:_HEADER_LEN])
    (
        uptime_days,
        uptime_hrs,
        uptime_mins,
        uptime_secs,
        mode,
        alert,
        ip_type,
        ip_addr,
        netmask,
        gateway,
        app_ver,
        boot_ver,
        hw_ver,
        name,
    ) = struct.unpack(_IDENTIFY_REPLY_FMT, msg[_HEADER_LEN:])
    return DeviceInfo(
        mac=header["mac"],
        mode=ord(mode),
        alert=ord(alert),
        ip_type=IpType(ord(ip_type)),
        ip_address=socket.inet_ntoa(ip_addr),
        netmask=socket.inet_ntoa(netmask),
        gateway=socket.inet_ntoa(gateway),
        app_version=socket.inet_ntoa(app_ver),
        boot_version=socket.inet_ntoa(boot_ver),
        hw_version=socket.inet_ntoa(hw_ver),
        name=name.split(b"\x00", 1)[0].decode("ascii", errors="replace"),
        uptime_days=uptime_days,
        uptime_hours=ord(uptime_hrs),
        uptime_minutes=ord(uptime_mins),
        uptime_seconds=ord(uptime_secs),
    )


def _pack_assignment(
    seq: int,
    mac: bytes,
    ip_type: IpType,
    ip_address: str,
    netmask: str,
    gateway: str,
) -> bytes:
    return _pack_header(NfCommand.ASSIGNMENT, seq, mac) + struct.pack(
        _ASSIGNMENT_FMT,
        bytes([int(ip_type)]),  # fixed: was chr(ip_type) in vendortools/nfutil.py
        socket.inet_aton(ip_address),
        socket.inet_aton(netmask),
        socket.inet_aton(gateway),
    )


def _unpack_assignment_reply(msg: bytes) -> NfResult:
    (result,) = struct.unpack(_ASSIGNMENT_REPLY_FMT, msg[_HEADER_LEN:])
    return NfResult(ord(result))


def _pack_test(seq: int, mac: bytes) -> bytes:
    return _pack_header(NfCommand.TEST, seq, mac)


def _unpack_test_reply(msg: bytes) -> str:
    (raw,) = struct.unpack(_TEST_REPLY_FMT, msg[_HEADER_LEN:])
    return raw.split(b"\x00", 1)[0].decode("ascii", errors="replace")


def _validate_reply_header(
    reply: bytes, expected_command: NfCommand, seq: int, expected_len: int
) -> bool:
    if len(reply) != expected_len:
        return False
    header = _unpack_header(reply[:_HEADER_LEN])
    return (
        header["magic"] == NF_MAGIC
        and header["command"] == int(expected_command)
        and header["seq"] == seq
    )


def _local_ipv4_addresses() -> list[str]:
    """Local, non-loopback IPv4 addresses to broadcast netfinder requests from."""
    import platform

    if platform.system() in ("Windows", "Microsoft"):
        addresses = socket.gethostbyname_ex(socket.gethostname())[2]
    else:
        addresses = _local_ipv4_addresses_unix()
    return [ip for ip in addresses if ip != "127.0.0.1"]


def _local_ipv4_addresses_unix() -> list[str]:
    import array
    import fcntl

    siocgifconf = 0x8912
    struct_size = 40 if sys.maxsize > 2**32 else 32
    max_interfaces = 128
    inbytes = max_interfaces * struct_size
    ifreq = array.array("B", b"\0" * inbytes)
    ifconf = struct.pack("iL", inbytes, ifreq.buffer_info()[0])

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        ifconf = fcntl.ioctl(sock.fileno(), siocgifconf, ifconf)
    finally:
        sock.close()

    outbytes = struct.unpack("iL", ifconf)[0]
    return [
        socket.inet_ntoa(ifreq[i + 20 : i + 24])
        for i in range(0, outbytes, struct_size)
    ]


@contextlib.contextmanager
def _netfinder_sockets(interface: str):
    """A bound (send, recv) UDP socket pair for broadcasting on `interface`."""
    send_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    recv_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        send_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        send_sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        send_sock.bind((interface, 0))
        port = send_sock.getsockname()[1]

        recv_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        recv_sock.settimeout(0.1)
        recv_sock.bind(("", port))

        yield send_sock, recv_sock
    finally:
        send_sock.close()
        recv_sock.close()


def _broadcast_request(
    send_sock: socket.socket, recv_sock: socket.socket, packet: bytes, timeout: float
) -> Iterator[bytes]:
    """Broadcast `packet` and yield every reply datagram received within `timeout`."""
    send_sock.sendto(packet, ("<broadcast>", NETFINDER_PORT))
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        try:
            yield recv_sock.recv(256)
        except socket.timeout:
            continue


def discover(
    *,
    timeout: float = 0.5,
    attempts: int = 2,
    interfaces: list[str] | None = None,
    on_progress: Callable[[str], None] | None = None,
) -> list[DeviceInfo]:
    """Broadcast-discover Prologix GPIB-ETHERNET controllers on the LAN.

    Sends an Identify request on each local IPv4 interface (or `interfaces`,
    if given) and collects IdentifyReply packets for `timeout` seconds,
    `attempts` times per interface. Devices found on more than one interface
    are deduplicated by MAC address (last interface wins for `host_ip`).
    """
    if interfaces is None:
        interfaces = _local_ipv4_addresses()
    if not interfaces:
        raise DiscoveryError("no local IPv4 interfaces found to broadcast from")

    reply_len = _HEADER_LEN + struct.calcsize(_IDENTIFY_REPLY_FMT)
    devices: dict[bytes, DeviceInfo] = {}

    for interface in interfaces:
        if on_progress:
            on_progress(f"scanning via {interface}")
        with _netfinder_sockets(interface) as (send_sock, recv_sock):
            for _ in range(attempts):
                seq = random.randint(1, 65535)
                packet = _pack_identify(seq)
                for reply in _broadcast_request(send_sock, recv_sock, packet, timeout):
                    if not _validate_reply_header(
                        reply, NfCommand.IDENTIFY_REPLY, seq, reply_len
                    ):
                        continue
                    try:
                        info = _unpack_identify_reply(reply)
                    except struct.error:
                        continue
                    info = dataclasses.replace(info, host_ip=interface)
                    devices[info.mac] = info
                    if on_progress:
                        on_progress(f"found {info.mac_str} at {info.ip_address}")

    return list(devices.values())


def identify(
    mac: bytes,
    *,
    timeout: float = 2.0,
    attempts: int = 2,
    interfaces: list[str] | None = None,
) -> DeviceInfo | None:
    """Discover and return the device with the given MAC, or `None`."""
    for info in discover(timeout=timeout, attempts=attempts, interfaces=interfaces):
        if info.mac == mac:
            return info
    return None


def set_network_config(
    mac: bytes,
    *,
    ip_type: IpType,
    ip_address: str = "0.0.0.0",
    netmask: str = "0.0.0.0",
    gateway: str = "0.0.0.0",
    host_ip: str | None = None,
    timeout: float = 0.5,
    attempts: int = 10,
) -> bool:
    """Set the network configuration of the device with the given MAC.

    `host_ip` should be the local interface that can reach the device (e.g.
    `DeviceInfo.host_ip` from a prior `discover()`/`identify()` call); if
    omitted, every local interface is tried in turn. Returns whether the
    device acknowledged success. Raises `DiscoveryError` if no reply is
    received at all after exhausting `attempts` (the vendor tool instead
    returns a falsy empty dict, which is easy to silently mishandle).
    """
    interfaces = [host_ip] if host_ip else _local_ipv4_addresses()
    if not interfaces:
        raise DiscoveryError("no local IPv4 interfaces found to broadcast from")

    reply_len = _HEADER_LEN + struct.calcsize(_ASSIGNMENT_REPLY_FMT)

    for interface in interfaces:
        with _netfinder_sockets(interface) as (send_sock, recv_sock):
            for _ in range(attempts):
                seq = random.randint(1, 65535)
                packet = _pack_assignment(seq, mac, ip_type, ip_address, netmask, gateway)
                for reply in _broadcast_request(send_sock, recv_sock, packet, timeout):
                    if not _validate_reply_header(
                        reply, NfCommand.ASSIGNMENT_REPLY, seq, reply_len
                    ):
                        continue
                    header = _unpack_header(reply[:_HEADER_LEN])
                    if header["mac"] != mac:
                        continue
                    return _unpack_assignment_reply(reply) == NfResult.SUCCESS

    raise DiscoveryError(f"no response configuring {mac.hex()}")


def test(mac: bytes, *, host_ip: str | None = None, timeout: float = 0.5) -> str | None:
    """Send a Test request to the device with the given MAC; a simple liveness
    check, e.g. after `set_network_config`. Returns the reply string, or
    `None` if there was no reply.
    """
    interfaces = [host_ip] if host_ip else _local_ipv4_addresses()
    reply_len = _HEADER_LEN + struct.calcsize(_TEST_REPLY_FMT)

    for interface in interfaces:
        with _netfinder_sockets(interface) as (send_sock, recv_sock):
            seq = random.randint(1, 65535)
            packet = _pack_test(seq, mac)
            for reply in _broadcast_request(send_sock, recv_sock, packet, timeout):
                if not _validate_reply_header(reply, NfCommand.TEST_REPLY, seq, reply_len):
                    continue
                header = _unpack_header(reply[:_HEADER_LEN])
                if header["mac"] != mac:
                    continue
                return _unpack_test_reply(reply)

    return None