celebright/celebright_controller.py

#!/usr/bin/env python3
"""
Celebright Holiday Lights Controller
Based on packet capture analysis of the Celebright protocol.
"""
import questionary
from questionary import Style, Choice
import socket
import json
import asyncio
import logging
import websockets
from typing import Optional, Dict, List, Self
import netifaces
from dataclasses import dataclass

DISCOVERY_PORT = 49999
DISCOVERY_TIMEOUT = 5.0
DISCOVERY_MESSAGE = b'App Broadcast Message'
WEBSOCKET_ENDPOINT = "/ws"

log = logging.getLogger("CelebrightController")


@dataclass(frozen=True)
class CelebrightPreset:
    preset_uuid: str
    preset_name: str
    preset_desc: str
    preset_mode: int
    brightness: int
    web_id: int
    md5: str
    controller: "CelebrightController"

    async def select(self) -> None:
        """Select/load this preset on the controller."""
        await self.controller.load_preset(self)


@dataclass
class CelebrightController:
    """Controller for Celebright holiday lights system."""
    address: str
    model_name: str
    device_id: str
    firmware: str
    hardware: int
    num_leds: int
    location_id: int
    timezoneId: int
    led_strings: list[int]
    isRGBW: bool
    colorOrder: str
    RGBW_type: str
    bulbType: str
    _ws: Optional[websockets.ClientConnection] = None

    @classmethod
    def _discover_broadcast_addresses(cls) -> list[str]:
        """ Get list of broadcast addresses for all network interfaces. """
        broadcast_addrs = list()
        for iface in netifaces.interfaces():
            addrs = netifaces.ifaddresses(iface)
            if netifaces.AF_INET in addrs:
                for addr in addrs[netifaces.AF_INET]:
                    if 'broadcast' in addr:
                        broadcast_addrs.append(addr['broadcast'])
        broadcast_addrs = sorted(set(broadcast_addrs))
        log.debug("Discovered System Broadcast addresses: %s", broadcast_addrs)
        return broadcast_addrs

    @classmethod
    def _discover_controller(cls, bcast: str) -> Self | None:
        """
        Discover Celebright controller on the network by sending UDP broadcast.

        Args:
            broadcast_addr: Broadcast address to use (default: 255.255.255.255)

        Returns:
            A list of dictionary containing device information
        """

        log.debug("Discovering Celebright controller on: %s", bcast)

        # Create socket for sending broadcast
        soc = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        soc.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
        soc.bind(('', DISCOVERY_PORT))

        # Send broadcast discovery request ("App Broadcast Message" payload)
        soc.sendto(DISCOVERY_MESSAGE, (bcast, DISCOVERY_PORT))
        log.debug(f"Sent Discovery Broadcast: {bcast}:{DISCOVERY_PORT}")

        try:
            # Wait for response from controller
            while True:
                try:
                    data, addr = soc.recvfrom(4096)
                    log.debug(f"Received response from {addr}: {data}")
                    if data:
                        data = json.loads(data.decode('utf-8'))
                        if "model_name" in data:
                            log.info("Discovered Celebright Controller at %s",
                                     addr[0])
                            return cls(address=addr[0], **data)
                except json.JSONDecodeError:
                    log.debug("Failed to decode JSON from response: %s", data)

        except socket.timeout:
            log.debug(f"Controller Timeout on %s", bcast)
        except Exception as e:
            log.warning(f"Unhandled Discovery Error: {e}")

        finally:
            soc.close()

    @classmethod
    def _discover_controllers(cls) -> List[Self]:
        """
        Discover Celebright controllers on the network by sending UDP broadcast.

        Returns:
            A list of dictionary containing device information
        """

        controllers = list()
        for address in cls._discover_broadcast_addresses():
            result = cls._discover_controller(address)
            if result:
                controllers.append(result)

        return controllers

    @classmethod
    def from_discovery(cls) -> list[Self]:
        """
        Factory method to create CelebrightController instances by
            discovering controllers on the network.
        """
        return cls._discover_controllers()

    async def connect(self) -> None:
        """
        Establish WebSocket connection to the controller.
        """
        ws_url = f"ws://{self.address}{WEBSOCKET_ENDPOINT}"
        log.debug(f"establishing websocket to {ws_url}...")

        self._ws = await websockets.connect(ws_url)
        log.info("websocket connection established: %s", ws_url)

    async def send_command(self, topic: str, message: Dict = None) -> None:
        """
        Send a command to the controller via WebSocket.

        Args:
            topic: Command topic (e.g., "getPresetsAndEventsPaginated")
            message: Command parameters dictionary (default: {})
        """
        if not self._ws:
            await self.connect()

        if message is None:
            message = {}

        command = {"topic": topic, "message": message}

        log.debug(f"sending to '%s': %s", topic, message)
        await self._ws.send(json.dumps(command))

    async def receive_response(self) -> Optional[Dict]:
        """
        Receive a response from the controller.
        """
        if not self._ws:
            await self.connect()

        response = await self._ws.recv()
        data = json.loads(response)
        log.debug(f"Received Response topic: {data.get('topic')}")
        return data

    async def get_presets(self) -> list[CelebrightPreset]:
        """
        Get list of available scene presets.

        Returns:
            List of preset dictionaries, or None if error
        """
        log.debug("Requesting preset list...")
        await self.send_command("getPresetsAndEventsPaginated")
        presets = list()
        response = await self.receive_response()
        if response and response.get('topic') == 'presetsPage':
            for preset in response.get('message', {}).get('presets', []):
                presets.append(CelebrightPreset(**preset, controller=self))
        return presets

    async def load_preset(self, preset: CelebrightPreset) -> None:
        """
        Load/activate a specific scene preset.

        Args:
            preset_uuid: UUID of the preset to load
        """
        log.info("Loading preset: %s", preset.preset_name)
        await self.send_command(
            "loadPreset",
            {"presetUuid": preset.preset_uuid}
        )
        log.info("Setting 12-hour sleep timer...")
        await self.send_command(
            "setSleepTimer",
            {"timerMinutes": 720}
        )

    async def turn_off(self) -> None:
        """Turn off the lights and disable schedule."""
        log.info("Turning off lights and disabling schedule...")
        await self.send_command("setTurnOffAndDisableSchedule")

    async def disconnect(self) -> None:
        """Close the WebSocket connection."""
        if self._ws:
            log.debug("Closing WebSocket connection...")
            await self._ws.close()
            self._ws = None
            log.debug("WebSocket connection closed.")

    async def interact(self):
        """Interactive menu for selecting presets or turning off lights."""
        custom_style = Style([('selected', 'fg:ansicyan bold')])
        presets = await self.get_presets()
        choices = [
            Choice(title=preset.preset_name, value=preset)
            for preset in presets
        ]
        choices.append(
            Choice(title="Turn Off Lights", value="__TURN_OFF__")
        )
        choices.append(
            Choice(title="Exit", value="__EXIT__")
        )

        while True:
            selected = await questionary.select(
                "Select a preset or action:",
                choices=choices,
                style=custom_style,
                use_indicator=True,
                use_shortcuts=False
            ).ask_async()

            if selected is None:
                break
            if selected == "__TURN_OFF__":
                await self.turn_off()
                return
            elif selected == "__EXIT__":
                return
            else:
                preset = selected
                await preset.select()


async def main():
    """Main function demonstrating the Celebright controller usage."""
    controllers = CelebrightController.from_discovery()
    controller = controllers[0]
    await controller.interact()


if __name__ == "__main__":
    asyncio.run(main())