spinifex/spinifex/utils/utils.go

package utils

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net"
	"net/http"
	"os"
	"os/signal"
	"reflect"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/private/protocol/xml/xmlutil"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/mulgadc/spinifex/spinifex/awserrors"
	"github.com/pterm/pterm"
)

// GenerateResourceID generates a unique resource ID with the given prefix.
// Format: {prefix}-{17 hex chars} using crypto/rand.
func GenerateResourceID(prefix string) string {
	b := make([]byte, 9)
	if _, err := rand.Read(b); err != nil {
		//nolint:forbidigo // Continuing after the CSPRNG fails could create predictable resource IDs.
		panic("crypto/rand failed: " + err.Error())
	}
	return prefix + "-" + hex.EncodeToString(b)[:17]
}

func MarshalToXML(payload any) ([]byte, error) {
	var buf bytes.Buffer
	enc := xml.NewEncoder(&buf)

	if err := xmlutil.BuildXML(payload, enc); err != nil {
		slog.Error("BuildXML failed", "err", err)
		return nil, err
	}

	if err := enc.Flush(); err != nil {
		slog.Error("Flush failed", "err", err)
		return nil, err
	}

	return buf.Bytes(), nil
}

// GenerateXMLPayload wraps payload with the requested locationName tag.
func GenerateXMLPayload(locationName string, payload any) any {
	t := reflect.StructOf([]reflect.StructField{
		{
			Name: "Value",
			Type: reflect.TypeOf(payload),
			Tag:  reflect.StructTag(`locationName:"` + locationName + `"`),
		},
	})

	v := reflect.New(t).Elem()
	v.Field(0).Set(reflect.ValueOf(payload))
	return v.Interface()
}

// NormalizeXMLOutput returns a copy of output with every nil slice field
// (recursively) replaced by a non-nil empty slice, since aws-sdk-go's
// xmlutil.BuildXML omits a nil slice's container element entirely but
// renders an empty one for a non-nil empty slice, unlike real AWS.
func NormalizeXMLOutput(output any) any {
	v := reflect.ValueOf(output)
	if !v.IsValid() {
		return output
	}
	// Work on an addressable copy: callers pass struct values, and reflection
	// can only Set fields through an addressable Value.
	ptr := reflect.New(v.Type())
	ptr.Elem().Set(v)
	normalizeNilSlices(ptr.Elem())
	return ptr.Elem().Interface()
}

// normalizeNilSlices walks v in place, turning nil slice fields into empty
// ones and recursing into structs, pointers, and existing slice elements.
func normalizeNilSlices(v reflect.Value) {
	switch v.Kind() {
	case reflect.Pointer:
		if !v.IsNil() {
			normalizeNilSlices(v.Elem())
		}
	case reflect.Struct:
		for _, field := range v.Fields() {
			switch field.Kind() {
			case reflect.Slice:
				if field.IsNil() {
					if field.CanSet() {
						field.Set(reflect.MakeSlice(field.Type(), 0, 0))
					}
				} else {
					for j := 0; j < field.Len(); j++ {
						normalizeNilSlices(field.Index(j))
					}
				}
			case reflect.Pointer, reflect.Struct:
				normalizeNilSlices(field)
			}
		}
	}
}

// WithRequestID returns a copy of payload's structure with a synthetic
// RequestId field prepended, since the SDK's generated output structs never
// carry one. payload must be a struct or pointer to struct; anything else is
// returned unchanged.
func WithRequestID(payload any, requestID string) any {
	pv := reflect.ValueOf(payload)
	for pv.Kind() == reflect.Pointer {
		pv = pv.Elem()
	}
	if pv.Kind() != reflect.Struct {
		return payload
	}
	pt := pv.Type()

	fields := []reflect.StructField{
		{
			Name: "RequestId",
			Type: reflect.TypeFor[string](),
			Tag:  reflect.StructTag(`locationName:"requestId" type:"string"`),
		},
	}
	for f := range pt.Fields() {
		if f.PkgPath == "" { // skip unexported marker fields (e.g. "_")
			fields = append(fields, f)
		}
	}

	composite := reflect.New(reflect.StructOf(fields)).Elem()
	composite.FieldByName("RequestId").SetString(requestID)
	for i := 0; i < pt.NumField(); i++ {
		if f := pt.Field(i); f.PkgPath == "" {
			composite.FieldByName(f.Name).Set(pv.Field(i))
		}
	}

	return composite.Interface()
}

// GenerateIAMXMLPayload wraps IAM output in the <ActionResponse><ActionResult>...</ActionResult></ActionResponse> structure.
func GenerateIAMXMLPayload(action string, payload any) any {
	resultName := action + "Result"
	resultWrapper := reflect.StructOf([]reflect.StructField{
		{
			Name: "Result",
			Type: reflect.TypeOf(payload),
			Tag:  reflect.StructTag(`locationName:"` + resultName + `"`),
		},
	})
	resultV := reflect.New(resultWrapper).Elem()
	resultV.Field(0).Set(reflect.ValueOf(payload))

	responseName := action + "Response"
	responseWrapper := reflect.StructOf([]reflect.StructField{
		{
			Name: "Response",
			Type: resultWrapper,
			Tag:  reflect.StructTag(`locationName:"` + responseName + `"`),
		},
	})
	responseV := reflect.New(responseWrapper).Elem()
	responseV.Field(0).Set(resultV)

	return responseV.Interface()
}

// GenerateErrorPayload serializes an ec2.ResponseError with the given code as JSON.
func GenerateErrorPayload(code string) (jsonResponse []byte) {
	return GenerateErrorPayloadWithMessage(code, "")
}

// GenerateErrorPayloadWithMessage serializes an ec2.ResponseError carrying both
// the sanitized code and the original error message, so the client can surface
// the actionable reason instead of a bare code. Message is omitted when empty
// or identical to the code.
func GenerateErrorPayloadWithMessage(code, message string) (jsonResponse []byte) {
	var responseError ec2.ResponseError
	responseError.Code = aws.String(code)
	if message != "" && message != code {
		responseError.Message = aws.String(message)
	}
	jsonResponse, err := json.Marshal(responseError)
	if err != nil {
		slog.Error("GenerateErrorPayload could not marshal JSON payload", "err", err)
		return nil
	}

	return jsonResponse
}

// ValidateErrorPayload decodes payload as an ec2.ResponseError and returns an error when a non-nil Code is detected.
func ValidateErrorPayload(payload []byte) (responseError ec2.ResponseError, err error) {
	decoder := json.NewDecoder(bytes.NewReader(payload))
	decoder.DisallowUnknownFields()

	err = decoder.Decode(&responseError)

	if err == nil && responseError.Code != nil {
		return responseError, errors.New("ResponseError detected")
	}
	return responseError, nil
}

// UnmarshalJsonPayload decodes jsonData into input (already a pointer) using strict field checking.
func UnmarshalJsonPayload(input any, jsonData []byte) []byte {
	decoder := json.NewDecoder(bytes.NewReader(jsonData))
	decoder.DisallowUnknownFields()
	err := decoder.Decode(input)
	if err != nil {
		return GenerateErrorPayload(awserrors.ErrorValidationError)
	}

	return nil
}

// ValidateKeyPairName validates that a key pair name contains only [A-Za-z0-9._-].
// Rejects empty names and returns ErrorInvalidKeyPairFormat on any invalid character.
func ValidateKeyPairName(name string) error {
	if name == "" {
		return errors.New("key name cannot be empty")
	}

	for _, char := range name {
		valid := (char >= 'A' && char <= 'Z') ||
			(char >= 'a' && char <= 'z') ||
			(char >= '0' && char <= '9') ||
			char == '-' ||
			char == '_' ||
			char == '.'

		if !valid {
			return errors.New(awserrors.ErrorInvalidKeyPairFormat)
		}
	}

	return nil
}

func DownloadFileWithProgress(url string, name string, filename string, timeout time.Duration) (err error) {
	ctx, cancel := context.WithCancel(context.Background())
	if timeout > 0 {
		ctx, cancel = context.WithTimeout(ctx, timeout)
	}
	defer cancel()
	intCh := make(chan os.Signal, 1)
	signal.Notify(intCh, os.Interrupt)
	go func() {
		<-intCh
		cancel()
	}()

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return fmt.Errorf("request error: %w", err)
	}

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return fmt.Errorf("http error: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode > 299 {
		return fmt.Errorf("unexpected status: %s", resp.Status)
	}

	f, err := os.Create(filename)
	if err != nil {
		return fmt.Errorf("file create error: %w", err)
	}
	defer f.Close()

	cl := resp.ContentLength

	if cl > 0 {
		total := SafeInt64ToUint64(cl)
		bar, update := NewByteProgressBar(fmt.Sprintf("Downloading %s", name), total)

		// io.Copy writes ~32 KiB per TeeReader call, so gate rendering on
		// integer-percentage change to cap the bar at ≤101 renders instead of
		// re-rendering tens of thousands of times.
		var current uint64
		lastPct := -1
		reader := io.TeeReader(resp.Body, progressWriter(func(n int) {
			current += SafeIntToUint64(n)
			pct := SafeUint64ToInt(current * 100 / total)
			if pct > lastPct {
				lastPct = pct
				update(current)
			}
		}))

		_, err = io.Copy(f, reader)
		_, _ = bar.Stop()
		if err != nil {
			return fmt.Errorf("copy error: %w", err)
		}
		return err
	} else {
		spin, _ := pterm.DefaultSpinner.
			WithText("Downloading (size unknown)...").
			Start()
		var written int64
		reader := io.TeeReader(resp.Body, progressWriter(func(n int) {
			written += int64(n)
			spin.UpdateText(fmt.Sprintf("Downloading %s (%s) ...", name, HumanBytes(SafeInt64ToUint64(written))))
		}))
		_, err = io.Copy(f, reader)
		_ = spin.Stop()

		if err != nil {
			return fmt.Errorf("copy error: %w", err)
		}
	}

	return nil
}

// progressWriter turns byte counts into a callback for UI updates.
type progressWriter func(n int)

func (pw progressWriter) Write(p []byte) (int, error) {
	pw(len(p))
	return len(p), nil
}

// NewByteProgressBar starts a pterm progress bar that renders human-readable
// sizes in its title instead of raw byte counts, and returns an update func
// that performs one render per call. Callers must invoke update at a throttled
// cadence (integer-percentage change) — the render itself is expensive, so an
// unthrottled per-write call regresses throughput.
//
// pterm's elapsed-time display normally spawns a background goroutine that
// re-renders every second with no lock, tearing the line against our own
// low-frequency renders. Starting with it off means no timer is spawned; the
// flag is then set on the returned bar so pterm still appends its "| 9s"
// suffix, now emitted only on our (single) renders.
func NewByteProgressBar(title string, total uint64) (*pterm.ProgressbarPrinter, func(current uint64)) {
	totalHuman := HumanBytes(total)
	bar, _ := pterm.DefaultProgressbar.
		WithTitle(title).
		WithTotal(SafeUint64ToInt(total)).
		WithShowCount(false).       // hide raw ints; the size goes in the title
		WithShowElapsedTime(false). // suppress the async re-render (see above)
		Start()
	bar.ShowElapsedTime = true // keep pterm's elapsed, rendered only by us

	update := func(current uint64) {
		bar.Current = SafeUint64ToInt(current) // drives fill + percentage
		// UpdateTitle performs the single render for this step.
		bar.UpdateTitle(fmt.Sprintf("%s — %s / %s", title, HumanBytes(current), totalHuman))
	}
	return bar, update
}

// HumanBytes formats a byte count using IEC binary suffixes (KiB, MiB, ...).
// Values below 1024 render as exact bytes.
func HumanBytes(b uint64) string {
	const unit = 1024
	if b < unit {
		return fmt.Sprintf("%d B", b)
	}
	div, exp := uint64(unit), 0
	for n := b / unit; n >= unit; n /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.1f %ciB", float64(b)/float64(div), "KMGTPEZY"[exp])
}

// HashMAC returns a deterministic locally-administered unicast MAC for id (SHA-256; first octet 0x02).
// id must be globally unique; callers sharing a base id across resource classes must compose a class tag (e.g. "dev:"+id).
func HashMAC(id string) string {
	sum := sha256.Sum256([]byte(id))
	b := make([]byte, 6)
	b[0] = 0x02
	copy(b[1:], sum[:5])
	return net.HardwareAddr(b).String()
}

// ClientIP returns the IP from a RemoteAddr, stripping the port. Handles both
// IPv4 and IPv6, and tolerates a RemoteAddr that is already a bare IP (no port).
func ClientIP(remoteAddr string) string {
	host, _, err := net.SplitHostPort(remoteAddr)
	if err != nil {
		return remoteAddr
	}
	return host
}

func dirExists(path string) bool {
	info, err := os.Stat(path)
	if os.IsNotExist(err) {
		return false
	}
	if err != nil {
		return false
	}
	return info.IsDir()
}