jpictl/pkg/zones/wireguard.go

package zones

import (
	"io/fs"
	"os"
)

var (
	_ WireguardConfigPruner = (*Zones)(nil)
	_ WireguardConfigPruner = (*Zone)(nil)
	_ WireguardConfigPruner = (*Machine)(nil)

	_ WireguardConfigWriter = (*Zones)(nil)
	_ WireguardConfigWriter = (*Zone)(nil)
	_ WireguardConfigWriter = (*Machine)(nil)

	_ WireguardKeysWriter = (*Zones)(nil)
	_ WireguardKeysWriter = (*Zone)(nil)
	_ WireguardKeysWriter = (*Machine)(nil)
)

// A WireguardConfigPruner deletes wgN.conf on all machines under
// its scope with the specified ring disabled
type WireguardConfigPruner interface {
	PruneWireguardConfig(ring int) error
}

// PruneWireguardConfig removes wgN.conf files of machines with
// the corresponding ring disabled on all zones
func (m *Zones) PruneWireguardConfig(ring int) error {
	return pruneWireguardConfig(m, ring)
}

// PruneWireguardConfig removes wgN.conf files of machines with
// the corresponding ring disabled.
func (z *Zone) PruneWireguardConfig(ring int) error {
	return pruneWireguardConfig(z, ring)
}

func pruneWireguardConfig(m MachineIterator, ring int) error {
	var err error

	m.ForEachMachine(func(p *Machine) bool {
		err = p.zone.PruneWireguardConfig(ring)
		if os.IsNotExist(err) {
			// ignore
			err = nil
		}

		return err != nil
	})

	return err
}

// PruneWireguardConfig deletes the wgN.conf file if its
// presence on the ring is disabled
func (m *Machine) PruneWireguardConfig(ring int) error {
	_, ok := m.getRingInfo(ring)
	if !ok {
		return m.RemoveWireguardConfig(ring)
	}

	return nil
}

// A WireguardConfigWriter rewrites all wgN.conf on all machines under
// its scope attached to that ring
type WireguardConfigWriter interface {
	WriteWireguardConfig(ring int) error
}

// WriteWireguardConfig rewrites all wgN.conf on all machines
// attached to that ring
func (m *Zones) WriteWireguardConfig(ring int) error {
	switch ring {
	case 0:
		return writeWireguardConfig(m, m, ring)
	case 1:
		var err error
		m.ForEachZone(func(z *Zone) bool {
			err = writeWireguardConfig(m, z, ring)
			return err != nil
		})
		return err
	default:
		return fs.ErrInvalid
	}
}

// WriteWireguardConfig rewrites all wgN.conf on all machines
// on the Zone attached to that ring
func (z *Zone) WriteWireguardConfig(ring int) error {
	switch ring {
	case 0:
		return writeWireguardConfig(z.zones, z.zones, ring)
	case 1:
		return writeWireguardConfig(z.zones, z, ring)
	default:
		return fs.ErrInvalid
	}
}

func writeWireguardConfig(z ZoneIterator, m MachineIterator, ring int) error {
	r, err := NewRing(z, m, ring)
	if err != nil {
		return err
	}

	r.ForEachMachine(func(p *Machine) bool {
		err = p.writeWireguardRingConfig(r)
		return err != nil
	})

	return err
}

// WriteWireguardConfig rewrites the wgN.conf file of this Machine
// if enabled
func (m *Machine) WriteWireguardConfig(ring int) error {
	r, err := NewRing(m.zone.zones, m.zone, ring)
	if err != nil {
		return err
	}

	return m.writeWireguardRingConfig(r)
}

func (m *Machine) writeWireguardRingConfig(r *Ring) error {
	wg, err := r.ExportConfig(m)
	if err != nil {
		return nil
	}

	f, err := m.CreateTruncFile("wg%v.conf", r.ID)
	if err != nil {
		return err
	}
	defer f.Close()

	_, err = wg.WriteTo(f)
	return err
}

// A WireguardKeysWriter writes the Wireguard Keys for all machines
// under its scope for the specified ring
type WireguardKeysWriter interface {
	WriteWireguardKeys(ring int) error
}

// WriteWireguardKeys rewrites all wgN.{key,pub} files
func (m *Zones) WriteWireguardKeys(ring int) error {
	return writeWireguardKeys(m, ring)
}

// WriteWireguardKeys rewrites all wgN.{key,pub} files on this zone
func (z *Zone) WriteWireguardKeys(ring int) error {
	return writeWireguardKeys(z, ring)
}

func writeWireguardKeys(m MachineIterator, ring int) error {
	var err error

	m.ForEachMachine(func(p *Machine) bool {
		err = p.WriteWireguardKeys(ring)
		if os.IsNotExist(err) {
			// ignore
			err = nil
		}

		return err != nil
	})

	return err
}

// WriteWireguardKeys writes the wgN.key/wgN.pub files
func (m *Machine) WriteWireguardKeys(ring int) error {
	var err error
	var key, pub string
	var ri *RingInfo

	ri, _ = m.getRingInfo(ring)
	if ri != nil {
		key = ri.Keys.PrivateKey.String()
		pub = ri.Keys.PublicKey.String()
	}

	switch {
	case key == "":
		return fs.ErrNotExist
	case pub == "":
		pub = ri.Keys.PrivateKey.Public().String()
	}

	err = m.WriteStringFile(key+"\n", "wg%v.key", ring)
	if err != nil {
		return err
	}

	err = m.WriteStringFile(pub+"\n", "wg%v.pub", ring)
	if err != nil {
		return err
	}

	return nil
}