wild-central/internal/wireguard/manager.go

package wireguard

import (
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/google/uuid"
	"gopkg.in/yaml.v3"

	"github.com/wild-cloud/wild-central/internal/storage"
)

// Config holds the WireGuard server interface configuration.
type Config struct {
	Enabled    bool   `yaml:"enabled"`
	ListenPort int    `yaml:"listenPort"`
	Address    string `yaml:"address"`  // server VPN IP, e.g. "10.8.0.1/24"
	Endpoint   string `yaml:"endpoint"` // public host:port clients connect to
	DNS        string `yaml:"dns"`      // DNS server pushed to clients
	LanCIDR    string `yaml:"lanCIDR"`  // LAN subnet to route through VPN, e.g. "192.168.8.0/24"
}

// Peer represents a WireGuard client peer.
type Peer struct {
	ID         string    `yaml:"id"`
	Name       string    `yaml:"name"`
	PublicKey  string    `yaml:"publicKey"`
	PrivateKey string    `yaml:"privateKey"` // stored so we can regenerate client config
	AllowedIPs string    `yaml:"allowedIPs"` // IP assigned to this peer, e.g. "10.8.0.2/32"
	CreatedAt  time.Time `yaml:"createdAt"`
}

// secrets holds server keypair (never returned via API).
type secrets struct {
	PrivateKey string `yaml:"privateKey"`
	PublicKey  string `yaml:"publicKey"`
}

// Status represents the current runtime state of the WireGuard interface.
type Status struct {
	Running    bool   `json:"running"`
	Interface  string `json:"interface"`
	PublicKey  string `json:"publicKey"`
	ListenPort int    `json:"listenPort"`
	PeerCount  int    `json:"peerCount"`
	RawOutput  string `json:"rawOutput"`
}

// Manager handles WireGuard configuration and service management.
type Manager struct {
	configPath string // path to /etc/wireguard/wg0.conf
	dataDir    string // WILD_API_DATA_DIR
}

// NewManager creates a new WireGuard manager.
func NewManager(dataDir, configPath string) *Manager {
	if configPath == "" {
		configPath = "/etc/wireguard/wg0.conf"
	}
	return &Manager{configPath: configPath, dataDir: dataDir}
}

// --- Config ---

func (m *Manager) vpnDir() string {
	return filepath.Join(m.dataDir, "vpn")
}

func (m *Manager) configFilePath() string {
	return filepath.Join(m.vpnDir(), "config.yaml")
}

func (m *Manager) secretsFilePath() string {
	return filepath.Join(m.vpnDir(), "secrets.yaml")
}

func (m *Manager) peersDir() string {
	return filepath.Join(m.vpnDir(), "peers")
}

// GetConfig reads server interface config, returning defaults if not yet configured.
func (m *Manager) GetConfig() (*Config, error) {
	data, err := os.ReadFile(m.configFilePath())
	if os.IsNotExist(err) {
		return &Config{
			Enabled:    false,
			ListenPort: 51820,
			Address:    "10.8.0.1/24",
		}, nil
	}
	if err != nil {
		return nil, fmt.Errorf("read vpn config: %w", err)
	}
	var cfg Config
	if err := yaml.Unmarshal(data, &cfg); err != nil {
		return nil, fmt.Errorf("parse vpn config: %w", err)
	}
	return &cfg, nil
}

// ValidateConfig checks a WireGuard config for common errors before saving.
func ValidateConfig(cfg *Config) error {
	if cfg.ListenPort < 1 || cfg.ListenPort > 65535 {
		return fmt.Errorf("listenPort must be 1-65535, got %d", cfg.ListenPort)
	}
	if cfg.Address != "" {
		if _, _, err := net.ParseCIDR(cfg.Address); err != nil {
			return fmt.Errorf("invalid server address CIDR %q: %w", cfg.Address, err)
		}
	}
	if cfg.LanCIDR != "" {
		if _, _, err := net.ParseCIDR(cfg.LanCIDR); err != nil {
			return fmt.Errorf("invalid LAN CIDR %q: %w", cfg.LanCIDR, err)
		}
	}
	return nil
}

// SaveConfig validates and writes the server interface config.
func (m *Manager) SaveConfig(cfg *Config) error {
	if err := ValidateConfig(cfg); err != nil {
		return err
	}
	if err := os.MkdirAll(m.vpnDir(), 0755); err != nil {
		return fmt.Errorf("create vpn dir: %w", err)
	}
	data, err := yaml.Marshal(cfg)
	if err != nil {
		return fmt.Errorf("marshal vpn config: %w", err)
	}
	return storage.WriteFileAtomic(m.configFilePath(), data, 0644)
}

// --- Keys ---

// GenerateKeypair generates a new server keypair using the wg command and saves it.
func (m *Manager) GenerateKeypair() error {
	privKey, err := runWgGenkey()
	if err != nil {
		return fmt.Errorf("generate private key: %w", err)
	}
	pubKey, err := runWgPubkey(privKey)
	if err != nil {
		return fmt.Errorf("derive public key: %w", err)
	}
	if err := os.MkdirAll(m.vpnDir(), 0755); err != nil {
		return fmt.Errorf("create vpn dir: %w", err)
	}
	s := secrets{PrivateKey: privKey, PublicKey: pubKey}
	data, err := yaml.Marshal(s)
	if err != nil {
		return err
	}
	return storage.WriteFileAtomic(m.secretsFilePath(), data, 0600)
}

// GetPublicKey returns the server public key, or empty string if not yet generated.
func (m *Manager) GetPublicKey() string {
	data, err := os.ReadFile(m.secretsFilePath())
	if err != nil {
		return ""
	}
	var s secrets
	if err := yaml.Unmarshal(data, &s); err != nil {
		return ""
	}
	return s.PublicKey
}

func (m *Manager) getSecrets() (*secrets, error) {
	data, err := os.ReadFile(m.secretsFilePath())
	if err != nil {
		return nil, fmt.Errorf("read vpn secrets: %w", err)
	}
	var s secrets
	if err := yaml.Unmarshal(data, &s); err != nil {
		return nil, fmt.Errorf("parse vpn secrets: %w", err)
	}
	return &s, nil
}

// --- Peers ---

// ListPeers returns all configured peers.
func (m *Manager) ListPeers() ([]*Peer, error) {
	entries, err := os.ReadDir(m.peersDir())
	if os.IsNotExist(err) {
		return []*Peer{}, nil
	}
	if err != nil {
		return nil, fmt.Errorf("read peers dir: %w", err)
	}
	var peers []*Peer
	for _, e := range entries {
		if e.IsDir() || !strings.HasSuffix(e.Name(), ".yaml") {
			continue
		}
		p, err := m.readPeerFile(filepath.Join(m.peersDir(), e.Name()))
		if err != nil {
			continue
		}
		peers = append(peers, p)
	}
	return peers, nil
}

// GetPeer returns a single peer by ID.
func (m *Manager) GetPeer(id string) (*Peer, error) {
	return m.readPeerFile(filepath.Join(m.peersDir(), id+".yaml"))
}

func (m *Manager) readPeerFile(path string) (*Peer, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read peer file: %w", err)
	}
	var p Peer
	if err := yaml.Unmarshal(data, &p); err != nil {
		return nil, fmt.Errorf("parse peer file: %w", err)
	}
	return &p, nil
}

// AddPeer generates a new peer keypair, assigns the next available IP in the VPN
// subnet, saves the peer, and returns it.
func (m *Manager) AddPeer(name string) (*Peer, error) {
	cfg, err := m.GetConfig()
	if err != nil {
		return nil, err
	}
	if cfg.Address == "" {
		return nil, fmt.Errorf("server address not configured")
	}

	assignedIP, err := m.nextAvailableIP(cfg.Address)
	if err != nil {
		return nil, fmt.Errorf("assign peer IP: %w", err)
	}

	privKey, err := runWgGenkey()
	if err != nil {
		return nil, fmt.Errorf("generate peer private key: %w", err)
	}
	pubKey, err := runWgPubkey(privKey)
	if err != nil {
		return nil, fmt.Errorf("derive peer public key: %w", err)
	}

	peer := &Peer{
		ID:         uuid.New().String(),
		Name:       name,
		PublicKey:  pubKey,
		PrivateKey: privKey,
		AllowedIPs: assignedIP + "/32",
		CreatedAt:  time.Now().UTC(),
	}

	if err := m.savePeer(peer); err != nil {
		return nil, err
	}
	return peer, nil
}

// DeletePeer removes a peer by ID.
func (m *Manager) DeletePeer(id string) error {
	path := filepath.Join(m.peersDir(), id+".yaml")
	err := os.Remove(path)
	if os.IsNotExist(err) {
		return fmt.Errorf("peer not found")
	}
	return err
}

func (m *Manager) savePeer(p *Peer) error {
	if err := os.MkdirAll(m.peersDir(), 0755); err != nil {
		return fmt.Errorf("create peers dir: %w", err)
	}
	data, err := yaml.Marshal(p)
	if err != nil {
		return err
	}
	return storage.WriteFileAtomic(filepath.Join(m.peersDir(), p.ID+".yaml"), data, 0600)
}

// nextAvailableIP finds the next unused host IP in the given CIDR (skipping the network
// address and the server's own address).
func (m *Manager) nextAvailableIP(serverCIDR string) (string, error) {
	serverIP, ipNet, err := net.ParseCIDR(serverCIDR)
	if err != nil {
		return "", fmt.Errorf("parse server address %q: %w", serverCIDR, err)
	}

	// Collect already-taken IPs (server IP + existing peer IPs).
	taken := map[string]bool{serverIP.String(): true}
	peers, _ := m.ListPeers()
	for _, p := range peers {
		ip, _, _ := net.ParseCIDR(p.AllowedIPs)
		if ip != nil {
			taken[ip.String()] = true
		}
	}

	// Iterate hosts in subnet and find the first untaken one.
	for ip := cloneIP(ipNet.IP); ipNet.Contains(ip); incrementIP(ip) {
		candidate := ip.String()
		if candidate == networkAddr(ipNet) || candidate == broadcastAddr(ipNet) {
			continue
		}
		if !taken[candidate] {
			return candidate, nil
		}
	}
	return "", fmt.Errorf("no available IPs in subnet %s", ipNet)
}

// --- Peer config text ---

// GeneratePeerConfigText builds the wg-quick config block a client needs to connect.
func (m *Manager) GeneratePeerConfigText(peerID string) (string, error) {
	cfg, err := m.GetConfig()
	if err != nil {
		return "", err
	}
	srv, err := m.getSecrets()
	if err != nil {
		return "", fmt.Errorf("server keys not generated yet")
	}
	peer, err := m.GetPeer(peerID)
	if err != nil {
		return "", err
	}

	// Derive the peer's /32 address for the [Interface] section.
	peerIP, _, _ := net.ParseCIDR(peer.AllowedIPs)
	if peerIP == nil {
		peerIP = net.ParseIP(peer.AllowedIPs)
	}

	// Build endpoint string: if Endpoint already contains a port, use it as-is.
	endpoint := cfg.Endpoint
	if endpoint != "" && !strings.Contains(endpoint, ":") {
		endpoint = fmt.Sprintf("%s:%d", endpoint, cfg.ListenPort)
	}

	// Determine AllowedIPs for the client: VPN subnet + LAN CIDR if configured.
	// DNS must be set in the config for Android to activate VPN routing.
	_, vpnNet, _ := net.ParseCIDR(cfg.Address)
	clientAllowedIPs := "0.0.0.0/0, ::/0" // fallback: full tunnel
	if vpnNet != nil {
		clientAllowedIPs = vpnNet.String()
		if cfg.LanCIDR != "" {
			clientAllowedIPs += ", " + cfg.LanCIDR
		}
	}

	var sb strings.Builder
	sb.WriteString("[Interface]\n")
	fmt.Fprintf(&sb, "PrivateKey = %s\n", peer.PrivateKey)
	if peerIP != nil {
		fmt.Fprintf(&sb, "Address = %s/32\n", peerIP)
	}
	if cfg.DNS != "" {
		fmt.Fprintf(&sb, "DNS = %s\n", cfg.DNS)
	}
	sb.WriteString("\n[Peer]\n")
	fmt.Fprintf(&sb, "PublicKey = %s\n", srv.PublicKey)
	if endpoint != "" {
		fmt.Fprintf(&sb, "Endpoint = %s\n", endpoint)
	}
	fmt.Fprintf(&sb, "AllowedIPs = %s\n", clientAllowedIPs)
	sb.WriteString("PersistentKeepalive = 25\n")

	return sb.String(), nil
}

// --- Service management ---

// GetStatus returns the current WireGuard interface status.
func (m *Manager) GetStatus() (*Status, error) {
	cmd := exec.Command("sudo", "wg", "show", "wg0")
	out, err := cmd.CombinedOutput()
	raw := strings.TrimSpace(string(out))

	status := &Status{Interface: "wg0", RawOutput: raw}
	if err != nil {
		// wg show exits non-zero when interface is down.
		status.Running = false
		return status, nil
	}
	status.Running = true

	// Parse peer count from output.
	for line := range strings.SplitSeq(raw, "\n") {
		if strings.HasPrefix(strings.TrimSpace(line), "peer:") {
			status.PeerCount++
		}
	}
	status.PublicKey = m.GetPublicKey()
	if cfg, err := m.GetConfig(); err == nil {
		status.ListenPort = cfg.ListenPort
	}
	return status, nil
}

// Apply writes the wg0.conf file and brings the interface up (or reloads it).
func (m *Manager) Apply() error {
	cfg, err := m.GetConfig()
	if err != nil {
		return err
	}
	srv, err := m.getSecrets()
	if err != nil {
		return fmt.Errorf("server keys not generated yet; run keygen first")
	}
	peers, err := m.ListPeers()
	if err != nil {
		return err
	}

	content := m.renderWgConfig(cfg, srv, peers)

	// Bring down first (using the OLD config so PreDown hooks match what was set up).
	// Ignore error — interface may not exist yet.
	_ = exec.Command("sudo", "wg-quick", "down", "wg0").Run()

	// Now write the new config and bring it up.
	if err := os.MkdirAll(filepath.Dir(m.configPath), 0755); err != nil {
		return fmt.Errorf("create wireguard config dir: %w", err)
	}
	if err := storage.WriteFileAtomic(m.configPath, []byte(content), 0600); err != nil {
		return fmt.Errorf("write wireguard config: %w", err)
	}
	upCmd := exec.Command("sudo", "wg-quick", "up", "wg0")
	if out, err := upCmd.CombinedOutput(); err != nil {
		return fmt.Errorf("wg-quick up: %w\n%s", err, string(out))
	}
	return nil
}

// renderWgConfig generates the /etc/wireguard/wg0.conf content.
func (m *Manager) renderWgConfig(cfg *Config, srv *secrets, peers []*Peer) string {
	var sb strings.Builder
	sb.WriteString("# Managed by Wild Cloud Central — do not edit manually\n\n")
	sb.WriteString("[Interface]\n")
	fmt.Fprintf(&sb, "PrivateKey = %s\n", srv.PrivateKey)
	fmt.Fprintf(&sb, "Address = %s\n", cfg.Address)
	fmt.Fprintf(&sb, "ListenPort = %d\n", cfg.ListenPort)
	// DNS is intentionally omitted from the server config — it belongs only in
	// the client peer config (GeneratePeerConfigText). Setting it here would
	// cause wg-quick to reconfigure the server's own DNS resolver.
	// When LAN CIDR is configured: allow forwarding between wg0 and LAN, and masquerade.
	// Uses iptables (iptables-nft on modern systems) for reliable PostUp/PreDown syntax.
	if cfg.LanCIDR != "" {
		_, vpnNet, _ := net.ParseCIDR(cfg.Address)
		if vpnNet != nil {
			vpnCIDR := vpnNet.String()
			fmt.Fprintf(&sb, "PostUp = iptables -I FORWARD -i wg0 -j ACCEPT; iptables -I FORWARD -o wg0 -m state --state RELATED,ESTABLISHED -j ACCEPT; iptables -t nat -A POSTROUTING -s %s ! -d %s -j MASQUERADE\n", vpnCIDR, vpnCIDR)
			fmt.Fprintf(&sb, "PreDown = iptables -D FORWARD -i wg0 -j ACCEPT 2>/dev/null; iptables -D FORWARD -o wg0 -m state --state RELATED,ESTABLISHED -j ACCEPT 2>/dev/null; iptables -t nat -D POSTROUTING -s %s ! -d %s -j MASQUERADE 2>/dev/null; true\n", vpnCIDR, vpnCIDR)
		}
	}
	for _, p := range peers {
		sb.WriteString("\n[Peer]\n")
		fmt.Fprintf(&sb, "# %s\n", p.Name)
		fmt.Fprintf(&sb, "PublicKey = %s\n", p.PublicKey)
		fmt.Fprintf(&sb, "AllowedIPs = %s\n", p.AllowedIPs)
	}
	return sb.String()
}

// --- wg command helpers ---

func runWgGenkey() (string, error) {
	out, err := exec.Command("wg", "genkey").Output()
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(string(out)), nil
}

func runWgPubkey(privKey string) (string, error) {
	cmd := exec.Command("wg", "pubkey")
	cmd.Stdin = strings.NewReader(privKey)
	out, err := cmd.Output()
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(string(out)), nil
}

// --- IP arithmetic helpers ---

func cloneIP(ip net.IP) net.IP {
	clone := make(net.IP, len(ip))
	copy(clone, ip)
	return clone
}

func incrementIP(ip net.IP) {
	for i := len(ip) - 1; i >= 0; i-- {
		ip[i]++
		if ip[i] != 0 {
			break
		}
	}
}

func networkAddr(n *net.IPNet) string {
	return n.IP.String()
}

func broadcastAddr(n *net.IPNet) string {
	ip := cloneIP(n.IP)
	for i := range ip {
		ip[i] |= ^n.Mask[i]
	}
	return ip.String()
}