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 {
	devices map[string]*deviceClient
}

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

// 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 {
		// Use Host as the key.
		// Ensure scheme is set
		if cfg.Scheme == "" {
			cfg.Scheme = "https"
		}

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

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

		devices[cfg.Host] = &deviceClient{
			config: cfg,
			client: client,
		}
	}

	return &Client{
		devices: devices,
	}
}

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
}