feat: add server stats

server/app/ws/stats/ssh.go

@@ -0,0 +1,209 @@
+package stats
+
+import (
+	"context"
+	"fmt"
+	"log"
+	"strconv"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/gofiber/contrib/websocket"
+	"rul.sh/vaulterm/lib"
+)
+
+func HandleSSHStats(c *websocket.Conn, client *lib.SSHClient) error {
+	if err := client.Connect(); err != nil {
+		log.Printf("error connecting to SSH: %v", err)
+		return err
+	}
+	defer client.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	msgCh := make(chan string)
+
+	go func() {
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			default:
+				wg := &sync.WaitGroup{}
+				wg.Add(4)
+				go getCPUUsage(client, wg, msgCh)
+				go getMemoryUsage(client, wg, msgCh)
+				go getDiskUsage(client, wg, msgCh)
+				go getNetworkUsage(client, wg, msgCh)
+				wg.Wait()
+			}
+		}
+	}()
+
+	go func() {
+		for msg := range msgCh {
+			if err := c.WriteMessage(websocket.TextMessage, []byte(msg)); err != nil {
+				break
+			}
+		}
+	}()
+
+	for {
+		_, _, err := c.ReadMessage()
+		if err != nil {
+			break
+		}
+	}
+
+	return nil
+}
+
+func getCPUUsage(client *lib.SSHClient, wg *sync.WaitGroup, result chan<- string) {
+	defer wg.Done()
+
+	cpuData, err := client.Exec("cat /proc/stat | grep '^cpu '")
+	if err != nil {
+		return
+	}
+	total1, idle1, err := parseCPUStats(cpuData)
+	if err != nil {
+		return
+	}
+
+	time.Sleep(time.Second)
+
+	cpuData, err = client.Exec("cat /proc/stat | grep '^cpu '")
+	if err != nil {
+		return
+	}
+	total2, idle2, err := parseCPUStats(cpuData)
+	if err != nil {
+		return
+	}
+
+	totalDiff := total2 - total1
+	idleDiff := idle2 - idle1
+	usage := (float64(totalDiff-idleDiff) / float64(totalDiff)) * 100
+
+	result <- fmt.Sprintf("\x01%.2f", usage)
+}
+
+func parseCPUStats(data string) (int64, int64, error) {
+	fields := strings.Fields(data)
+	if len(fields) < 8 {
+		return 0, 0, fmt.Errorf("unexpected format in /proc/stat")
+	}
+
+	user, _ := strconv.ParseInt(fields[1], 10, 64)
+	nice, _ := strconv.ParseInt(fields[2], 10, 64)
+	system, _ := strconv.ParseInt(fields[3], 10, 64)
+	idle, _ := strconv.ParseInt(fields[4], 10, 64)
+	iowait, _ := strconv.ParseInt(fields[5], 10, 64)
+	irq, _ := strconv.ParseInt(fields[6], 10, 64)
+	softirq, _ := strconv.ParseInt(fields[7], 10, 64)
+
+	total := user + nice + system + idle + iowait + irq + softirq
+	idle = idle + iowait
+
+	return total, idle, nil
+}
+
+func getMemoryUsage(client *lib.SSHClient, wg *sync.WaitGroup, result chan<- string) {
+	defer wg.Done()
+	data, err := client.Exec("cat /proc/meminfo")
+	if err != nil {
+		return
+	}
+
+	var total, available int
+	lines := strings.Split(data, "\n")
+
+	for _, line := range lines {
+		line = strings.TrimSpace(strings.ToLower(line))
+		fields := strings.Fields(line)
+		if len(fields) < 2 {
+			continue
+		}
+		value, _ := strconv.Atoi(fields[1])
+		if strings.HasPrefix(line, "memtotal") {
+			total = value / 1024
+		} else if strings.HasPrefix(line, "memavailable") {
+			available = value / 1024
+		}
+	}
+
+	result <- fmt.Sprintf("\x02%d,%d", total, available)
+}
+
+func getDiskUsage(client *lib.SSHClient, wg *sync.WaitGroup, result chan<- string) {
+	defer wg.Done()
+	data, err := client.Exec("df -h /")
+	if err != nil {
+		return
+	}
+	lines := strings.Split(data, "\n")
+	if len(lines) < 2 {
+		return
+	}
+
+	fields := strings.Fields(lines[1])
+	result <- fmt.Sprintf("\x03%s,%s,%s", fields[1], fields[2], fields[4])
+}
+
+func getNetworkUsage(client *lib.SSHClient, wg *sync.WaitGroup, result chan<- string) {
+	defer wg.Done()
+
+	cmd := `iface=$(ip route | awk '/^default/ {print $5}'); if [ -n "$iface" ]; then ip -s link show "$iface"; fi`
+	data, err := client.Exec(cmd)
+	if err != nil || strings.TrimSpace(data) == "" {
+		return
+	}
+
+	// Parse RX/TX values from the network data
+	rx, tx := parseNetwork(data)
+	result <- fmt.Sprintf("\x04%d,%d", rx/1024/1024, tx/1024/1024)
+}
+
+func parseNetwork(data string) (int, int) {
+	lines := strings.Split(data, "\n")
+	var rxBytes, txBytes int
+	rxMode, txMode := false, false
+
+	for _, line := range lines {
+		line := strings.TrimSpace(line)
+
+		// Check for RX and TX headers
+		if strings.HasPrefix(line, "RX:") {
+			rxMode = true
+			txMode = false
+			continue
+		}
+		if strings.HasPrefix(line, "TX:") {
+			txMode = true
+			rxMode = false
+			continue
+		}
+
+		// Parse RX bytes if in RX mode
+		if rxMode {
+			fields := strings.Fields(line)
+			if len(fields) > 0 {
+				rxBytes, _ = strconv.Atoi(fields[0])
+			}
+			rxMode = false // Reset RX mode after capturing data
+		}
+
+		// Parse TX bytes if in TX mode
+		if txMode {
+			fields := strings.Fields(line)
+			if len(fields) > 0 {
+				txBytes, _ = strconv.Atoi(fields[0])
+			}
+			txMode = false // Reset TX mode after capturing data
+		}
+	}
+
+	return txBytes, rxBytes
+}

server/app/ws/stats/stats.go

@@ -0,0 +1,42 @@
+package stats
+
+import (
+	"github.com/gofiber/contrib/websocket"
+	"rul.sh/vaulterm/app/hosts"
+	"rul.sh/vaulterm/lib"
+	"rul.sh/vaulterm/models"
+	"rul.sh/vaulterm/utils"
+)
+
+func HandleStats(c *websocket.Conn) {
+	hostId := c.Query("hostId")
+
+	user := utils.GetUserWs(c)
+	hostRepo := hosts.NewRepository(&hosts.Hosts{User: user})
+	data, _ := hostRepo.Get(hostId)
+
+	if data == nil || !data.HasAccess(&user.User) {
+		c.WriteMessage(websocket.TextMessage, []byte("Host not found"))
+		return
+	}
+
+	switch data.Host.Type {
+	case "ssh":
+		sshHandler(c, data)
+	default:
+		c.WriteMessage(websocket.TextMessage, []byte("Invalid host type"))
+	}
+}
+
+func sshHandler(c *websocket.Conn, data *models.HostDecrypted) {
+	cfg := lib.NewSSHClient(&lib.SSHClientConfig{
+		HostName: data.Host.Host,
+		Port:     data.Port,
+		Key:      data.Key,
+		AltKey:   data.AltKey,
+	})
+
+	if err := HandleSSHStats(c, cfg); err != nil {
+		c.WriteMessage(websocket.TextMessage, []byte(err.Error()))
+	}
+}