ssh-netbox-wrapper/internal/netbox/client.go

package netbox

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

type Client struct {
	baseURL      string
	token        string
	tokenVersion int
	httpClient   *http.Client
}

// NewClient creates a NetBox API client. Pass tokenVersion=0 to auto-detect
// from the token string (1 for legacy, 2 for nbt_-prefixed tokens).
func NewClient(baseURL, token string, tokenVersion int) *Client {
	if tokenVersion == 0 {
		tokenVersion = TokenVersion(token)
	}
	return &Client{
		baseURL:      strings.TrimRight(baseURL, "/"),
		token:        token,
		tokenVersion: tokenVersion,
		httpClient:   &http.Client{Timeout: 15 * time.Second},
	}
}

// Search queries up to 50 devices and VMs in parallel and merges the results.
// Use SearchOptions to restrict by kind or tag.
func (c *Client) Search(ctx context.Context, query string, opts SearchOptions) ([]HostEntry, error) {
	if opts.Kind != "" && opts.Kind != "device" && opts.Kind != "vm" {
		return nil, fmt.Errorf("invalid host kind %q", opts.Kind)
	}
	var (
		mu      sync.Mutex
		results []HostEntry
		errs    []error
		wg      sync.WaitGroup
		started int
	)

	if opts.Kind != "vm" {
		started++
		wg.Add(1)
		go func() {
			defer wg.Done()
			devices, err := c.searchDevices(ctx, query, opts.Tag)
			mu.Lock()
			defer mu.Unlock()
			if err != nil {
				errs = append(errs, fmt.Errorf("devices: %w", err))
				return
			}
			results = append(results, devices...)
		}()
	}

	if opts.Kind != "device" {
		started++
		wg.Add(1)
		go func() {
			defer wg.Done()
			vms, err := c.searchVMs(ctx, query, opts.Tag)
			mu.Lock()
			defer mu.Unlock()
			if err != nil {
				errs = append(errs, fmt.Errorf("vms: %w", err))
				return
			}
			results = append(results, vms...)
		}()
	}

	wg.Wait()

	if len(errs) == started {
		if started == 1 {
			return nil, errs[0]
		}
		return nil, fmt.Errorf("netbox search failed: %v; %v", errs[0], errs[1])
	}
	return results, nil
}

// SearchAll paginates through all matching devices and VMs, fetching every page.
// Intended for cache refresh; use Search for interactive queries.
func (c *Client) SearchAll(ctx context.Context, query string, opts SearchOptions) ([]HostEntry, error) {
	if opts.Kind != "" && opts.Kind != "device" && opts.Kind != "vm" {
		return nil, fmt.Errorf("invalid host kind %q", opts.Kind)
	}
	var (
		mu      sync.Mutex
		results []HostEntry
		errs    []error
		wg      sync.WaitGroup
	)

	if opts.Kind != "vm" {
		wg.Add(1)
		go func() {
			defer wg.Done()
			devices, err := c.fetchAllDevices(ctx, query, opts.Tag)
			mu.Lock()
			defer mu.Unlock()
			if err != nil {
				errs = append(errs, fmt.Errorf("devices: %w", err))
				return
			}
			results = append(results, devices...)
		}()
	}

	if opts.Kind != "device" {
		wg.Add(1)
		go func() {
			defer wg.Done()
			vms, err := c.fetchAllVMs(ctx, query, opts.Tag)
			mu.Lock()
			defer mu.Unlock()
			if err != nil {
				errs = append(errs, fmt.Errorf("vms: %w", err))
				return
			}
			results = append(results, vms...)
		}()
	}

	wg.Wait()

	if len(errs) > 0 {
		return nil, errors.Join(errs...)
	}
	return results, nil
}

// GetIPs returns all IP addresses assigned to a host, used by resolver strategies
// that need more than just the primary IP.
func (c *Client) GetIPs(ctx context.Context, entry HostEntry) ([]string, error) {
	var apiURL string
	switch entry.Kind {
	case "device":
		apiURL = fmt.Sprintf("%s/api/ipam/ip-addresses/?device_id=%d&limit=100", c.baseURL, entry.ID)
	case "vm":
		apiURL = fmt.Sprintf("%s/api/ipam/ip-addresses/?virtual_machine_id=%d&limit=100", c.baseURL, entry.ID)
	default:
		return nil, fmt.Errorf("unknown host kind: %q", entry.Kind)
	}

	var resp netboxIPListResponse
	if err := c.get(ctx, apiURL, &resp); err != nil {
		return nil, err
	}

	ips := make([]string, 0, len(resp.Results))
	for _, r := range resp.Results {
		ips = append(ips, stripPrefix(r.Address))
	}
	return ips, nil
}

// GetIPsWithFilter calls /api/ipam/ip-addresses/ with arbitrary filter query parameters.
func (c *Client) GetIPsWithFilter(ctx context.Context, filterParams string) ([]string, error) {
	apiURL := fmt.Sprintf("%s/api/ipam/ip-addresses/?%s&limit=100", c.baseURL, filterParams)
	var resp netboxIPListResponse
	if err := c.get(ctx, apiURL, &resp); err != nil {
		return nil, err
	}
	ips := make([]string, 0, len(resp.Results))
	for _, r := range resp.Results {
		ips = append(ips, stripPrefix(r.Address))
	}
	return ips, nil
}

func (c *Client) searchDevices(ctx context.Context, query, tag string) ([]HostEntry, error) {
	apiURL := fmt.Sprintf("%s/api/dcim/devices/?name__ic=%s&limit=50", c.baseURL, url.QueryEscape(query))
	if tag != "" {
		apiURL += "&tag=" + url.QueryEscape(tag)
	}
	var resp netboxListResponse[netboxDevice]
	if err := c.get(ctx, apiURL, &resp); err != nil {
		return nil, err
	}
	entries := make([]HostEntry, 0, len(resp.Results))
	for _, d := range resp.Results {
		entries = append(entries, deviceToEntry(d))
	}
	return entries, nil
}

func (c *Client) searchVMs(ctx context.Context, query, tag string) ([]HostEntry, error) {
	apiURL := fmt.Sprintf("%s/api/virtualization/virtual-machines/?name__ic=%s&limit=50", c.baseURL, url.QueryEscape(query))
	if tag != "" {
		apiURL += "&tag=" + url.QueryEscape(tag)
	}
	var resp netboxListResponse[netboxVM]
	if err := c.get(ctx, apiURL, &resp); err != nil {
		return nil, err
	}
	entries := make([]HostEntry, 0, len(resp.Results))
	for _, v := range resp.Results {
		entries = append(entries, vmToEntry(v))
	}
	return entries, nil
}

func (c *Client) fetchAllDevices(ctx context.Context, query, tag string) ([]HostEntry, error) {
	const pageSize = 50
	var all []HostEntry
	for offset := 0; ; offset += pageSize {
		apiURL := fmt.Sprintf("%s/api/dcim/devices/?name__ic=%s&limit=%d&offset=%d",
			c.baseURL, url.QueryEscape(query), pageSize, offset)
		if tag != "" {
			apiURL += "&tag=" + url.QueryEscape(tag)
		}
		var resp netboxListResponse[netboxDevice]
		if err := c.get(ctx, apiURL, &resp); err != nil {
			return nil, err
		}
		for _, d := range resp.Results {
			all = append(all, deviceToEntry(d))
		}
		if len(resp.Results) == 0 || len(all) >= resp.Count {
			break
		}
	}
	return all, nil
}

func (c *Client) fetchAllVMs(ctx context.Context, query, tag string) ([]HostEntry, error) {
	const pageSize = 50
	var all []HostEntry
	for offset := 0; ; offset += pageSize {
		apiURL := fmt.Sprintf("%s/api/virtualization/virtual-machines/?name__ic=%s&limit=%d&offset=%d",
			c.baseURL, url.QueryEscape(query), pageSize, offset)
		if tag != "" {
			apiURL += "&tag=" + url.QueryEscape(tag)
		}
		var resp netboxListResponse[netboxVM]
		if err := c.get(ctx, apiURL, &resp); err != nil {
			return nil, err
		}
		for _, v := range resp.Results {
			all = append(all, vmToEntry(v))
		}
		if len(resp.Results) == 0 || len(all) >= resp.Count {
			break
		}
	}
	return all, nil
}

// TokenVersion returns 2 for NetBox v2 tokens (nbt_ prefix) or 1 for legacy tokens.
func TokenVersion(token string) int {
	if strings.HasPrefix(token, "nbt_") {
		return 2
	}
	return 1
}

func (c *Client) get(ctx context.Context, apiURL string, out any) error {
	_, err := c.request(ctx, http.MethodGet, apiURL, nil, "", out)
	return err
}

func (c *Client) request(ctx context.Context, method, apiURL string, body io.Reader, etag string, out any) (string, error) {
	req, err := http.NewRequestWithContext(ctx, method, apiURL, body)
	if err != nil {
		return "", fmt.Errorf("creating request: %w", err)
	}
	scheme := "Token "
	if c.tokenVersion == 2 {
		scheme = "Bearer "
	}
	req.Header.Set("Authorization", scheme+c.token)
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	if etag != "" {
		req.Header.Set("If-Match", etag)
	}
	req.Header.Set("Accept", "application/json")

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return "", fmt.Errorf("request to %s: %w", apiURL, err)
	}
	defer resp.Body.Close()

	if resp.StatusCode == http.StatusForbidden {
		hint := "check token permissions in NetBox"
		if c.tokenVersion == 1 {
			hint += " — legacy v1 token detected, consider upgrading to a v2 token (starts with nbt_)"
		}
		return "", fmt.Errorf("%s: %s", apiURL, hint)
	}

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("netbox returned %d for %s", resp.StatusCode, apiURL)
	}

	if err := json.NewDecoder(resp.Body).Decode(out); err != nil {
		return "", fmt.Errorf("decoding response: %w", err)
	}
	return resp.Header.Get("ETag"), nil
}

func deviceToEntry(d netboxDevice) HostEntry {
	e := HostEntry{ID: d.ID, Name: d.Name, Kind: "device"}
	if d.PrimaryIP4 != nil {
		e.PrimaryIP4 = stripPrefix(d.PrimaryIP4.Address)
	}
	if d.PrimaryIP6 != nil {
		e.PrimaryIP6 = stripPrefix(d.PrimaryIP6.Address)
	}
	for _, t := range d.Tags {
		e.Tags = append(e.Tags, t.Name)
	}
	return e
}

func vmToEntry(v netboxVM) HostEntry {
	e := HostEntry{ID: v.ID, Name: v.Name, Kind: "vm"}
	if v.PrimaryIP4 != nil {
		e.PrimaryIP4 = stripPrefix(v.PrimaryIP4.Address)
	}
	if v.PrimaryIP6 != nil {
		e.PrimaryIP6 = stripPrefix(v.PrimaryIP6.Address)
	}
	for _, t := range v.Tags {
		e.Tags = append(e.Tags, t.Name)
	}
	return e
}

// stripPrefix removes the CIDR prefix length from a NetBox IP (e.g. "10.0.1.5/24" → "10.0.1.5").
func stripPrefix(cidr string) string {
	if idx := strings.Index(cidr, "/"); idx != -1 {
		return cidr[:idx]
	}
	return cidr
}

// ErrHostNotFound distinguishes a missing exact match from API failures or ambiguity.
var ErrHostNotFound = errors.New("host not found")

// FindHost refuses ambiguous names so writes cannot target an arbitrary server.
func (c *Client) FindHost(ctx context.Context, name, kind string) (HostEntry, error) {
	if kind != "" && kind != "device" && kind != "vm" {
		return HostEntry{}, fmt.Errorf("invalid kind %q", kind)
	}
	entries, err := c.SearchAll(ctx, name, SearchOptions{Kind: kind})
	if err != nil {
		return HostEntry{}, err
	}
	var matches []HostEntry
	for _, e := range entries {
		if strings.EqualFold(e.Name, name) {
			matches = append(matches, e)
		}
	}
	if len(matches) == 0 {
		return HostEntry{}, fmt.Errorf("%w: %q", ErrHostNotFound, name)
	}
	if len(matches) != 1 {
		return HostEntry{}, fmt.Errorf("host %q: found %d exact matches; use --kind or a unique hostname", name, len(matches))
	}
	return matches[0], nil
}

func (c *Client) hostURL(host HostEntry) (string, error) {
	var path string
	switch host.Kind {
	case "device":
		path = "dcim/devices"
	case "vm":
		path = "virtualization/virtual-machines"
	default:
		return "", fmt.Errorf("invalid host kind %q", host.Kind)
	}
	if host.ID <= 0 {
		return "", fmt.Errorf("invalid host ID")
	}
	return fmt.Sprintf("%s/api/%s/%d/", c.baseURL, path, host.ID), nil
}

func (c *Client) GetComments(ctx context.Context, host HostEntry) (string, string, error) {
	u, err := c.hostURL(host)
	if err != nil {
		return "", "", err
	}
	var out struct {
		Comments string `json:"comments"`
	}
	etag, err := c.request(ctx, http.MethodGet, u, nil, "", &out)
	return out.Comments, etag, err
}

func (c *Client) SetComments(ctx context.Context, host HostEntry, comments, etag string) error {
	u, err := c.hostURL(host)
	if err != nil {
		return err
	}
	data, err := json.Marshal(map[string]string{"comments": comments})
	if err != nil {
		return err
	}
	var out json.RawMessage
	_, err = c.request(ctx, http.MethodPatch, u, bytes.NewReader(data), etag, &out)
	return err
}