jpictl/pkg/zones/env.go

package zones

import (
	"bytes"
	"fmt"
	"io"
	"strings"
)

// Env is a shell environment factory for this cluster
type Env struct {
	ZoneIterator

	export bool
}

// Env returns a shell environment factory
func (m *Zones) Env(export bool) *Env {
	return &Env{
		ZoneIterator: m,
		export:       export,
	}
}

// WriteTo generates environment variables for shell scripts
func (m *Env) WriteTo(w io.Writer) (int64, error) {
	var buf bytes.Buffer

	m.writeEnvVarFn(&buf, genEnvZones, "ZONES")
	m.ForEachZone(func(z *Zone) bool {
		m.writeEnvZone(&buf, z)
		return false
	})

	return buf.WriteTo(w)
}

func (m *Env) writeEnvZone(w io.Writer, z *Zone) {
	zoneID := z.ID

	// ZONE{zoneID}
	m.writeEnvVar(w, genEnvZoneNodes(z), "ZONE%v", zoneID)

	// ZONE{zoneID}_NAME
	m.writeEnvVar(w, z.Name, "ZONE%v_%s", zoneID, "NAME")

	// ZONE{zoneID}_GW
	gatewayID := getRingZeroGatewayID(z)
	if gatewayID > 0 {
		m.writeEnvVar(w, fmt.Sprintf("%v", gatewayID), "ZONE%v_%s", zoneID, "GW")

		// ZONE{zoneID}_IP
		if ip, ok := RingZeroAddress(zoneID, gatewayID); ok {
			m.writeEnvVar(w, ip.String(), "ZONE%v_%s", zoneID, "IP")
		}
	}
}

func (m *Env) writeEnvVarFn(w io.Writer, fn func(*Env) string, name string, args ...any) {
	var value string

	if fn != nil {
		value = fn(m)
	}

	m.writeEnvVar(w, value, name, args...)
}

func (m *Env) writeEnvVar(w io.Writer, value string, name string, args ...any) {
	var prefix string

	if m.export {
		prefix = "export "
	}

	if len(args) > 0 {
		name = fmt.Sprintf(name, args...)
	}

	if name != "" {
		value = strings.TrimSpace(value)

		_, _ = fmt.Fprintf(w, "%s%s=%q\n", prefix, name, value)
	}
}

func genEnvZones(m *Env) string {
	var s []string

	m.ForEachZone(func(z *Zone) bool {
		s = append(s, fmt.Sprintf("%v", z.ID))
		return false
	})

	return strings.Join(s, " ")
}

func genEnvZoneNodes(z *Zone) string {
	s := make([]string, 0, len(z.Machines))
	for _, p := range z.Machines {
		s = append(s, p.Name)
	}
	return strings.Join(s, " ")
}

func getRingZeroGatewayID(z *Zone) int {
	var gatewayID int

	z.ForEachMachine(func(p *Machine) bool {
		if p.IsGateway() {
			gatewayID = p.ID
		}

		return gatewayID != 0
	})

	return gatewayID
}