wizlight.go

// Simple Golang script for controlling Wiz light bulbs.
// It uses UDP discovery and makes it blink randomly.

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"math"
	"math/rand"
	"net"
	"time"
)

type Params struct {
	State   bool  `json:"state"`
	SceneId uint8 `json:"sceneId,omitempty"` // TODO: until 1000
	R       uint8 `json:"r,omitempty"`
	G       uint8 `json:"g,omitempty"`
	B       uint8 `json:"b,omitempty"`
	C       uint8 `json:"c,omitempty"`
	W       uint8 `json:"w,omitempty"`
	Dimming uint8 `json:"dimming,omitempty"` // TODO: limit 100
	Success bool  `json:"success,omitempty"` // TODO: to Result
}

type Packet struct {
	Method string  `json:"method,omitempty"`
	Params *Params `json:"params,omitempty"`
	Env    string  `json:"env,omitempty"`
	Result *Params `json:"result,omitempty"`
}

var (
	count uint
	sleep time.Duration
)

func init() {
	flag.UintVar(&count, "count", 10, "")
	flag.DurationVar(&sleep, "sleep", time.Second, "")
}

func main() {
	flag.Parse()
	log.SetFlags(0)
	rand.Seed(time.Now().UnixNano())

	conn := must(net.ListenPacket("udp", ":0"))
	defer conn.Close()
	dst := must(net.ResolveUDPAddr("udp", "255.255.255.255:38899"))
	must(conn.WriteTo(must(json.Marshal(&Packet{Method: "getPilot"})), dst))

	cur := make(map[string]*Params)
	buf := make([]byte, math.MaxUint16)
	ban := make(map[string]chan struct{})
	for {
		n, addr, err := conn.ReadFrom(buf)
		if err != nil {
			log.Println("read:", err)
			continue
		}
		var res *Packet
		err = json.Unmarshal(buf[:n], &res)
		if err != nil {
			log.Println("load:", err)
			continue
		}
		ip := addr.(*net.UDPAddr).IP.String()
		switch res.Method {
		case "getPilot":
			fmt.Println("Found", ip)
			cur[ip] = res.Result
			ban[ip] = make(chan struct{})
			go func() {
				for i := uint(0); i < count; {
					<-ban[ip]
					r, g, b, err := roll(conn, addr)
					if err != nil {
						log.Println("roll:", err)
						return
					}
					fmt.Printf("Set %s to #%02x%02x%02x\n", ip, r, g, b)
					i++
					time.Sleep(sleep)
				}
				for err = send(conn, addr, cur[ip]); err != nil; {
					log.Println("send:", err)
				}
				fmt.Println("Reset", ip)
				delete(cur, ip)
			}()
			fallthrough
		case "setState":
			if res.Method != "getPilot" && res.Result != nil && !res.Result.Success {
				log.Println("recv:", res.Result)
				continue
			}
			go func() {
				ban[ip] <- struct{}{}
			}()
		}
	}
}

func roll(conn net.PacketConn, addr net.Addr) (r, g, b uint8, err error) {
	r = uint8(rand.Intn(256))
	g = uint8(rand.Intn(256))
	b = uint8(rand.Intn(256))
	err = send(conn, addr, &Params{R: r, G: g, B: b, State: true})
	return
}

func send(conn net.PacketConn, addr net.Addr, params *Params) (err error) {
	data := must(json.Marshal(&Packet{
		Method: "setState",
		Params: params,
	}))
	_, err = conn.WriteTo(data, addr)
	return
}

func must[T any](obj T, err error) T {
	if err != nil {
		panic(err)
	}
	return obj
}