ubiquiti-clients/pkg/edgeos/client.go

/*
Package edgeos provides a client for interacting with Ubiquiti EdgeOS devices
via their REST API. It supports authentication, token management, and
retrieval of system, interface, VLAN, and discovery information from
one or more devices.
*/
package edgeos

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

// Client handles communication with EdgeOS devices.
type Client struct {
	mu      sync.RWMutex
	devices map[string]*deviceClient
}

type deviceClient struct {
	config Config
	client *http.Client
	token  string
	mu     sync.Mutex
}

func newDeviceClient(cfg Config) *deviceClient {
	// Ensure scheme is set
	if cfg.Scheme == "" {
		cfg.Scheme = "https"
	}

	var tr http.RoundTripper
	if cfg.Transport != nil {
		tr = cfg.Transport
	} else {
		defaultTr := http.DefaultTransport.(*http.Transport).Clone()
		if defaultTr.TLSClientConfig == nil {
			defaultTr.TLSClientConfig = &tls.Config{}
		}
		defaultTr.TLSClientConfig.InsecureSkipVerify = cfg.Insecure
		tr = defaultTr
	}

	client := &http.Client{
		Transport: tr,
		Timeout:   cfg.Timeout,
	}

	return &deviceClient{
		config: cfg,
		client: client,
	}
}

// MustNew creates a new Client with the given configurations.
// It panics if a configuration is invalid (though currently we just accept all).
func MustNew(ctx context.Context, configs []Config) *Client {
	devices := make(map[string]*deviceClient)

	for _, cfg := range configs {
		devices[cfg.Host] = newDeviceClient(cfg)
	}

	return &Client{
		devices: devices,
	}
}

// Add adds a new device to the client.
// It returns an error if a device with the same host already exists.
func (c *Client) Add(cfg *Config) error {
	if cfg == nil {
		return fmt.Errorf("config cannot be nil")
	}

	d := newDeviceClient(*cfg)

	c.mu.Lock()
	defer c.mu.Unlock()

	if _, ok := c.devices[cfg.Host]; ok {
		return fmt.Errorf("device already exists: %s", cfg.Host)
	}

	c.devices[cfg.Host] = d
	return nil
}

// Del removes a device from the client.
// It returns an error if the device does not exist.
func (c *Client) Del(host string) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if _, ok := c.devices[host]; !ok {
		return fmt.Errorf("device not found: %s", host)
	}

	delete(c.devices, host)
	return nil
}

func (c *Client) getDeviceByHost(host string) (*deviceClient, error) {
	c.mu.RLock()
	defer c.mu.RUnlock()

	d, ok := c.devices[host]
	if !ok {
		return nil, fmt.Errorf("device not found: %s", host)
	}
	return d, nil
}

func (d *deviceClient) login(ctx context.Context) error {
	d.mu.Lock()
	defer d.mu.Unlock()

	reqUrl := fmt.Sprintf("%s://%s/api/v1.0/user/login", d.config.Scheme, d.config.Host)
	payload := map[string]string{
		"username": d.config.Username,
		"password": d.config.Password,
	}
	body, err := json.Marshal(payload)
	if err != nil {
		return err
	}

	req, err := http.NewRequestWithContext(ctx, "POST", reqUrl, bytes.NewBuffer(body))
	if err != nil {
		return err
	}

	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")
	req.URL.User = url.UserPassword(d.config.Username, d.config.Password)

	resp, err := d.client.Do(req)
	if err != nil {
		return err
	}

	defer resp.Body.Close()
	respPayload, err := io.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("login failed [%d]: %s", resp.StatusCode, string(respPayload))
	}

	token := resp.Header.Get("x-auth-token")
	if token == "" {
		return fmt.Errorf("login failed: no token in response")
	}

	d.token = token
	return nil
}

func (d *deviceClient) do(ctx context.Context, method, path string, body any, out any) error {
	// First attempt
	err := d.doRequest(ctx, method, path, body, out)
	if err == nil {
		return nil
	}

	// If unauthorized, try to login and retry
	if strings.Contains(err.Error(), "status 401") || strings.Contains(err.Error(), "unauthorized") {
		if loginErr := d.login(ctx); loginErr != nil {
			return fmt.Errorf("re-login failed: %w", loginErr)
		}
		return d.doRequest(ctx, method, path, body, out)
	}

	return err
}

func (d *deviceClient) doRequest(ctx context.Context, method, path string, body any, out any) error {
	url := fmt.Sprintf("%s://%s%s", d.config.Scheme, d.config.Host, path)

	var reqBody io.Reader
	if body != nil {
		b, err := json.Marshal(body)
		if err != nil {
			return err
		}
		reqBody = bytes.NewBuffer(b)
	}

	req, err := http.NewRequestWithContext(ctx, method, url, reqBody)
	if err != nil {
		return err
	}

	d.mu.Lock()
	token := d.token
	d.mu.Unlock()

	if token != "" {
		req.Header.Set("x-auth-token", token)
	}

	// Some endpoints might require Content-Type even for GET if we were strict, but usually only for POST/PUT
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}

	resp, err := d.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusUnauthorized {
		return fmt.Errorf("status 401")
	}

	if resp.StatusCode != http.StatusOK {
		// Read body to see error message
		b, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("request failed: status %d, body: %s", resp.StatusCode, string(b))
	}

	if out != nil {
		if err := json.NewDecoder(resp.Body).Decode(out); err != nil {
			return err
		}
	}

	return nil
}