wg-admin/internal/wireguard/client.go

package wireguard

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/calmcacil/wg-admin/internal/backup"
	"github.com/calmcacil/wg-admin/internal/config"
)

// Client represents a WireGuard peer configuration
type Client struct {
	Name       string // Client name extracted from filename
	IPv4       string // IPv4 address from AllowedIPs
	IPv6       string // IPv6 address from AllowedIPs
	PublicKey  string // WireGuard public key
	HasPSK     bool   // Whether PresharedKey is configured
	ConfigPath string // Path to the client config file
}

// ParseClientConfig parses a single WireGuard client configuration file
func ParseClientConfig(path string) (*Client, error) {
	// Read the file
	content, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("failed to read config file %s: %w", path, err)
	}

	// Extract client name from filename
	base := filepath.Base(path)
	name := strings.TrimPrefix(base, "client-")
	name = strings.TrimSuffix(name, ".conf")

	if name == "" || name == "client-" {
		return nil, fmt.Errorf("invalid client filename: %s", base)
	}

	client := &Client{
		Name:       name,
		ConfigPath: path,
	}

	// Parse the INI-style config
	inPeerSection := false
	hasPublicKey := false

	for i, line := range strings.Split(string(content), "\n") {
		line = strings.TrimSpace(line)

		// Skip empty lines and comments
		if line == "" || strings.HasPrefix(line, "#") || strings.HasPrefix(line, ";") {
			continue
		}

		// Check for Peer section
		if line == "[Peer]" {
			inPeerSection = true
			continue
		}

		// Parse key-value pairs within Peer section
		if inPeerSection {
			parts := strings.SplitN(line, "=", 2)
			if len(parts) != 2 {
				log.Printf("Warning: malformed line %d in %s: %s", i+1, path, line)
				continue
			}

			key := strings.TrimSpace(parts[0])
			value := strings.TrimSpace(parts[1])

			switch key {
			case "PublicKey":
				client.PublicKey = value
				hasPublicKey = true
			case "PresharedKey":
				client.HasPSK = true
			case "AllowedIPs":
				if err := parseAllowedIPs(client, value); err != nil {
					log.Printf("Warning: %v (file: %s, line: %d)", err, path, i+1)
				}
			}
		}
	}

	// Validate required fields
	if !hasPublicKey {
		return nil, fmt.Errorf("missing required PublicKey in %s", path)
	}

	if client.IPv4 == "" && client.IPv6 == "" {
		return nil, fmt.Errorf("no valid IP addresses found in AllowedIPs in %s", path)
	}

	return client, nil
}

// parseAllowedIPs extracts IPv4 and IPv6 addresses from AllowedIPs value
func parseAllowedIPs(client *Client, allowedIPs string) error {
	// AllowedIPs format: "ipv4/32, ipv6/128"
	addresses := strings.Split(allowedIPs, ",")

	for _, addr := range addresses {
		addr = strings.TrimSpace(addr)
		if addr == "" {
			continue
		}

		// Split IP from CIDR suffix
		parts := strings.Split(addr, "/")
		if len(parts) != 2 {
			return fmt.Errorf("invalid AllowedIP format: %s", addr)
		}

		ip := strings.TrimSpace(parts[0])

		// Detect if IPv4 or IPv6 based on presence of colon
		if strings.Contains(ip, ":") {
			client.IPv6 = ip
		} else {
			client.IPv4 = ip
		}
	}

	return nil
}

// ListClients finds and parses all client configurations from /etc/wireguard/conf.d/
func ListClients() ([]Client, error) {
	configDir := "/etc/wireguard/conf.d"

	// Check if directory exists
	if _, err := os.Stat(configDir); os.IsNotExist(err) {
		return nil, fmt.Errorf("wireguard config directory does not exist: %s", configDir)
	}

	// Find all client-*.conf files
	pattern := filepath.Join(configDir, "client-*.conf")
	matches, err := filepath.Glob(pattern)
	if err != nil {
		return nil, fmt.Errorf("failed to find client config files: %w", err)
	}

	if len(matches) == 0 {
		return []Client{}, nil // No clients found, return empty slice
	}

	// Parse each config file
	var clients []Client
	var parseErrors []string

	for _, match := range matches {
		client, err := ParseClientConfig(match)
		if err != nil {
			parseErrors = append(parseErrors, err.Error())
			log.Printf("Warning: failed to parse %s: %v", match, err)
			continue
		}
		clients = append(clients, *client)
	}

	// If all files failed to parse, return an error
	if len(clients) == 0 && len(parseErrors) > 0 {
		return nil, fmt.Errorf("failed to parse any client configs: %s", strings.Join(parseErrors, "; "))
	}

	return clients, nil
}

// GetClientConfigContent reads the raw configuration content for a client
func GetClientConfigContent(name string) (string, error) {
	configDir := "/etc/wireguard/clients"
	configPath := filepath.Join(configDir, fmt.Sprintf("%s.conf", name))

	content, err := os.ReadFile(configPath)
	if err != nil {
		if os.IsNotExist(err) {
			return "", fmt.Errorf("client config not found: %s", configPath)
		}
		return "", fmt.Errorf("failed to read client config %s: %w", configPath, err)
	}

	return string(content), nil
}

// DeleteClient removes a WireGuard client configuration and associated files
func DeleteClient(name string) error {
	// First, find the client config to get public key for removal from interface
	configDir := "/etc/wireguard/conf.d"
	configPath := filepath.Join(configDir, fmt.Sprintf("client-%s.conf", name))

	client, err := ParseClientConfig(configPath)
	if err != nil {
		return fmt.Errorf("failed to parse client config for deletion: %w", err)
	}

	log.Printf("Deleting client: %s (public key: %s)", name, client.PublicKey)

	// Create backup before deletion
	backupPath, err := backup.BackupConfig(fmt.Sprintf("delete-%s", name))
	if err != nil {
		log.Printf("Warning: failed to create backup before deletion: %v", err)
	} else {
		log.Printf("Created backup: %s", backupPath)
	}

	// Remove peer from WireGuard interface using wg command
	if err := removePeerFromInterface(client.PublicKey); err != nil {
		log.Printf("Warning: failed to remove peer from interface: %v", err)
	}

	// Remove client config from /etc/wireguard/conf.d/
	if err := os.Remove(configPath); err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("failed to remove client config %s: %w", configPath, err)
	}
	log.Printf("Removed client config: %s", configPath)

	// Remove client files from /etc/wireguard/clients/
	clientsDir := "/etc/wireguard/clients"
	clientFile := filepath.Join(clientsDir, fmt.Sprintf("%s.conf", name))
	if err := os.Remove(clientFile); err != nil && !os.IsNotExist(err) {
		log.Printf("Warning: failed to remove client file %s: %v", clientFile, err)
	} else {
		log.Printf("Removed client file: %s", clientFile)
	}

	// Remove QR code PNG if it exists
	qrFile := filepath.Join(clientsDir, fmt.Sprintf("%s.png", name))
	if err := os.Remove(qrFile); err != nil && !os.IsNotExist(err) {
		log.Printf("Warning: failed to remove QR code %s: %v", qrFile, err)
	} else {
		log.Printf("Removed QR code: %s", qrFile)
	}

	log.Printf("Successfully deleted client: %s", name)
	return nil
}

// removePeerFromInterface removes a peer from the WireGuard interface
func removePeerFromInterface(publicKey string) error {
	// Use wg command to remove peer
	cmd := exec.Command("wg", "set", "wg0", "peer", publicKey, "remove")
	output, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("wg set peer remove failed: %w, output: %s", err, string(output))
	}
	return nil
}

// CreateClient creates a new WireGuard client configuration
func CreateClient(name, dns string, usePSK bool) error {
	log.Printf("Creating client: %s (PSK: %v)", name, usePSK)

	// Create backup before creating client
	backupPath, err := backup.BackupConfig(fmt.Sprintf("create-%s", name))
	if err != nil {
		log.Printf("Warning: failed to create backup before creating client: %v", err)
	} else {
		log.Printf("Created backup: %s", backupPath)
	}

	// Generate keys
	privateKey, publicKey, err := generateKeyPair()
	if err != nil {
		return fmt.Errorf("failed to generate key pair: %w", err)
	}

	var psk string
	if usePSK {
		psk, err = generatePSK()
		if err != nil {
			return fmt.Errorf("failed to generate PSK: %w", err)
		}
	}

	// Get next available IP addresses
	cfg, err := config.LoadConfig()
	if err != nil {
		return fmt.Errorf("failed to load config: %w", err)
	}

	clients, err := ListClients()
	if err != nil {
		return fmt.Errorf("failed to list existing clients: %w", err)
	}

	ipv4, err := getNextAvailableIP(cfg.VPNIPv4Range, clients)
	if err != nil {
		return fmt.Errorf("failed to get IPv4 address: %w", err)
	}

	ipv6, err := getNextAvailableIP(cfg.VPNIPv6Range, clients)
	if err != nil {
		return fmt.Errorf("failed to get IPv6 address: %w", err)
	}

	// Create server config
	serverConfigPath := fmt.Sprintf("/etc/wireguard/conf.d/client-%s.conf", name)
	serverConfig, err := generateServerConfig(name, publicKey, ipv4, ipv6, psk, cfg)
	if err != nil {
		return fmt.Errorf("failed to generate server config: %w", err)
	}

	if err := os.WriteFile(serverConfigPath, []byte(serverConfig), 0600); err != nil {
		return fmt.Errorf("failed to write server config: %w", err)
	}
	log.Printf("Created server config: %s", serverConfigPath)

	// Create client config
	clientsDir := "/etc/wireguard/clients"
	if err := os.MkdirAll(clientsDir, 0700); err != nil {
		return fmt.Errorf("failed to create clients directory: %w", err)
	}

	clientConfigPath := filepath.Join(clientsDir, fmt.Sprintf("%s.conf", name))
	clientConfig, err := generateClientConfig(name, privateKey, ipv4, ipv6, dns, cfg)
	if err != nil {
		return fmt.Errorf("failed to generate client config: %w", err)
	}

	if err := os.WriteFile(clientConfigPath, []byte(clientConfig), 0600); err != nil {
		return fmt.Errorf("failed to write client config: %w", err)
	}
	log.Printf("Created client config: %s", clientConfigPath)

	// Add peer to WireGuard interface
	if err := addPeerToInterface(publicKey, ipv4, ipv6, psk); err != nil {
		log.Printf("Warning: failed to add peer to interface: %v", err)
	} else {
		log.Printf("Added peer to WireGuard interface")
	}

	// Generate QR code
	qrPath := filepath.Join(clientsDir, fmt.Sprintf("%s.png", name))
	if err := generateQRCode(clientConfigPath, qrPath); err != nil {
		log.Printf("Warning: failed to generate QR code: %v", err)
	} else {
		log.Printf("Generated QR code: %s", qrPath)
	}

	log.Printf("Successfully created client: %s", name)
	return nil
}

// generateKeyPair generates a WireGuard private and public key pair
func generateKeyPair() (privateKey, publicKey string, err error) {
	// Generate private key
	privateKeyBytes, err := exec.Command("wg", "genkey").Output()
	if err != nil {
		return "", "", fmt.Errorf("wg genkey failed: %w", err)
	}
	privateKey = strings.TrimSpace(string(privateKeyBytes))

	// Derive public key
	pubKeyCmd := exec.Command("wg", "pubkey")
	pubKeyCmd.Stdin = strings.NewReader(privateKey)
	publicKeyBytes, err := pubKeyCmd.Output()
	if err != nil {
		return "", "", fmt.Errorf("wg pubkey failed: %w", err)
	}
	publicKey = strings.TrimSpace(string(publicKeyBytes))

	return privateKey, publicKey, nil
}

// generatePSK generates a WireGuard preshared key
func generatePSK() (string, error) {
	psk, err := exec.Command("wg", "genpsk").Output()
	if err != nil {
		return "", fmt.Errorf("wg genpsk failed: %w", err)
	}
	return strings.TrimSpace(string(psk)), nil
}

// getNextAvailableIP finds the next available IP address in the given CIDR range
func getNextAvailableIP(cidr string, existingClients []Client) (string, error) {
	// Parse CIDR to get network
	parts := strings.Split(cidr, "/")
	if len(parts) != 2 {
		return "", fmt.Errorf("invalid CIDR format: %s", cidr)
	}

	network := strings.TrimSpace(parts[0])

	// For IPv4, extract base network and assign next available host
	if !strings.Contains(network, ":") {
		// IPv4: Simple implementation - use .1, .2, etc.
		// In production, this would parse the CIDR properly
		usedHosts := make(map[string]bool)
		for _, client := range existingClients {
			if client.IPv4 != "" {
				ipParts := strings.Split(client.IPv4, ".")
				if len(ipParts) == 4 {
					usedHosts[ipParts[3]] = true
				}
			}
		}

		// Find next available host (skip 0 and 1 as they may be reserved)
		for i := 2; i < 255; i++ {
			host := fmt.Sprintf("%d", i)
			if !usedHosts[host] {
				return fmt.Sprintf("%s.%s", network, host), nil
			}
		}

		return "", fmt.Errorf("no available IPv4 addresses in range: %s", cidr)
	}

	// IPv6: Similar simplified approach
	usedHosts := make(map[string]bool)
	for _, client := range existingClients {
		if client.IPv6 != "" {
			// Extract last segment for IPv6
			lastColon := strings.LastIndex(client.IPv6, ":")
			if lastColon > 0 {
				host := client.IPv6[lastColon+1:]
				usedHosts[host] = true
			}
		}
	}

	// Find next available host
	for i := 1; i < 65536; i++ {
		host := fmt.Sprintf("%x", i)
		if !usedHosts[host] {
			return fmt.Sprintf("%s:%s", network, host), nil
		}
	}

	return "", fmt.Errorf("no available IPv6 addresses in range: %s", cidr)
}

// generateServerConfig generates the server-side configuration for a client
func generateServerConfig(name, publicKey, ipv4, ipv6, psk string, cfg *config.Config) (string, error) {
	var builder strings.Builder

	builder.WriteString(fmt.Sprintf("# Client: %s\n", name))
	builder.WriteString(fmt.Sprintf("[Peer]\n"))
	builder.WriteString(fmt.Sprintf("PublicKey = %s\n", publicKey))

	allowedIPs := ""
	if ipv4 != "" {
		allowedIPs = ipv4 + "/32"
	}
	if ipv6 != "" {
		if allowedIPs != "" {
			allowedIPs += ", "
		}
		allowedIPs += ipv6 + "/128"
	}
	builder.WriteString(fmt.Sprintf("AllowedIPs = %s\n", allowedIPs))

	if psk != "" {
		builder.WriteString(fmt.Sprintf("PresharedKey = %s\n", psk))
	}

	return builder.String(), nil
}

// generateClientConfig generates the client-side configuration
func generateClientConfig(name, privateKey, ipv4, ipv6, dns string, cfg *config.Config) (string, error) {
	// Get server's public key from the main config
	serverConfigPath := "/etc/wireguard/wg0.conf"
	serverPublicKey, serverEndpoint, err := getServerConfig(serverConfigPath)
	if err != nil {
		return "", fmt.Errorf("failed to read server config: %w", err)
	}

	var builder strings.Builder

	builder.WriteString(fmt.Sprintf("# WireGuard client configuration for %s\n", name))
	builder.WriteString("[Interface]\n")
	builder.WriteString(fmt.Sprintf("PrivateKey = %s\n", privateKey))
	builder.WriteString(fmt.Sprintf("Address = %s/32", ipv4))
	if ipv6 != "" {
		builder.WriteString(fmt.Sprintf(", %s/128", ipv6))
	}
	builder.WriteString("\n")
	builder.WriteString(fmt.Sprintf("DNS = %s\n", dns))

	builder.WriteString("\n")
	builder.WriteString("[Peer]\n")
	builder.WriteString(fmt.Sprintf("PublicKey = %s\n", serverPublicKey))
	builder.WriteString(fmt.Sprintf("Endpoint = %s:%d\n", serverEndpoint, cfg.WGPort))
	builder.WriteString("AllowedIPs = 0.0.0.0/0, ::/0\n")

	return builder.String(), nil
}

// getServerConfig reads the server's public key and endpoint from the main config
func getServerConfig(path string) (publicKey, endpoint string, err error) {
	content, err := os.ReadFile(path)
	if err != nil {
		return "", "", fmt.Errorf("failed to read server config: %w", err)
	}

	inInterfaceSection := false
	for _, line := range strings.Split(string(content), "\n") {
		line = strings.TrimSpace(line)

		if line == "[Interface]" {
			inInterfaceSection = true
			continue
		}
		if line == "[Peer]" {
			inInterfaceSection = false
			continue
		}

		if inInterfaceSection {
			if strings.HasPrefix(line, "PublicKey") {
				parts := strings.SplitN(line, "=", 2)
				if len(parts) == 2 {
					publicKey = strings.TrimSpace(parts[1])
				}
			}
		}
	}

	// Use SERVER_DOMAIN as endpoint if available, otherwise fallback
	cfg, err := config.LoadConfig()
	if err == nil && cfg.ServerDomain != "" {
		endpoint = cfg.ServerDomain
	} else {
		endpoint = "0.0.0.0"
	}

	return publicKey, endpoint, nil
}

// addPeerToInterface adds a peer to the WireGuard interface
func addPeerToInterface(publicKey, ipv4, ipv6, psk string) error {
	args := []string{"set", "wg0", "peer", publicKey}

	if ipv4 != "" {
		args = append(args, "allowed-ips", ipv4+"/32")
	}
	if ipv6 != "" {
		args = append(args, ipv6+"/128")
	}

	if psk != "" {
		args = append(args, "preshared-key", "/dev/stdin")
		cmd := exec.Command("wg", args...)
		cmd.Stdin = strings.NewReader(psk)
		output, err := cmd.CombinedOutput()
		if err != nil {
			return fmt.Errorf("wg set peer failed: %w, output: %s", err, string(output))
		}
	} else {
		cmd := exec.Command("wg", args...)
		output, err := cmd.CombinedOutput()
		if err != nil {
			return fmt.Errorf("wg set peer failed: %w, output: %s", err, string(output))
		}
	}

	return nil
}

// generateQRCode generates a QR code from the client config
func generateQRCode(configPath, qrPath string) error {
	// Read config file
	content, err := os.ReadFile(configPath)
	if err != nil {
		return fmt.Errorf("failed to read config file: %w", err)
	}

	// Generate QR code using qrencode or similar
	// For now, use a simple approach with qrencode if available
	cmd := exec.Command("qrencode", "-o", qrPath, "-t", "PNG")
	cmd.Stdin = bytes.NewReader(content)
	if err := cmd.Run(); err != nil {
		// qrencode not available, try alternative method
		return fmt.Errorf("qrencode not available: %w", err)
	}

	return nil
}