feat: add server stats

frontend/components/containers/interactive-session.tsx

@@ -3,6 +3,8 @@ import Terminal from "./terminal";
 import VNCViewer from "./vncviewer";
 import { useAuthStore } from "@/stores/auth";
 import { AppServer, useServer } from "@/stores/app";
+import { useWebsocketUrl } from "@/hooks/useWebsocket";
+import ServerStatsBar from "./server-stats-bar";
 
 type SSHSessionProps = {
   type: "ssh";
@@ -30,20 +32,25 @@ export type InteractiveSessionProps = {
 
 const InteractiveSession = ({ type, params }: InteractiveSessionProps) => {
   const { token } = useAuthStore();
-  const server = useServer();
-  const query = new URLSearchParams({ ...params, sid: token || "" });
-  const url = `${getBaseUrl(server)}/ws/term?${query}`;
+  const ws = useWebsocketUrl({ ...params, sid: token || "" });
+  const termUrl = ws("term");
+  const statsUrl = ws("stats");
 
   switch (type) {
     case "ssh":
-      return <Terminal url={url} />;
+      return (
+        <>
+          <Terminal url={termUrl} />
+          <ServerStatsBar url={statsUrl} />
+        </>
+      );
 
     case "pve":
     case "incus":
       return params.client === "vnc" ? (
-        <VNCViewer url={url} />
+        <VNCViewer url={termUrl} />
       ) : (
-        <Terminal url={url} />
+        <Terminal url={termUrl} />
       );
 
     default:
@@ -51,8 +58,4 @@ const InteractiveSession = ({ type, params }: InteractiveSessionProps) => {
   }
 };
 
-function getBaseUrl(server?: AppServer | null) {
-  return server?.url.replace("http://", "ws://") || "";
-}
-
 export default InteractiveSession;

frontend/components/containers/server-stats-bar.tsx

@@ -0,0 +1,84 @@
+import { View, Text, XStack, Separator } from "tamagui";
+import React, { useState } from "react";
+import { useWebSocket } from "@/hooks/useWebsocket";
+import Icons from "../ui/icons";
+
+type Props = {
+  url: string;
+};
+
+const ServerStatsBar = ({ url }: Props) => {
+  const [cpu, setCPU] = useState(0);
+  const [memory, setMemory] = useState({ total: 0, used: 0, available: 0 });
+  const [disk, setDisk] = useState({ total: "0", used: "0", percent: "0%" });
+  const [network, setNetwork] = useState({ tx: 0, rx: 0 });
+
+  const { isConnected } = useWebSocket(url, {
+    onMessage: (msg) => {
+      const type = msg.substring(0, 1);
+      const value = msg.substring(1);
+      let values: string[];
+
+      switch (type) {
+        case "\x01":
+          setCPU(parseFloat(value));
+          break;
+
+        case "\x02":
+          values = value.split(",");
+          const total = parseInt(values[0]) || 0;
+          const available = parseInt(values[1]) || 0;
+          const used = total - available;
+          setMemory({ total, used, available });
+          break;
+
+        case "\x03":
+          values = value.split(",");
+          setDisk({ total: values[0], used: values[1], percent: values[2] });
+          break;
+
+        case "\x04":
+          values = value.split(",");
+          setNetwork({
+            tx: parseInt(values[0]) || 0,
+            rx: parseInt(values[1]) || 0,
+          });
+          break;
+      }
+    },
+  });
+
+  if (!isConnected || !memory.total) {
+    return null;
+  }
+
+  return (
+    <XStack gap="$1" p="$2" alignItems="center">
+      <XStack gap="$1" alignItems="center" minWidth={48}>
+        <Icons name="desktop-tower" size={16} />
+        <Text fontSize="$2">{Math.round(cpu)}%</Text>
+      </XStack>
+
+      <Separator vertical h="100%" mx="$2" borderColor="$color" />
+      <Icons name="memory" size={16} />
+      <Text fontSize="$2">
+        {memory.used} MB / {memory.total} MB (
+        {Math.round((memory.used / memory.total) * 100) || 0}%)
+      </Text>
+
+      <Separator vertical h="100%" mx="$2" borderColor="$color" />
+      <Icons name="harddisk" size={16} />
+      <Text fontSize="$2">
+        {disk.used} / {disk.total} ({disk.percent})
+      </Text>
+
+      <Separator vertical h="100%" mx="$2" borderColor="$color" />
+      <Icons name="download" size={16} />
+      <Text fontSize="$2">{network.rx} MB</Text>
+      <Icons name="upload" size={16} />
+      <Text fontSize="$2">{network.tx} MB</Text>
+    </XStack>
+  );
+};
+
+export default ServerStatsBar;

frontend/hooks/useWebsocket.ts

@@ -0,0 +1,59 @@
+import { useServer } from "@/stores/app";
+
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
+
+type UseWebsocketOptions = {
+  onMessage?: (message: string) => void;
+};
+
+export const useWebSocket = (url: string, opt?: UseWebsocketOptions) => {
+  const [isConnected, setIsConnected] = useState(false);
+  const websocketRef = useRef<WebSocket | null>(null);
+
+  useEffect(() => {
+    // Create WebSocket connection
+    const ws = new WebSocket(url);
+    websocketRef.current = ws;
+
+    // Connection opened
+    ws.onopen = () => {
+      setIsConnected(true);
+      console.log("WebSocket connected");
+    };
+
+    // Listen for messages
+    ws.onmessage = (event) => {
+      opt?.onMessage?.(event.data);
+    };
+
+    // Connection closed
+    ws.onclose = () => {
+      setIsConnected(false);
+      console.log("WebSocket disconnected");
+    };
+
+    // Cleanup on unmount
+    return () => {
+      ws.close();
+    };
+  }, [url]);
+
+  // Send message function
+  const send = (msg: string) => {
+    if (isConnected && websocketRef.current) {
+      websocketRef.current.send(msg);
+    }
+  };
+
+  return { isConnected, send };
+};
+
+export const useWebsocketUrl = (initParams: any = {}) => {
+  const server = useServer();
+  const baseUrl = server?.url.replace("http://", "ws://") || "";
+
+  return (url: string, params: any = {}) => {
+    const query = new URLSearchParams({ ...initParams, ...params });
+    return `${baseUrl}/ws/${url}?${query}`;
+  };
+};

server/app/hosts/utils.go

@@ -2,6 +2,7 @@ package hosts
 
 import (
 	"fmt"
+	"log"
 
 	"github.com/gofiber/fiber/v2"
 	"rul.sh/vaulterm/app/keychains"
@@ -40,12 +41,14 @@ func tryConnect(c *fiber.Ctx, host *models.Host) (string, error) {
 			AltKey:   altKey,
 		})
 
-		con, err := c.Connect()
-		if err != nil {
+		if err := c.Connect(); err != nil {
 			return "", err
 		}
+		defer c.Close()
 
-		os, err := c.GetOS(c, con)
+		log.Println("Test", c.Conn)
+
+		os, err := c.GetOS(c)
 		if err != nil {
 			return "", err
 		}

server/app/ws/router.go

@@ -3,6 +3,8 @@ package ws
 import (
 	"github.com/gofiber/contrib/websocket"
 	"github.com/gofiber/fiber/v2"
+	"rul.sh/vaulterm/app/ws/stats"
+	"rul.sh/vaulterm/app/ws/term"
 )
 
 func Router(app fiber.Router) {
@@ -15,5 +17,6 @@ func Router(app fiber.Router) {
 		return fiber.ErrUpgradeRequired
 	})
 
-	router.Get("/term", websocket.New(HandleTerm))
+	router.Get("/term", websocket.New(term.HandleTerm))
+	router.Get("/stats", websocket.New(stats.HandleStats))
 }

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()))
+	}
+}

server/app/ws/term/incus.go

@@ -1,4 +1,4 @@
-package ws
+package term
 
 import (
 	"crypto/tls"

server/app/ws/term/pve.go

@@ -1,4 +1,4 @@
-package ws
+package term
 
 import (
 	"crypto/tls"

server/app/ws/term/ssh.go

@@ -1,4 +1,4 @@
-package ws
+package term
 
 import (
 	"io"
@@ -11,14 +11,13 @@ import (
 )
 
 func NewSSHWebsocketSession(c *websocket.Conn, client *lib.SSHClient) error {
-	con, err := client.Connect()
-	if err != nil {
+	if err := client.Connect(); err != nil {
 		log.Printf("error connecting to SSH: %v", err)
 		return err
 	}
-	defer con.Close()
+	defer client.Close()
 
-	shell, err := client.StartPtyShell(con)
+	shell, err := client.StartPtyShell()
 	if err != nil {
 		log.Printf("error starting SSH shell: %v", err)
 		return err

server/app/ws/term/term.go

@@ -1,4 +1,4 @@
-package ws
+package term
 
 import (
 	"log"

server/lib/ssh.go

@@ -14,6 +14,8 @@ type SSHClient struct {
 	Port                 int
 	PrivateKey           string
 	PrivateKeyPassphrase string
+
+	Conn *ssh.Client
 }
 
 type SSHClientConfig struct {
@@ -39,7 +41,7 @@ func NewSSHClient(cfg *SSHClientConfig) *SSHClient {
 	}
 }
 
-func (s *SSHClient) Connect() (*ssh.Client, error) {
+func (s *SSHClient) Connect() error {
 	// Set up SSH client configuration
 	port := s.Port
 	if port == 0 {
@@ -60,7 +62,7 @@ func (s *SSHClient) Connect() (*ssh.Client, error) {
 		}
 
 		if err != nil {
-			return nil, fmt.Errorf("unable to parse private key: %v", err)
+			return fmt.Errorf("unable to parse private key: %v", err)
 		}
 		auth = append(auth, ssh.PublicKeys(signer))
 	}
@@ -75,10 +77,18 @@ func (s *SSHClient) Connect() (*ssh.Client, error) {
 	hostName := fmt.Sprintf("%s:%d", s.HostName, port)
 	sshConn, err := ssh.Dial("tcp", hostName, sshConfig)
 	if err != nil {
-		return nil, err
+		return err
 	}
 
-	return sshConn, nil
+	s.Conn = sshConn
+	return nil
+}
+
+func (s *SSHClient) Close() error {
+	if s.Conn != nil {
+		return s.Conn.Close()
+	}
+	return nil
 }
 
 type PtyShellRes struct {
@@ -88,9 +98,13 @@ type PtyShellRes struct {
 	Session *ssh.Session
 }
 
-func (s *SSHClient) StartPtyShell(sshConn *ssh.Client) (res *PtyShellRes, err error) {
+func (s *SSHClient) StartPtyShell() (res *PtyShellRes, err error) {
+	if s.Conn == nil {
+		return nil, fmt.Errorf("SSH client is not connected")
+	}
+
 	// Start an SSH shell session
-	session, err := sshConn.NewSession()
+	session, err := s.Conn.NewSession()
 	if err != nil {
 		return nil, err
 	}
@@ -127,9 +141,13 @@ func (s *SSHClient) StartPtyShell(sshConn *ssh.Client) (res *PtyShellRes, err er
 	}, nil
 }
 
-func (s *SSHClient) Exec(sshConn *ssh.Client, command string) (string, error) {
+func (s *SSHClient) Exec(command string) (string, error) {
+	if s.Conn == nil {
+		return "", fmt.Errorf("SSH client is not connected")
+	}
+
 	// Start an SSH shell session
-	session, err := sshConn.NewSession()
+	session, err := s.Conn.NewSession()
 	if err != nil {
 		return "", err
 	}
@@ -144,8 +162,8 @@ func (s *SSHClient) Exec(sshConn *ssh.Client, command string) (string, error) {
 	return string(output), nil
 }
 
-func (s *SSHClient) GetOS(client *SSHClient, con *ssh.Client) (string, error) {
-	out, err := client.Exec(con, "cat /etc/os-release || uname -a || systeminfo")
+func (s *SSHClient) GetOS(client *SSHClient) (string, error) {
+	out, err := client.Exec("cat /etc/os-release || uname -a || systeminfo")
 	if err != nil {
 		return "", err
 	}