wg-admin/internal/wireguard/status.go

package wireguard

import (
	"bytes"
	"fmt"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"time"
)

const (
	// StatusConnected indicates a peer has an active connection
	StatusConnected = "Connected"
	// StatusDisconnected indicates a peer is not connected
	StatusDisconnected = "Disconnected"

	// QualityExcellent indicates handshake was very recent (< 30s)
	QualityExcellent = "Excellent"
	// QualityGood indicates handshake was recent (< 2m)
	QualityGood = "Good"
	// QualityFair indicates handshake was acceptable (< 5m)
	QualityFair = "Fair"
	// QualityPoor indicates handshake was old (> 5m)
	QualityPoor = "Poor"
)

// PeerStatus represents the status of a WireGuard peer
type PeerStatus struct {
	PublicKey       string    `json:"public_key"`
	Endpoint        string    `json:"endpoint"`
	AllowedIPs      string    `json:"allowed_ips"`
	LatestHandshake time.Time `json:"latest_handshake"`
	TransferRx      string    `json:"transfer_rx"`
	TransferTx      string    `json:"transfer_tx"`
	Status          string    `json:"status"`            // "Connected" or "Disconnected"
	Quality         string    `json:"quality,omitempty"` // "Excellent", "Good", "Fair", "Poor" (if connected)
}

// GetClientStatus checks if a specific client is connected
// Returns "Connected" if the peer appears in the active peers list, "Disconnected" otherwise
func GetClientStatus(publicKey string) (string, error) {
	peers, err := GetAllPeers()
	if err != nil {
		return StatusDisconnected, fmt.Errorf("failed to get peer status: %w", err)
	}

	for _, peer := range peers {
		if peer.PublicKey == publicKey {
			return peer.Status, nil
		}
	}

	return StatusDisconnected, nil
}

// GetAllPeers retrieves all peers with their current status from WireGuard
func GetAllPeers() ([]PeerStatus, error) {
	output, err := exec.Command("wg", "show", "wg0").Output()
	if err != nil {
		return nil, fmt.Errorf("failed to execute wg show: %w", err)
	}

	return parsePeersOutput(string(output)), nil
}

// parsePeersOutput parses the output of 'wg show wg0' command
func parsePeersOutput(output string) []PeerStatus {
	var peers []PeerStatus
	var currentPeer *PeerStatus
	var handshake string
	var transfer string

	lines := strings.Split(output, "\n")
	peerLineRegex := regexp.MustCompile(`^peer:\s*(.+)$`)
	handshakeRegex := regexp.MustCompile(`^latest handshake:\s*(.+)\s+ago$`)
	transferRegex := regexp.MustCompile(`^transfer:\s*(.+),\s+(.+)$`)

	for _, line := range lines {
		line = strings.TrimSpace(line)

		// Check for new peer
		if match := peerLineRegex.FindStringSubmatch(line); match != nil {
			// Save previous peer if exists
			if currentPeer != nil {
				peers = append(peers, finalizePeerStatus(currentPeer, handshake, transfer))
			}

			// Start new peer
			currentPeer = &PeerStatus{
				PublicKey: match[1],
			}
			handshake = ""
			transfer = ""
			continue
		}

		if currentPeer == nil {
			continue
		}

		// Parse endpoint
		if strings.HasPrefix(line, "endpoint:") {
			currentPeer.Endpoint = strings.TrimSpace(strings.TrimPrefix(line, "endpoint:"))
		}

		// Parse allowed ips
		if strings.HasPrefix(line, "allowed ips:") {
			currentPeer.AllowedIPs = strings.TrimSpace(strings.TrimPrefix(line, "allowed ips:"))
		}

		// Parse latest handshake
		if match := handshakeRegex.FindStringSubmatch(line); match != nil {
			handshake = match[1]
		}

		// Parse transfer
		if match := transferRegex.FindStringSubmatch(line); match != nil {
			transfer = fmt.Sprintf("%s, %s", match[1], match[2])
		}
	}

	// Don't forget the last peer
	if currentPeer != nil {
		peers = append(peers, finalizePeerStatus(currentPeer, handshake, transfer))
	}

	return peers
}

// CalculateQuality returns the connection quality based on handshake time
func CalculateQuality(timeSinceHandshake time.Duration) string {
	if timeSinceHandshake < 30*time.Second {
		return QualityExcellent
	}
	if timeSinceHandshake < 2*time.Minute {
		return QualityGood
	}
	if timeSinceHandshake < 5*time.Minute {
		return QualityFair
	}
	return QualityPoor
}

// finalizePeerStatus determines the peer's status based on handshake time
func finalizePeerStatus(peer *PeerStatus, handshake string, transfer string) PeerStatus {
	peer.TransferRx = ""
	peer.TransferTx = ""

	// Parse transfer
	if transfer != "" {
		parts := strings.Split(transfer, ", ")
		if len(parts) == 2 {
			// Extract received and sent values
			rxParts := strings.Fields(parts[0])
			if len(rxParts) >= 2 {
				peer.TransferRx = strings.Join(rxParts[:2], " ")
			}
			txParts := strings.Fields(parts[1])
			if len(txParts) >= 2 {
				peer.TransferTx = strings.Join(txParts[:2], " ")
			}
		}
	}

	// Determine status and quality based on handshake
	if handshake != "" {
		peer.LatestHandshake = parseHandshake(handshake)
		timeSinceHandshake := time.Since(peer.LatestHandshake)

		// Peer is considered connected if handshake is recent (within 5 minutes)
		// This allows for ~12 missed keepalive intervals (at 25 seconds each)
		if timeSinceHandshake < 5*time.Minute {
			peer.Status = StatusConnected
			peer.Quality = calculateQuality(timeSinceHandshake)
		} else {
			peer.Status = StatusDisconnected
		}
	} else {
		peer.Status = StatusDisconnected
	}

	return *peer
}

// parseHandshake converts handshake string to time.Time
func parseHandshake(handshake string) time.Time {
	now := time.Now()
	var totalDuration time.Duration
	parts := strings.Fields(handshake)

	for i, part := range parts {
		// Clean up commas
		cleanPart := strings.TrimSuffix(part, ",")

		// Check if this part is a time unit
		if strings.HasSuffix(cleanPart, "second") || strings.HasSuffix(cleanPart, "seconds") {
			// The number is the previous part
			if i > 0 {
				if val, err := strconv.Atoi(parts[i-1]); err == nil {
					totalDuration += time.Duration(val) * time.Second
				}
			}
		}
		if strings.HasSuffix(cleanPart, "minute") || strings.HasSuffix(cleanPart, "minutes") {
			if i > 0 {
				if val, err := strconv.Atoi(parts[i-1]); err == nil {
					totalDuration += time.Duration(val) * time.Minute
				}
			}
		}
		if strings.HasSuffix(cleanPart, "hour") || strings.HasSuffix(cleanPart, "hours") {
			if i > 0 {
				if val, err := strconv.Atoi(parts[i-1]); err == nil {
					totalDuration += time.Duration(val) * time.Hour
				}
			}
		}
		if strings.HasSuffix(cleanPart, "day") || strings.HasSuffix(cleanPart, "days") {
			if i > 0 {
				if val, err := strconv.Atoi(parts[i-1]); err == nil {
					totalDuration += time.Duration(val) * 24 * time.Hour
				}
			}
		}
	}

	if totalDuration == 0 {
		return now.Add(-time.Hour) // Default to 1 hour ago if parsing fails
	}

	return now.Add(-totalDuration)
}

// CheckInterface verifies if the WireGuard interface exists and is accessible
func CheckInterface(interfaceName string) error {
	cmd := exec.Command("wg", "show", interfaceName)
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		return fmt.Errorf("wireguard interface '%s' not accessible: %w", interfaceName, err)
	}
	return nil
}