DebTik/internal/vpn/wireguard.go

package vpn

import (
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/vishvananda/netlink"
)

const (
	WGInterfacesStorePath = "/etc/proxmox-router/vpn/wireguard_interfaces.json"
	WGPeersStorePath      = "/etc/proxmox-router/vpn/wireguard_peers.json"
)

type WireGuardManager struct {
	mu         sync.Mutex
	interfaces []WireGuardInterface
	peers      []WireGuardPeer
}

func NewWireGuardManager() *WireGuardManager {
	m := &WireGuardManager{
		interfaces: make([]WireGuardInterface, 0),
		peers:      make([]WireGuardPeer, 0),
	}
	m.load()
	return m
}

func (m *WireGuardManager) load() {
	if data, err := os.ReadFile(WGInterfacesStorePath); err == nil {
		_ = json.Unmarshal(data, &m.interfaces)
	}
	if data, err := os.ReadFile(WGPeersStorePath); err == nil {
		_ = json.Unmarshal(data, &m.peers)
	}
	if m.interfaces == nil {
		m.interfaces = make([]WireGuardInterface, 0)
	}
	if m.peers == nil {
		m.peers = make([]WireGuardPeer, 0)
	}
}

func (m *WireGuardManager) save() error {
	_ = os.MkdirAll(filepath.Dir(WGInterfacesStorePath), 0755)
	if data, err := json.MarshalIndent(m.interfaces, "", "  "); err == nil {
		_ = os.WriteFile(WGInterfacesStorePath, data, 0600)
	}
	if data, err := json.MarshalIndent(m.peers, "", "  "); err == nil {
		_ = os.WriteFile(WGPeersStorePath, data, 0600)
	}
	return nil
}

func (m *WireGuardManager) GenerateKeyPair() (string, string, error) {
	privCmd := exec.Command("wg", "genkey")
	privOut, err := privCmd.Output()
	if err != nil {
		return "", "", fmt.Errorf("failed to generate private key: %w", err)
	}
	privKey := strings.TrimSpace(string(privOut))

	pubCmd := exec.Command("wg", "pubkey")
	pubCmd.Stdin = strings.NewReader(privKey)
	pubOut, err := pubCmd.Output()
	if err != nil {
		return "", "", fmt.Errorf("failed to generate public key: %w", err)
	}
	pubKey := strings.TrimSpace(string(pubOut))

	return privKey, pubKey, nil
}

// ListInterfaces returns all WireGuard interfaces with MikroTik running flags
func (m *WireGuardManager) ListInterfaces() []WireGuardInterface {
	m.mu.Lock()
	defer m.mu.Unlock()

	res := make([]WireGuardInterface, len(m.interfaces))
	for i, iface := range m.interfaces {
		res[i] = iface
		// Check running state via netlink
		if _, err := netlink.LinkByName(iface.Name); err == nil {
			res[i].Running = true
		} else {
			res[i].Running = false
		}

		// Count attached peers
		pCount := 0
		for _, p := range m.peers {
			if p.Interface == iface.Name {
				pCount++
			}
		}
		res[i].PeersCount = pCount
	}
	return res
}

func (m *WireGuardManager) SaveInterface(iface WireGuardInterface) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if iface.Name == "" {
		return fmt.Errorf("interface name is required (e.g. wg0)")
	}
	if iface.ListenPort <= 0 {
		iface.ListenPort = 51820
	}
	if iface.Address == "" {
		iface.Address = "10.10.0.1/24"
	}
	if iface.MTU <= 0 {
		iface.MTU = 1420
	}

	// Auto generate keys if not provided
	if iface.PrivateKey == "" {
		priv, pub, err := m.GenerateKeyPair()
		if err == nil {
			iface.PrivateKey = priv
			iface.PublicKey = pub
		}
	}

	found := false
	for i, existing := range m.interfaces {
		if existing.Name == iface.Name {
			m.interfaces[i] = iface
			found = true
			break
		}
	}
	if !found {
		m.interfaces = append(m.interfaces, iface)
	}

	_ = m.save()
	_ = m.generateWGConfigFile(iface.Name)

	if !iface.Disabled {
		go m.startWGInterface(iface.Name)
	} else {
		go m.stopWGInterface(iface.Name)
	}

	return nil
}

func (m *WireGuardManager) DeleteInterface(name string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	_ = m.stopWGInterface(name)
	_ = os.Remove(fmt.Sprintf("/etc/wireguard/%s.conf", name))

	updatedIfaces := make([]WireGuardInterface, 0)
	for _, iface := range m.interfaces {
		if iface.Name != name {
			updatedIfaces = append(updatedIfaces, iface)
		}
	}
	m.interfaces = updatedIfaces

	// Remove associated peers
	updatedPeers := make([]WireGuardPeer, 0)
	for _, p := range m.peers {
		if p.Interface != name {
			updatedPeers = append(updatedPeers, p)
		}
	}
	m.peers = updatedPeers

	return m.save()
}

func (m *WireGuardManager) SetInterfaceEnabled(name string, enabled bool) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	for i, iface := range m.interfaces {
		if iface.Name == name {
			m.interfaces[i].Disabled = !enabled
			_ = m.save()
			if enabled {
				go m.startWGInterface(name)
			} else {
				go m.stopWGInterface(name)
			}
			return nil
		}
	}
	return fmt.Errorf("interface %q not found", name)
}

func (m *WireGuardManager) StartInterface(name string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.startWGInterface(name)
}

func (m *WireGuardManager) StopInterface(name string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.stopWGInterface(name)
}

func (m *WireGuardManager) startWGInterface(name string) error {
	_ = m.generateWGConfigFile(name)
	_ = exec.Command("wg-quick", "down", name).Run()

	cmd := exec.Command("wg-quick", "up", name)
	if out, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("failed to start %s: %s: %w", name, string(out), err)
	}
	return nil
}

func (m *WireGuardManager) stopWGInterface(name string) error {
	cmd := exec.Command("wg-quick", "down", name)
	_ = cmd.Run()
	return nil
}

// ----------------- Peers Management (MikroTik style) -----------------

func (m *WireGuardManager) ListPeers(ifaceFilter string) []WireGuardPeer {
	m.mu.Lock()
	defer m.mu.Unlock()

	// Parse runtime dump from `wg show <iface> dump`
	dumpMap := make(map[string]map[string]WireGuardPeer) // iface -> pubkey -> peer stats

	for _, iface := range m.interfaces {
		if _, err := netlink.LinkByName(iface.Name); err == nil {
			cmd := exec.Command("wg", "show", iface.Name, "dump")
			out, err := cmd.Output()
			if err == nil {
				dumpMap[iface.Name] = make(map[string]WireGuardPeer)
				lines := strings.Split(strings.TrimSpace(string(out)), "\n")
				for idx, line := range lines {
					if idx == 0 {
						continue // skip interface line
					}
					parts := strings.Fields(line)
					if len(parts) >= 8 {
						pubKey := parts[0]
						endpoint := parts[2]
						if endpoint == "(none)" {
							endpoint = ""
						}
						handshakeUnix, _ := strconv.ParseInt(parts[4], 10, 64)
						rx, _ := strconv.ParseUint(parts[5], 10, 64)
						tx, _ := strconv.ParseUint(parts[6], 10, 64)

						handshakeStr := "Never"
						if handshakeUnix > 0 {
							handshakeStr = time.Since(time.Unix(handshakeUnix, 0)).Round(time.Second).String() + " ago"
						}

						dumpMap[iface.Name][pubKey] = WireGuardPeer{
							EndpointAddress: endpoint,
							LatestHandshake: handshakeStr,
							TransferRxBytes: rx,
							TransferTxBytes: tx,
							TransferRxHuman: formatTrafficBytes(rx),
							TransferTxHuman: formatTrafficBytes(tx),
						}
					}
				}
			}
		}
	}

	result := make([]WireGuardPeer, 0)
	for _, p := range m.peers {
		if ifaceFilter != "" && p.Interface != ifaceFilter {
			continue
		}
		item := p
		// Overlay runtime stats if available
		if ifaceDumps, ok := dumpMap[p.Interface]; ok {
			if stats, has := ifaceDumps[p.PublicKey]; has {
				item.EndpointAddress = stats.EndpointAddress
				item.LatestHandshake = stats.LatestHandshake
				item.TransferRxBytes = stats.TransferRxBytes
				item.TransferTxBytes = stats.TransferTxBytes
				item.TransferRxHuman = stats.TransferRxHuman
				item.TransferTxHuman = stats.TransferTxHuman
			}
		}
		result = append(result, item)
	}
	return result
}

func (m *WireGuardManager) SavePeer(peer WireGuardPeer) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if peer.Interface == "" {
		return fmt.Errorf("target WireGuard interface is required")
	}
	if peer.PublicKey == "" {
		return fmt.Errorf("public key is required")
	}

	found := false
	for i, existing := range m.peers {
		if existing.PublicKey == peer.PublicKey && existing.Interface == peer.Interface {
			m.peers[i] = peer
			found = true
			break
		}
	}
	if !found {
		m.peers = append(m.peers, peer)
	}

	_ = m.save()
	_ = m.generateWGConfigFile(peer.Interface)

	// Apply dynamically to kernel if running
	if _, err := netlink.LinkByName(peer.Interface); err == nil && !peer.Disabled {
		args := []string{"set", peer.Interface, "peer", peer.PublicKey}
		if len(peer.AllowedIPs) > 0 {
			args = append(args, "allowed-ips", strings.Join(peer.AllowedIPs, ","))
		}
		if peer.EndpointAddress != "" {
			ep := peer.EndpointAddress
			if peer.EndpointPort > 0 {
				ep = fmt.Sprintf("%s:%d", peer.EndpointAddress, peer.EndpointPort)
			}
			args = append(args, "endpoint", ep)
		}
		if peer.PersistentKeepalive > 0 {
			args = append(args, "persistent-keepalive", strconv.Itoa(peer.PersistentKeepalive))
		}
		_ = exec.Command("wg", args...).Run()
	}

	return nil
}

func (m *WireGuardManager) DeletePeer(ifaceName, pubKey string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	updated := make([]WireGuardPeer, 0)
	for _, p := range m.peers {
		if p.Interface == ifaceName && p.PublicKey == pubKey {
			continue
		}
		updated = append(updated, p)
	}
	m.peers = updated

	_ = m.save()
	_ = m.generateWGConfigFile(ifaceName)

	// Remove from runtime
	if _, err := netlink.LinkByName(ifaceName); err == nil {
		_ = exec.Command("wg", "set", ifaceName, "peer", pubKey, "remove").Run()
	}
	return nil
}

func (m *WireGuardManager) generateWGConfigFile(ifaceName string) error {
	var targetIface *WireGuardInterface
	for _, i := range m.interfaces {
		if i.Name == ifaceName {
			targetIface = &i
			break
		}
	}
	if targetIface == nil {
		return nil
	}

	_ = os.MkdirAll("/etc/wireguard", 0700)

	var sb strings.Builder
	sb.WriteString(fmt.Sprintf("# MikroTik-Style WireGuard Interface: %s\n", targetIface.Name))
	sb.WriteString("[Interface]\n")
	sb.WriteString(fmt.Sprintf("PrivateKey = %s\n", targetIface.PrivateKey))
	sb.WriteString(fmt.Sprintf("Address = %s\n", targetIface.Address))
	sb.WriteString(fmt.Sprintf("ListenPort = %d\n", targetIface.ListenPort))
	if targetIface.MTU > 0 {
		sb.WriteString(fmt.Sprintf("MTU = %d\n", targetIface.MTU))
	}
	sb.WriteString("\n")

	for _, p := range m.peers {
		if p.Interface == ifaceName && !p.Disabled {
			sb.WriteString(fmt.Sprintf("# Peer: %s\n", p.Name))
			sb.WriteString("[Peer]\n")
			sb.WriteString(fmt.Sprintf("PublicKey = %s\n", p.PublicKey))
			if len(p.AllowedIPs) > 0 {
				sb.WriteString(fmt.Sprintf("AllowedIPs = %s\n", strings.Join(p.AllowedIPs, ", ")))
			}
			if p.EndpointAddress != "" {
				ep := p.EndpointAddress
				if p.EndpointPort > 0 {
					ep = fmt.Sprintf("%s:%d", p.EndpointAddress, p.EndpointPort)
				}
				sb.WriteString(fmt.Sprintf("Endpoint = %s\n", ep))
			}
			if p.PersistentKeepalive > 0 {
				sb.WriteString(fmt.Sprintf("PersistentKeepalive = %d\n", p.PersistentKeepalive))
			}
			if p.PresharedKey != "" {
				sb.WriteString(fmt.Sprintf("PresharedKey = %s\n", p.PresharedKey))
			}
			sb.WriteString("\n")
		}
	}

	path := fmt.Sprintf("/etc/wireguard/%s.conf", ifaceName)
	return os.WriteFile(path, []byte(sb.String()), 0600)
}