DebTik/internal/vpn/wireguard.go

package vpn

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

	"github.com/vishvananda/netlink"
)

const (
	WGInterfacesStorePath = "/etc/debtik/vpn/wireguard_interfaces.json"
	WGPeersStorePath      = "/etc/debtik/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()

	// Autostart all enabled WireGuard interfaces
	for _, iface := range m.interfaces {
		if !iface.Disabled {
			if link, err := netlink.LinkByName(iface.Name); err == nil && link.Type() == "wireguard" {
				continue
			}
			name := iface.Name
			go func(n string) {
				_ = m.startWGInterface(n)
			}(name)
		}
	}

	return m
}

// CleanEndpoint normalizes an endpoint address and port.
// If addr contains a port (e.g. "1.2.3.4:51820", "[2001:db8::1]:51820"), it strips the port
// from addr and returns the clean host and the parsed port.
func CleanEndpoint(addr string, port int) (string, int) {
	addr = strings.TrimSpace(addr)
	if addr == "" || addr == "(none)" {
		return "", 0
	}

	// Remove multiple colons/ports if corrupted (e.g. "10.2.145.253:37616:37616")
	for {
		if h, p, err := net.SplitHostPort(addr); err == nil {
			if parsedPort, err := strconv.Atoi(p); err == nil && parsedPort > 0 {
				port = parsedPort
			}
			addr = h
		} else if strings.Contains(addr, ":") && !strings.HasPrefix(addr, "[") {
			lastColon := strings.LastIndex(addr, ":")
			maybePort := addr[lastColon+1:]
			if parsedPort, err := strconv.Atoi(maybePort); err == nil && parsedPort > 0 {
				port = parsedPort
				addr = addr[:lastColon]
			} else {
				break
			}
		} else {
			break
		}
	}

	return addr, port
}

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)
	}
	// Sanitize endpoints for all stored peers
	for i := range m.peers {
		m.peers[i].EndpointAddress, m.peers[i].EndpointPort = CleanEndpoint(m.peers[i].EndpointAddress, m.peers[i].EndpointPort)
	}
}

func (m *WireGuardManager) save() error {
	_ = os.MkdirAll(filepath.Dir(WGInterfacesStorePath), 0755)
	if data, err := json.MarshalIndent(m.interfaces, "", "  "); err != nil {
		return err
	} else if err = os.WriteFile(WGInterfacesStorePath, data, 0600); err != nil {
		return err
	}
	if data, err := json.MarshalIndent(m.peers, "", "  "); err != nil {
		return err
	} else if err = os.WriteFile(WGPeersStorePath, data, 0600); err != nil {
		return err
	}
	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))

	pubKey, err := m.CalcPublicKey(privKey)
	if err != nil {
		return "", "", err
	}

	return privKey, pubKey, nil
}

func (m *WireGuardManager) CalcPublicKey(privateKey string) (string, error) {
	pubCmd := exec.Command("wg", "pubkey")
	pubCmd.Stdin = strings.NewReader(strings.TrimSpace(privateKey))
	pubOut, err := pubCmd.Output()
	if err != nil {
		return "", fmt.Errorf("failed to derive public key: %w", err)
	}
	return strings.TrimSpace(string(pubOut)), 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(req SaveWireGuardInterfaceRequest) error {
	m.mu.Lock()
	defer m.mu.Unlock()

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

	var old *WireGuardInterface
	for i := range m.interfaces {
		if m.interfaces[i].Name == req.Name || (req.OrigName != "" && m.interfaces[i].Name == req.OrigName) {
			copy := m.interfaces[i]
			old = &copy
			break
		}
	}
	if req.PrivateKey == "" && old != nil {
		req.PrivateKey = old.PrivateKey
		req.PublicKey = old.PublicKey
	}
	// Auto generate or derive public key
	if req.PrivateKey != "" {
		pub, err := m.CalcPublicKey(req.PrivateKey)
		if err != nil {
			return err
		}
		if pub == "" {
			return fmt.Errorf("invalid private key")
		}
		req.PublicKey = pub
	} else {
		priv, pub, err := m.GenerateKeyPair()
		if err != nil {
			return err
		}
		req.PrivateKey = priv
		req.PublicKey = pub
	}

	// Support interface renaming: clean up old interface if renamed
	if req.OrigName != "" && req.OrigName != req.Name {
		_ = m.stopWGInterface(req.OrigName)
		_ = os.Remove(fmt.Sprintf("/etc/wireguard/%s.conf", req.OrigName))

		// Update all peers associated with OrigName
		for i := range m.peers {
			if m.peers[i].Interface == req.OrigName {
				m.peers[i].Interface = req.Name
			}
		}

		// Remove old interface from slice
		updated := make([]WireGuardInterface, 0)
		for _, iface := range m.interfaces {
			if iface.Name != req.OrigName {
				updated = append(updated, iface)
			}
		}
		m.interfaces = updated
	}

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

	if err := m.save(); err != nil {
		return err
	}
	if err := m.generateWGConfigFile(req.Name); err != nil {
		return err
	}

	if old != nil && old.Name == req.Name && old.Disabled == req.Disabled && old.PrivateKey == req.PrivateKey && old.ListenPort == req.ListenPort && old.Address == req.Address && old.MTU == req.MTU {
		return nil
	}
	if !req.Disabled {
		return m.startWGInterface(req.Name)
	}
	return m.stopWGInterface(req.Name)
}

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
			if err := m.save(); err != nil {
				return err
			}
			if enabled {
				return m.startWGInterface(name)
			} else {
				return m.stopWGInterface(name)
			}
		}
	}
	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 {
	if err := m.generateWGConfigFile(name); err != nil {
		return err
	}
	_ = 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)
	}
	_ = exec.Command("systemctl", "enable", fmt.Sprintf("wg-quick@%s", name)).Run()
	return nil
}

func (m *WireGuardManager) stopWGInterface(name string) error {
	_ = exec.Command("systemctl", "disable", fmt.Sprintf("wg-quick@%s", name)).Run()
	if _, err := netlink.LinkByName(name); err != nil {
		if _, absent := err.(netlink.LinkNotFoundError); absent {
			return nil
		}
		return err
	}
	cmd := exec.Command("wg-quick", "down", name)
	if out, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("failed to stop %s: %s: %w", name, out, err)
	}
	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"
						}

						epHost, epPort := CleanEndpoint(endpoint, 0)
						dumpMap[iface.Name][pubKey] = WireGuardPeer{
							EndpointAddress: epHost,
							EndpointPort:    epPort,
							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 {
				if stats.EndpointAddress != "" {
					item.EndpointAddress = stats.EndpointAddress
					if stats.EndpointPort > 0 {
						item.EndpointPort = stats.EndpointPort
					}
				}
				item.LatestHandshake = stats.LatestHandshake
				item.TransferRxBytes = stats.TransferRxBytes
				item.TransferTxBytes = stats.TransferTxBytes
				item.TransferRxHuman = stats.TransferRxHuman
				item.TransferTxHuman = stats.TransferTxHuman
			}
		}
		item.EndpointAddress, item.EndpointPort = CleanEndpoint(item.EndpointAddress, item.EndpointPort)
		result = append(result, item)
	}
	return result
}

func (m *WireGuardManager) SavePeer(req SaveWireGuardPeerRequest) error {
	m.mu.Lock()
	defer m.mu.Unlock()

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

	// Validate that peer public key is not the same as the interface's own public key
	for _, iface := range m.interfaces {
		if iface.Name == req.Interface && iface.PublicKey == req.PublicKey {
			return fmt.Errorf("peer public key cannot be identical to interface %s public key", iface.Name)
		}
	}

	// Handle modifying an existing peer's public key or moving to a different interface
	if req.OrigPublicKey != "" && (req.OrigPublicKey != req.PublicKey || (req.OrigInterface != "" && req.OrigInterface != req.Interface)) {
		oldIface := req.OrigInterface
		if oldIface == "" {
			oldIface = req.Interface
		}
		// Remove old peer from runtime
		if _, err := netlink.LinkByName(oldIface); err == nil {
			if out, err := exec.Command("wg", "set", oldIface, "peer", req.OrigPublicKey, "remove").CombinedOutput(); err != nil {
				return fmt.Errorf("wg command failed: %s: %w", out, err)
			}
		}

		// Remove old peer entry
		updated := make([]WireGuardPeer, 0)
		for _, p := range m.peers {
			if !(p.Interface == oldIface && p.PublicKey == req.OrigPublicKey) {
				updated = append(updated, p)
			}
		}
		m.peers = updated

		if oldIface != req.Interface {
			_ = m.generateWGConfigFile(oldIface)
		}
	}

	req.EndpointAddress, req.EndpointPort = CleanEndpoint(req.EndpointAddress, req.EndpointPort)
	req.WireGuardPeer.EndpointAddress = req.EndpointAddress
	req.WireGuardPeer.EndpointPort = req.EndpointPort

	clearEndpoint := false
	for _, old := range m.peers {
		if old.Interface == req.Interface && old.PublicKey == req.PublicKey && old.EndpointAddress != "" && req.EndpointAddress == "" {
			clearEndpoint = true
		}
	}
	found := false
	for i, existing := range m.peers {
		if existing.PublicKey == req.PublicKey && existing.Interface == req.Interface {
			m.peers[i] = req.WireGuardPeer
			found = true
			break
		}
	}
	if !found {
		m.peers = append(m.peers, req.WireGuardPeer)
	}

	if err := m.save(); err != nil {
		return err
	}
	if err := m.generateWGConfigFile(req.Interface); err != nil {
		return err
	}

	// Apply dynamically to kernel if running
	if _, err := netlink.LinkByName(req.Interface); err == nil {
		if req.Disabled {
			if out, err := exec.Command("wg", "set", req.Interface, "peer", req.PublicKey, "remove").CombinedOutput(); err != nil {
				return fmt.Errorf("wg command failed: %s: %w", out, err)
			}
		} else {
			if clearEndpoint {
				if out, err := exec.Command("wg", "set", req.Interface, "peer", req.PublicKey, "remove").CombinedOutput(); err != nil {
					return fmt.Errorf("clear WireGuard endpoint: %s: %w", out, err)
				}
			}
			args := []string{"set", req.Interface, "peer", req.PublicKey}
			args = append(args, "allowed-ips", strings.Join(req.AllowedIPs, ","))
			if req.EndpointAddress != "" {
				ep := req.EndpointAddress
				if req.EndpointPort > 0 {
					ep = net.JoinHostPort(req.EndpointAddress, strconv.Itoa(req.EndpointPort))
				}
				args = append(args, "endpoint", ep)
			}
			args = append(args, "persistent-keepalive", strconv.Itoa(req.PersistentKeepalive))
			if true {
				pskFile, err := os.CreateTemp("", "wg-psk-*")
				if err == nil {
					_, _ = pskFile.WriteString(req.PresharedKey)
					_ = pskFile.Close()
					args = append(args, "preshared-key", pskFile.Name())
					defer os.Remove(pskFile.Name())
				}
			}
			if out, err := exec.Command("wg", args...).CombinedOutput(); err != nil {
				return fmt.Errorf("wg set failed: %s: %w", out, err)
			}
		}
	}

	return nil
}

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

	found := false
	for i, p := range m.peers {
		if p.Interface == ifaceName && p.PublicKey == pubKey {
			m.peers[i].Disabled = !enabled
			found = true
			break
		}
	}
	if !found {
		return fmt.Errorf("peer with public key %s not found on interface %s", pubKey, ifaceName)
	}

	if err := m.save(); err != nil {
		return err
	}
	_ = m.generateWGConfigFile(ifaceName)

	if _, err := netlink.LinkByName(ifaceName); err == nil {
		if !enabled {
			if out, err := exec.Command("wg", "set", ifaceName, "peer", pubKey, "remove").CombinedOutput(); err != nil {
				return fmt.Errorf("wg command failed: %s: %w", out, err)
			}
		} else {
			for _, p := range m.peers {
				if p.Interface == ifaceName && p.PublicKey == pubKey {
					args := []string{"set", ifaceName, "peer", pubKey}
					if len(p.AllowedIPs) > 0 {
						args = append(args, "allowed-ips", strings.Join(p.AllowedIPs, ","))
					}
					epHost, epPort := CleanEndpoint(p.EndpointAddress, p.EndpointPort)
					if epHost != "" {
						ep := epHost
						if epPort > 0 {
							ep = net.JoinHostPort(epHost, strconv.Itoa(epPort))
						}
						args = append(args, "endpoint", ep)
					}
					if p.PersistentKeepalive > 0 {
						args = append(args, "persistent-keepalive", strconv.Itoa(p.PersistentKeepalive))
					}
					if p.PresharedKey != "" {
						pskFile, err := os.CreateTemp("", "wg-psk-*")
						if err == nil {
							_, _ = pskFile.WriteString(p.PresharedKey)
							_ = pskFile.Close()
							args = append(args, "preshared-key", pskFile.Name())
							defer os.Remove(pskFile.Name())
						}
					}
					if out, err := exec.Command("wg", args...).CombinedOutput(); err != nil {
						return fmt.Errorf("wg command failed: %s: %w", out, err)
					}
					break
				}
			}
		}
	}
	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

	if err := m.save(); err != nil {
		return err
	}
	_ = m.generateWGConfigFile(ifaceName)

	// Remove from runtime
	if _, err := netlink.LinkByName(ifaceName); err == nil {
		if out, err := exec.Command("wg", "set", ifaceName, "peer", pubKey, "remove").CombinedOutput(); err != nil {
			return fmt.Errorf("wg command failed: %s: %w", out, err)
		}
	}
	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))
	if targetIface.Comment != "" {
		sb.WriteString(fmt.Sprintf("# Comment: %s\n", targetIface.Comment))
	}
	sb.WriteString("[Interface]\n")
	sb.WriteString(fmt.Sprintf("PrivateKey = %s\n", targetIface.PrivateKey))
	if targetIface.Address != "" && targetIface.Address != "0.0.0.0/0" {
		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("Table = off\n")
	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, ", ")))
			}
			epHost, epPort := CleanEndpoint(p.EndpointAddress, p.EndpointPort)
			if epHost != "" {
				ep := epHost
				if epPort > 0 {
					ep = net.JoinHostPort(epHost, strconv.Itoa(epPort))
				}
				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)
}

// GenerateClientConfig creates a ready-to-use client .conf for mobile/desktop
func (m *WireGuardManager) GenerateClientConfig(ifaceName, clientPrivKey, clientAddress, serverEndpoint, dns string) (string, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	var targetIface *WireGuardInterface
	for _, i := range m.interfaces {
		if i.Name == ifaceName {
			targetIface = &i
			break
		}
	}
	if targetIface == nil {
		return "", fmt.Errorf("interface %s not found", ifaceName)
	}

	if dns == "" {
		dns = "1.1.1.1, 8.8.8.8"
	}
	if serverEndpoint == "" {
		serverEndpoint = fmt.Sprintf("YOUR_ROUTER_IP:%d", targetIface.ListenPort)
	}

	var sb strings.Builder
	sb.WriteString("[Interface]\n")
	sb.WriteString(fmt.Sprintf("PrivateKey = %s\n", clientPrivKey))
	sb.WriteString(fmt.Sprintf("Address = %s\n", clientAddress))
	sb.WriteString(fmt.Sprintf("DNS = %s\n", dns))
	sb.WriteString("\n[Peer]\n")
	sb.WriteString(fmt.Sprintf("PublicKey = %s\n", targetIface.PublicKey))
	sb.WriteString(fmt.Sprintf("Endpoint = %s\n", serverEndpoint))
	sb.WriteString("AllowedIPs = 0.0.0.0/0, ::/0\n")
	sb.WriteString("PersistentKeepalive = 25\n")

	return sb.String(), nil
}