xinfra/server/internal/service/awx.go

package service

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

type AWXClient struct {
	baseURL      string
	staticToken  string
	username     string
	password     string
	cachedToken  string
	tokenExpires time.Time
	tokenMu      sync.Mutex
	client       *http.Client
}

type AWXLaunchRequest struct {
	InventoryID uint64
	Limit       string
	ExtraVars   map[string]any
}

type AWXJob struct {
	ID            uint64         `json:"id"`
	Status        string         `json:"status"`
	Failed        bool           `json:"failed"`
	IgnoredFields map[string]any `json:"ignored_fields"`
}

type AWXJobTemplate struct {
	ID                   uint64 `json:"id"`
	Name                 string `json:"name"`
	Description          string `json:"description"`
	Inventory            uint64 `json:"inventory"`
	AskVariablesOnLaunch bool   `json:"ask_variables_on_launch"`
	AskLimitOnLaunch     bool   `json:"ask_limit_on_launch"`
}

type AWXInventoryHost struct {
	ID        uint64 `json:"id"`
	Name      string `json:"name"`
	Enabled   bool   `json:"enabled"`
	Variables string `json:"variables"`
}

type awxListResponse[T any] struct {
	Next    string `json:"next"`
	Results []T    `json:"results"`
}

type awxTokenResponse struct {
	Token   string `json:"token"`
	Expires string `json:"expires"`
}

func NewAWXClient(baseURL, token string, credentials ...string) *AWXClient {
	username := ""
	password := ""
	if len(credentials) > 0 {
		username = credentials[0]
	}
	if len(credentials) > 1 {
		password = credentials[1]
	}
	return &AWXClient{
		baseURL:     strings.TrimRight(baseURL, "/"),
		staticToken: strings.TrimSpace(token),
		username:    strings.TrimSpace(username),
		password:    strings.TrimSpace(password),
		client:      &http.Client{Timeout: 30 * time.Second},
	}
}

func (c *AWXClient) Configured() bool {
	return c.baseURL != "" && (c.staticToken != "" || (c.username != "" && c.password != ""))
}

func (c *AWXClient) Launch(ctx context.Context, templateID uint64, input AWXLaunchRequest) (*AWXJob, error) {
	body := map[string]any{}
	if input.InventoryID != 0 {
		body["inventory"] = input.InventoryID
	}
	if input.ExtraVars != nil {
		body["extra_vars"] = input.ExtraVars
	}
	if input.Limit != "" {
		body["limit"] = input.Limit
	}
	var job AWXJob
	if err := c.request(ctx, http.MethodPost, fmt.Sprintf("/api/v2/job_templates/%d/launch/", templateID), body, &job); err != nil {
		return nil, err
	}
	if len(job.IgnoredFields) > 0 {
		return nil, fmt.Errorf("AWX ignored launch fields %v; enable Prompt on launch for Variables and Limit on the job template", ignoredFieldNames(job.IgnoredFields))
	}
	if job.ID == 0 {
		return nil, fmt.Errorf("AWX launch response did not include a job id")
	}
	return &job, nil
}

func ignoredFieldNames(fields map[string]any) []string {
	names := make([]string, 0, len(fields))
	for name := range fields {
		names = append(names, name)
	}
	return names
}

func (c *AWXClient) GetJob(ctx context.Context, jobID string) (*AWXJob, error) {
	if _, err := strconv.ParseUint(jobID, 10, 64); err != nil {
		return nil, fmt.Errorf("invalid AWX job id %q", jobID)
	}
	var job AWXJob
	if err := c.request(ctx, http.MethodGet, "/api/v2/jobs/"+jobID+"/", nil, &job); err != nil {
		return nil, err
	}
	return &job, nil
}

func (c *AWXClient) JobStdout(ctx context.Context, jobID string) (string, error) {
	if _, err := strconv.ParseUint(jobID, 10, 64); err != nil {
		return "", fmt.Errorf("invalid AWX job id %q", jobID)
	}
	raw, err := c.requestRaw(ctx, http.MethodGet, "/api/v2/jobs/"+jobID+"/stdout/?format=txt", nil, "text/plain")
	if err != nil {
		return "", err
	}
	return string(raw), nil
}

func (c *AWXClient) Cancel(ctx context.Context, jobID string) error {
	if _, err := strconv.ParseUint(jobID, 10, 64); err != nil {
		return fmt.Errorf("invalid AWX job id %q", jobID)
	}
	return c.request(ctx, http.MethodPost, "/api/v2/jobs/"+jobID+"/cancel/", map[string]any{}, nil)
}

func (c *AWXClient) WaitJob(ctx context.Context, jobID string, timeout time.Duration) (*AWXJob, error) {
	if timeout <= 0 {
		timeout = 45 * time.Second
	}
	deadlineCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	ticker := time.NewTicker(2 * time.Second)
	defer ticker.Stop()
	for {
		job, err := c.GetJob(deadlineCtx, jobID)
		if err != nil {
			return nil, err
		}
		switch job.Status {
		case "successful", "failed", "error", "canceled":
			return job, nil
		}
		select {
		case <-deadlineCtx.Done():
			return nil, fmt.Errorf("AWX job %s did not finish within %s", jobID, timeout)
		case <-ticker.C:
		}
	}
}

func (c *AWXClient) ListJobTemplates(ctx context.Context) ([]AWXJobTemplate, error) {
	var out []AWXJobTemplate
	path := "/api/v2/job_templates/?page_size=200"
	for path != "" {
		var page awxListResponse[AWXJobTemplate]
		if err := c.request(ctx, http.MethodGet, path, nil, &page); err != nil {
			return nil, err
		}
		out = append(out, page.Results...)
		path = awxNextPath(page.Next)
	}
	return out, nil
}

func (c *AWXClient) GetJobTemplate(ctx context.Context, templateID uint64) (*AWXJobTemplate, error) {
	if templateID == 0 {
		return nil, fmt.Errorf("invalid AWX job template id")
	}
	var item AWXJobTemplate
	if err := c.request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/job_templates/%d/", templateID), nil, &item); err != nil {
		return nil, err
	}
	if item.ID == 0 {
		return nil, fmt.Errorf("AWX job template %d was not found", templateID)
	}
	return &item, nil
}

func (c *AWXClient) ListInventoryHosts(ctx context.Context, inventoryID uint64) ([]AWXInventoryHost, error) {
	if inventoryID == 0 {
		return nil, fmt.Errorf("AWX job template does not bind an inventory")
	}
	var out []AWXInventoryHost
	path := fmt.Sprintf("/api/v2/inventories/%d/hosts/?page_size=200", inventoryID)
	for path != "" {
		var page awxListResponse[AWXInventoryHost]
		if err := c.request(ctx, http.MethodGet, path, nil, &page); err != nil {
			return nil, err
		}
		out = append(out, page.Results...)
		path = awxNextPath(page.Next)
	}
	return out, nil
}

func (c *AWXClient) GetHostFacts(ctx context.Context, hostID uint64) (map[string]any, error) {
	if hostID == 0 {
		return nil, fmt.Errorf("invalid AWX host id")
	}
	var facts map[string]any
	if err := c.request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/hosts/%d/ansible_facts/", hostID), nil, &facts); err != nil {
		return nil, err
	}
	if nested, ok := facts["ansible_facts"].(map[string]any); ok {
		return nested, nil
	}
	return facts, nil
}

func (c *AWXClient) request(ctx context.Context, method, path string, payload any, output any) error {
	raw, err := c.requestRaw(ctx, method, path, payload, "application/json")
	if err != nil {
		return err
	}
	if output != nil && len(raw) > 0 {
		if err := json.Unmarshal(raw, output); err != nil {
			return fmt.Errorf("decode AWX response: %w", err)
		}
	}
	return nil
}

func (c *AWXClient) requestRaw(ctx context.Context, method, path string, payload any, accept string) ([]byte, error) {
	if !c.Configured() {
		return nil, fmt.Errorf("AWX_BASE_URL and AWX_TOKEN or AWX_USERNAME/AWX_PASSWORD must be configured")
	}
	token, err := c.bearerToken(ctx)
	if err != nil {
		return nil, err
	}
	var body io.Reader
	if payload != nil {
		raw, err := json.Marshal(payload)
		if err != nil {
			return nil, err
		}
		body = bytes.NewReader(raw)
	}
	req, err := http.NewRequestWithContext(ctx, method, c.baseURL+path, body)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", "Bearer "+token)
	if accept == "" {
		accept = "application/json"
	}
	req.Header.Set("Accept", accept)
	if payload != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	resp, err := c.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("AWX request: %w", err)
	}
	defer resp.Body.Close()
	raw, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
	if err != nil {
		return nil, err
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, fmt.Errorf("AWX returned %s: %s", resp.Status, strings.TrimSpace(string(raw)))
	}
	return raw, nil
}

func (c *AWXClient) bearerToken(ctx context.Context) (string, error) {
	if c.staticToken != "" {
		return c.staticToken, nil
	}
	c.tokenMu.Lock()
	defer c.tokenMu.Unlock()
	if c.cachedToken != "" && time.Until(c.tokenExpires) > 5*time.Minute {
		return c.cachedToken, nil
	}
	token, expires, err := c.createToken(ctx)
	if err != nil {
		return "", err
	}
	c.cachedToken = token
	c.tokenExpires = expires
	return token, nil
}

func (c *AWXClient) createToken(ctx context.Context) (string, time.Time, error) {
	raw, err := json.Marshal(map[string]any{"description": "xinfra delivery"})
	if err != nil {
		return "", time.Time{}, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/api/v2/tokens/", bytes.NewReader(raw))
	if err != nil {
		return "", time.Time{}, err
	}
	req.SetBasicAuth(c.username, c.password)
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Content-Type", "application/json")
	resp, err := c.client.Do(req)
	if err != nil {
		return "", time.Time{}, fmt.Errorf("AWX token request: %w", err)
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return "", time.Time{}, err
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return "", time.Time{}, fmt.Errorf("AWX token returned %s: %s", resp.Status, strings.TrimSpace(string(body)))
	}
	var data awxTokenResponse
	if err := json.Unmarshal(body, &data); err != nil {
		return "", time.Time{}, fmt.Errorf("decode AWX token response: %w", err)
	}
	if strings.TrimSpace(data.Token) == "" {
		return "", time.Time{}, fmt.Errorf("AWX token response did not include a token")
	}
	expires := time.Now().Add(24 * time.Hour)
	if data.Expires != "" {
		if parsed, err := time.Parse(time.RFC3339, data.Expires); err == nil {
			expires = parsed
		}
	}
	return strings.TrimSpace(data.Token), expires, nil
}

func awxNextPath(next string) string {
	if next == "" {
		return ""
	}
	if strings.HasPrefix(next, "http://") || strings.HasPrefix(next, "https://") {
		if idx := strings.Index(next, "/api/"); idx >= 0 {
			return next[idx:]
		}
		return ""
	}
	return next
}

var ansibleHostPattern = regexp.MustCompile(`(?m)^\s*ansible_host\s*:\s*"?([^"\s]+)"?\s*$`)

func AWXHostIP(host AWXInventoryHost) string {
	var parsed map[string]any
	if err := json.Unmarshal([]byte(host.Variables), &parsed); err == nil {
		for _, key := range []string{"xinfra_public_address", "xinfra_host_address"} {
			if value, ok := parsed[key].(string); ok && strings.TrimSpace(value) != "" {
				return strings.TrimSpace(value)
			}
		}
		if value, ok := parsed["ansible_host"].(string); ok && strings.TrimSpace(value) != "" {
			return strings.TrimSpace(value)
		}
	}
	if match := ansibleHostPattern.FindStringSubmatch(host.Variables); len(match) == 2 {
		return strings.TrimSpace(match[1])
	}
	return host.Name
}