claudePySnake/src/shared/discovery.py

"""Server discovery utilities using UDP multicast."""

import json
import socket
import struct
import time
from dataclasses import dataclass
from typing import Optional

from .constants import MULTICAST_GROUP, MULTICAST_PORT


@dataclass
class ServerInfo:
    """Information about a discovered server."""
    host: str
    port: int
    server_name: str
    players_count: int
    last_seen: float = 0.0

    def to_dict(self) -> dict:
        """Convert to dictionary for serialization."""
        return {
            "host": self.host,
            "port": self.port,
            "server_name": self.server_name,
            "players_count": self.players_count,
        }

    @classmethod
    def from_dict(cls, data: dict, host: str = None) -> "ServerInfo":
        """Create ServerInfo from dictionary.

        Args:
            data: Dictionary with server information
            host: Optional override for host address
        """
        return cls(
            host=host or data["host"],
            port=data["port"],
            server_name=data["server_name"],
            players_count=data["players_count"],
            last_seen=time.time(),
        )


class DiscoveryMessage:
    """Discovery protocol messages."""

    DISCOVER = "DISCOVER"
    SERVER_ANNOUNCE = "SERVER_ANNOUNCE"

    @staticmethod
    def create_discover() -> bytes:
        """Create a DISCOVER message.

        Returns:
            Encoded message bytes
        """
        msg = {"type": DiscoveryMessage.DISCOVER}
        return json.dumps(msg).encode("utf-8")

    @staticmethod
    def create_announce(server_info: ServerInfo) -> bytes:
        """Create a SERVER_ANNOUNCE message.

        Args:
            server_info: Server information to announce

        Returns:
            Encoded message bytes
        """
        msg = {
            "type": DiscoveryMessage.SERVER_ANNOUNCE,
            "data": server_info.to_dict(),
        }
        return json.dumps(msg).encode("utf-8")

    @staticmethod
    def parse(data: bytes) -> tuple[str, Optional[dict]]:
        """Parse a discovery message.

        Args:
            data: Raw message bytes

        Returns:
            Tuple of (message_type, data_dict)
        """
        try:
            msg = json.loads(data.decode("utf-8"))
            msg_type = msg.get("type")
            msg_data = msg.get("data")
            return msg_type, msg_data
        except Exception:
            return None, None


def create_multicast_socket(bind: bool = False) -> socket.socket:
    """Create a UDP socket configured for multicast.

    Args:
        bind: If True, bind to multicast group (for receiving)

    Returns:
        Configured socket
    """
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)

    # Allow multiple sockets to bind to the same port
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

    if bind:
        # Bind to the multicast port
        sock.bind(("", MULTICAST_PORT))

        # Join the multicast group
        mreq = struct.pack("4sl", socket.inet_aton(MULTICAST_GROUP), socket.INADDR_ANY)
        sock.setsockopt(socket.IPPROTO_IP, socket.IP_ADD_MEMBERSHIP, mreq)

    # Set multicast TTL (time-to-live) for sending
    sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, 2)

    return sock


def get_local_ip() -> str:
    """Get the local IP address for LAN communication.

    Returns:
        Local IP address as string
    """
    try:
        # Create a UDP socket and connect to a public DNS server
        # This doesn't actually send data, just determines routing
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        local_ip = s.getsockname()[0]
        s.close()
        return local_ip
    except Exception:
        return "127.0.0.1"