RTSP-Streamer/internal/protect/protect.go

// Package protect is a minimal client for the UniFi Protect local API.
//
// UniFi Protect has no officially documented public API, but the local
// endpoints used here (/api/auth/login and /proxy/protect/api/bootstrap) are
// stable and are the same ones the Home Assistant integration and the
// uiprotect/pyunifiprotect libraries rely on. We authenticate as a local
// Protect user, read the bootstrap document, and construct RTSPS URLs from
// each camera's enabled channel alias.
package protect

import (
	"bytes"
	"context"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/cookiejar"
	"strings"
	"time"
)

// Client talks to a single UniFi Protect controller.
type Client struct {
	host     string
	rtspPort int
	http     *http.Client
	csrf     string
}

// Camera is a discovered Protect camera with its resolved stream URLs.
type Camera struct {
	ID       string
	Name     string
	State    string // "CONNECTED", "DISCONNECTED", ...
	Channels []Channel
}

// Channel is one encoding profile (high/medium/low) on a camera.
type Channel struct {
	ID          int
	Name        string // "High", "Medium", "Low"
	Width       int
	Height      int
	RTSPEnabled bool
	RTSPAlias   string
}

// New builds a client. verifyTLS=false accepts the console's self-signed cert.
func New(host string, rtspPort int, verifyTLS bool) (*Client, error) {
	jar, err := cookiejar.New(nil)
	if err != nil {
		return nil, err
	}
	if rtspPort == 0 {
		rtspPort = 7441
	}
	return &Client{
		host:     host,
		rtspPort: rtspPort,
		http: &http.Client{
			Timeout: 20 * time.Second,
			Jar:     jar,
			Transport: &http.Transport{
				TLSClientConfig: &tls.Config{InsecureSkipVerify: !verifyTLS}, //nolint:gosec // self-signed console cert
			},
		},
	}, nil
}

// Login authenticates and captures the session cookie + CSRF token. UniFi OS
// returns the CSRF token in a response header on successful login.
func (c *Client) Login(ctx context.Context, username, password string) error {
	body, _ := json.Marshal(map[string]any{
		"username": username,
		"password": password,
		"rememberMe": true,
	})
	url := fmt.Sprintf("https://%s/api/auth/login", c.host)
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := c.http.Do(req)
	if err != nil {
		return fmt.Errorf("connecting to controller %s: %w", c.host, err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		snippet, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
		return fmt.Errorf("login failed (%s): %s", resp.Status, strings.TrimSpace(string(snippet)))
	}
	// UniFi OS exposes the CSRF token via header; capture whichever casing.
	if tok := resp.Header.Get("X-CSRF-Token"); tok != "" {
		c.csrf = tok
	} else if tok := resp.Header.Get("X-Updated-CSRF-Token"); tok != "" {
		c.csrf = tok
	}
	return nil
}

// bootstrap is the subset of the Protect bootstrap document we care about.
type bootstrap struct {
	Cameras []struct {
		ID              string `json:"id"`
		Name            string `json:"name"`
		State           string `json:"state"`
		IsRTSPEnabled   bool   `json:"isRtspEnabled"`
		ChannelsWrapper []struct {
			ID          int    `json:"id"`
			Name        string `json:"name"`
			Width       int    `json:"width"`
			Height      int    `json:"height"`
			IsRTSPEnabled bool `json:"isRtspEnabled"`
			RTSPAlias   string `json:"rtspAlias"`
		} `json:"channels"`
	} `json:"cameras"`
}

// Cameras fetches the bootstrap document and returns the camera list.
func (c *Client) Cameras(ctx context.Context) ([]Camera, error) {
	url := fmt.Sprintf("https://%s/proxy/protect/api/bootstrap", c.host)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	if c.csrf != "" {
		req.Header.Set("X-CSRF-Token", c.csrf)
	}
	resp, err := c.http.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		snippet, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
		return nil, fmt.Errorf("bootstrap failed (%s): %s", resp.Status, strings.TrimSpace(string(snippet)))
	}
	var b bootstrap
	if err := json.NewDecoder(resp.Body).Decode(&b); err != nil {
		return nil, fmt.Errorf("decoding bootstrap: %w", err)
	}

	out := make([]Camera, 0, len(b.Cameras))
	for _, bc := range b.Cameras {
		cam := Camera{ID: bc.ID, Name: bc.Name, State: bc.State}
		for _, ch := range bc.ChannelsWrapper {
			cam.Channels = append(cam.Channels, Channel{
				ID:          ch.ID,
				Name:        ch.Name,
				Width:       ch.Width,
				Height:      ch.Height,
				RTSPEnabled: ch.IsRTSPEnabled,
				RTSPAlias:   ch.RTSPAlias,
			})
		}
		out = append(out, cam)
	}
	return out, nil
}

// StreamURL builds the RTSPS URL for a channel alias on this controller.
// enableSrtp is required by Protect's RTSPS endpoint.
func (c *Client) StreamURL(alias string) string {
	return fmt.Sprintf("rtsps://%s:%d/%s?enableSrtp", c.host, c.rtspPort, alias)
}

// BestEnabledChannel returns the highest-resolution channel that has RTSP
// enabled, or nil if none are enabled. Preferring the top channel gives the
// sharpest wall tile; callers can pick a lower one for dense grids.
func (cam Camera) BestEnabledChannel() *Channel {
	var best *Channel
	for i := range cam.Channels {
		ch := &cam.Channels[i]
		if !ch.RTSPEnabled || ch.RTSPAlias == "" {
			continue
		}
		if best == nil || ch.Width*ch.Height > best.Width*best.Height {
			best = ch
		}
	}
	return best
}

// LowestEnabledChannel returns the lowest-resolution enabled channel, useful
// for dense grids where a substream is plenty.
func (cam Camera) LowestEnabledChannel() *Channel {
	var low *Channel
	for i := range cam.Channels {
		ch := &cam.Channels[i]
		if !ch.RTSPEnabled || ch.RTSPAlias == "" {
			continue
		}
		if low == nil || ch.Width*ch.Height < low.Width*low.Height {
			low = ch
		}
	}
	return low
}

// ChannelByPreference picks a channel to enable RTSP on, by preference
// "high" | "medium" | "low". It matches on the channel name first (Protect
// labels them "High"/"Medium"/"Low") and falls back to resolution ranking so
// it still works on cameras with unusual channel names.
func (cam Camera) ChannelByPreference(pref string) *Channel {
	pref = strings.ToLower(strings.TrimSpace(pref))
	for i := range cam.Channels {
		if strings.ToLower(cam.Channels[i].Name) == pref {
			return &cam.Channels[i]
		}
	}
	if len(cam.Channels) == 0 {
		return nil
	}
	var pick *Channel
	for i := range cam.Channels {
		ch := &cam.Channels[i]
		if pick == nil {
			pick = ch
			continue
		}
		switch pref {
		case "low":
			if ch.Width*ch.Height < pick.Width*pick.Height {
				pick = ch
			}
		default: // treat anything else as "highest resolution"
			if ch.Width*ch.Height > pick.Width*pick.Height {
				pick = ch
			}
		}
	}
	return pick
}

// do issues an authenticated request to a Protect API path (e.g.
// "/proxy/protect/api/cameras/<id>"), attaching the session cookie (via the
// client's jar) and the CSRF token. The caller closes the response body.
func (c *Client) do(ctx context.Context, method, path string, body any) (*http.Response, error) {
	var rdr io.Reader
	if body != nil {
		raw, err := json.Marshal(body)
		if err != nil {
			return nil, err
		}
		rdr = bytes.NewReader(raw)
	}
	url := fmt.Sprintf("https://%s%s", c.host, path)
	req, err := http.NewRequestWithContext(ctx, method, url, rdr)
	if err != nil {
		return nil, err
	}
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	if c.csrf != "" {
		req.Header.Set("X-CSRF-Token", c.csrf)
	}
	return c.http.Do(req)
}

// EnableRTSP turns on RTSP for one channel of a camera and returns the newly
// assigned rtspAlias. It reads the camera's current channels as raw JSON and
// flips only isRtspEnabled on the target channel before PATCHing them back, so
// no other encoder settings (bitrate, fps, ...) are disturbed.
func (c *Client) EnableRTSP(ctx context.Context, cameraID string, channelID int) (string, error) {
	getResp, err := c.do(ctx, http.MethodGet, "/proxy/protect/api/cameras/"+cameraID, nil)
	if err != nil {
		return "", err
	}
	defer getResp.Body.Close()
	if getResp.StatusCode != http.StatusOK {
		snippet, _ := io.ReadAll(io.LimitReader(getResp.Body, 512))
		return "", fmt.Errorf("fetching camera %s (%s): %s", cameraID, getResp.Status, strings.TrimSpace(string(snippet)))
	}
	// Keep channels as raw maps to preserve every field we don't touch.
	var cam struct {
		Channels []map[string]any `json:"channels"`
	}
	if err := json.NewDecoder(getResp.Body).Decode(&cam); err != nil {
		return "", fmt.Errorf("decoding camera %s: %w", cameraID, err)
	}
	found := false
	for _, ch := range cam.Channels {
		id, ok := ch["id"].(float64)
		if ok && int(id) == channelID {
			ch["isRtspEnabled"] = true
			found = true
		}
	}
	if !found {
		return "", fmt.Errorf("camera %s has no channel %d", cameraID, channelID)
	}

	patchResp, err := c.do(ctx, http.MethodPatch, "/proxy/protect/api/cameras/"+cameraID,
		map[string]any{"channels": cam.Channels})
	if err != nil {
		return "", err
	}
	defer patchResp.Body.Close()
	if patchResp.StatusCode != http.StatusOK {
		snippet, _ := io.ReadAll(io.LimitReader(patchResp.Body, 512))
		return "", fmt.Errorf("enabling RTSP on camera %s (%s): %s", cameraID, patchResp.Status, strings.TrimSpace(string(snippet)))
	}
	var updated struct {
		Channels []struct {
			ID        int    `json:"id"`
			RTSPAlias string `json:"rtspAlias"`
		} `json:"channels"`
	}
	if err := json.NewDecoder(patchResp.Body).Decode(&updated); err != nil {
		return "", fmt.Errorf("decoding RTSP-enable response: %w", err)
	}
	for _, ch := range updated.Channels {
		if ch.ID == channelID {
			if ch.RTSPAlias == "" {
				return "", fmt.Errorf("camera %s channel %d still has no rtspAlias after enabling", cameraID, channelID)
			}
			return ch.RTSPAlias, nil
		}
	}
	return "", fmt.Errorf("channel %d missing from RTSP-enable response", channelID)
}

// EnableMissing enables RTSP on the preferred channel for every camera that
// currently has no RTSP-enabled channel, mutating cams in place so their
// BestEnabledChannel/LowestEnabledChannel become usable. It returns the names
// it enabled and, per camera name, any error encountered.
func (c *Client) EnableMissing(ctx context.Context, cams []Camera, pref string) (enabled []string, failed map[string]error) {
	failed = map[string]error{}
	for i := range cams {
		cam := &cams[i]
		if cam.BestEnabledChannel() != nil {
			continue
		}
		target := cam.ChannelByPreference(pref)
		if target == nil {
			failed[cam.Name] = fmt.Errorf("no channels to enable")
			continue
		}
		alias, err := c.EnableRTSP(ctx, cam.ID, target.ID)
		if err != nil {
			failed[cam.Name] = err
			continue
		}
		// Reflect the change in the in-memory model.
		for j := range cam.Channels {
			if cam.Channels[j].ID == target.ID {
				cam.Channels[j].RTSPEnabled = true
				cam.Channels[j].RTSPAlias = alias
			}
		}
		enabled = append(enabled, cam.Name)
	}
	return enabled, failed
}