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

package netbox

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strings"
	"sync"
)

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{},
	}
}

// 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) {
	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) {
	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.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" {
		started++
		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) == started {
		if started == 1 {
			return nil, errs[0]
		}
		return nil, fmt.Errorf("netbox search failed: %v; %v", errs[0], errs[1])
	}
	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 {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, apiURL, nil)
	if err != nil {
		return fmt.Errorf("creating request: %w", err)
	}
	req.Header.Set("Authorization", "Token "+c.token)
	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 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
}