initrd/vpn.go

package main

import (
	"net"
	"os"
	"path/filepath"

	"github.com/rs/zerolog/log"
	"golang.zx2c4.com/wireguard/wgctrl"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"

	"novit.tech/direktil/pkg/config"
)

func setupVPN(vpn config.VPNDef, localGenDir string) {
	log := log.With().Str("vpn", vpn.Name).Logger()

	log.Info().Msg("VPN: setting up")

	vpnDir := filepath.Join(localGenDir, vpn.Name)
	os.MkdirAll(vpnDir, 0750)

	logMsg := log.Info()
	if vpn.ListenPort != nil {
		logMsg.Int("ListenPort", *vpn.ListenPort)
	}

	// public/private key
	keyFile := filepath.Join(vpnDir, "key")
	keyBytes, err := os.ReadFile(keyFile)
	if os.IsNotExist(err) {
		key, err := wgtypes.GeneratePrivateKey()
		if err != nil {
			fatalf("failed to generate VPN key: %v", err)
		}

		keyBytes = []byte(key.String())

		os.WriteFile(keyFile, keyBytes, 0600)
	} else if err != nil {
		fatalf("failed to read VPN key: %v", err)
	}

	key, err := wgtypes.ParseKey(string(keyBytes))
	if err != nil {
		fatalf("bad VPN key: %v", err)
	}

	logMsg.Stringer("PublicKey", key.PublicKey())

	// pre-shared key
	pskeyFile := filepath.Join(vpnDir, "pskey")
	pskeyBytes, err := os.ReadFile(pskeyFile)
	if os.IsNotExist(err) {
		key, err := wgtypes.GenerateKey()
		if err != nil {
			fatalf("failed to generate VPN pre-shared key: %v", err)
		}

		pskeyBytes = []byte(key.String())

		os.WriteFile(pskeyFile, pskeyBytes, 0600)
	} else if err != nil {
		fatalf("failed to read VPN pre-shared key: %v", err)
	}

	pskey, err := wgtypes.ParseKey(string(pskeyBytes))
	if err != nil {
		fatalf("bad VPN pre-shared key: %v", err)
	}

	{
		keyStr := key.String()
		logMsg.Str("PresharedKey", keyStr[0:4]+"..."+keyStr[len(keyStr)-4:])
	}

	// setup interface
	cfg := wgtypes.Config{
		PrivateKey: &key,
		ListenPort: vpn.ListenPort,
		Peers:      make([]wgtypes.PeerConfig, 0, len(vpn.Peers)),
	}

	for idx, vpnPeer := range vpn.Peers {
		vpnPeer := vpnPeer

		wgPeer := wgtypes.PeerConfig{
			Endpoint:   vpnPeer.Endpoint,
			AllowedIPs: make([]net.IPNet, 0, len(vpnPeer.AllowedIPs)),

			PersistentKeepaliveInterval: &vpnPeer.KeepAlive,
		}

		if vpnPeer.WithPreSharedKey {
			wgPeer.PresharedKey = &pskey
		}

		pubkey, err := wgtypes.ParseKey(vpnPeer.PublicKey)
		if err != nil {
			fatalf("bad VPN peer[%d] public key: %v", idx, err)
		}

		wgPeer.PublicKey = pubkey

		for _, ipnetStr := range vpnPeer.AllowedIPs {
			_, ipnet, err := net.ParseCIDR(ipnetStr)
			if err != nil {
				fatalf("bad IP/net: %q: %v", ipnetStr, err)
			}

			wgPeer.AllowedIPs = append(wgPeer.AllowedIPs, *ipnet)
		}

		cfg.Peers = append(cfg.Peers, wgPeer)
	}

	wg, err := wgctrl.New()
	if err != nil {
		fatalf("failed to setup WireGuard client: %v", err)
	}
	defer wg.Close()

	log.Info().Strs("ips", vpn.IPs).Msg("VPN: creating interface")
	run("ip", "link", "add", vpn.Name, "type", "wireguard")

	for _, ip := range vpn.IPs {
		run("ip", "addr", "add", ip, "dev", vpn.Name)
	}

	logMsg.Msg("VPN: configuring interface")
	err = wg.ConfigureDevice(vpn.Name, cfg)
	if err != nil {
		fatalf("failed to setup VPN %s: %v", vpn.Name, err)
	}

	run("ip", "link", "set", vpn.Name, "up")
}