jpictl/pkg/zones/env.go

package zones

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

// WriteEnv generates environment variables for shell scripts
func (m *Zones) WriteEnv(w io.Writer) error {
	var buf bytes.Buffer

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

	_, err := buf.WriteTo(w)
	return err
}

func (m *Zones) 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)
	m.writeEnvVar(w, fmt.Sprintf("%v", gatewayID), "ZONE%v_%s", zoneID, "GW")

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

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

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

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

func (*Zones) writeEnvVar(w io.Writer, value string, name string, args ...any) {
	if len(args) > 0 {
		name = fmt.Sprintf(name, args...)
	}

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

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

func genEnvZones(m *Zones) string {
	s := make([]string, 0, len(m.Zones))
	for _, z := range m.Zones {
		s = append(s, fmt.Sprintf("%v", z.ID))
	}

	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 firstNodeID, gatewayID int

	z.ForEachMachine(func(p *Machine) bool {
		if firstNodeID == 0 {
			firstNodeID = p.ID
		}

		if p.IsGateway() {
			gatewayID = p.ID
		}

		return gatewayID != 0
	})

	switch {
	case gatewayID == 0:
		return firstNodeID
	default:
		return gatewayID
	}
}