initrd/vpn.go

package main

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

	"golang.zx2c4.com/wireguard/wgctrl"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"

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

func setupVPN(vpn config.VPNDef, localGenDir string) {
	log.Printf("setting up VPN %s", vpn.Name)

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

	// 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)
	}

	log.Printf("VPN %s public key is %s", vpn.Name, key.PublicKey().String())

	// 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)
	}

	log.Printf("VPN %s pre-shared key is %s", vpn.Name, key.String())

	// 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()

	run("ip", "link", "add", vpn.Name, "type", "wireguard")

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

	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")
}