feat: add server stats

This commit is contained in:
Khairul Hidayat 2024-11-13 22:45:03 +00:00
parent 2d4c81e15d
commit 7a00992ff9
12 changed files with 452 additions and 32 deletions

View File

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

View File

@ -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;

View File

@ -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}`;
};
};

View File

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

View File

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

209
server/app/ws/stats/ssh.go Normal file
View File

@ -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
}

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@ -14,6 +14,8 @@ type SSHClient struct {
Port int Port int
PrivateKey string PrivateKey string
PrivateKeyPassphrase string PrivateKeyPassphrase string
Conn *ssh.Client
} }
type SSHClientConfig struct { 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 // Set up SSH client configuration
port := s.Port port := s.Port
if port == 0 { if port == 0 {
@ -60,7 +62,7 @@ func (s *SSHClient) Connect() (*ssh.Client, error) {
} }
if err != nil { 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)) 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) hostName := fmt.Sprintf("%s:%d", s.HostName, port)
sshConn, err := ssh.Dial("tcp", hostName, sshConfig) sshConn, err := ssh.Dial("tcp", hostName, sshConfig)
if err != nil { 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 { type PtyShellRes struct {
@ -88,9 +98,13 @@ type PtyShellRes struct {
Session *ssh.Session 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 // Start an SSH shell session
session, err := sshConn.NewSession() session, err := s.Conn.NewSession()
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -127,9 +141,13 @@ func (s *SSHClient) StartPtyShell(sshConn *ssh.Client) (res *PtyShellRes, err er
}, nil }, 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 // Start an SSH shell session
session, err := sshConn.NewSession() session, err := s.Conn.NewSession()
if err != nil { if err != nil {
return "", err return "", err
} }
@ -144,8 +162,8 @@ func (s *SSHClient) Exec(sshConn *ssh.Client, command string) (string, error) {
return string(output), nil return string(output), nil
} }
func (s *SSHClient) GetOS(client *SSHClient, con *ssh.Client) (string, error) { func (s *SSHClient) GetOS(client *SSHClient) (string, error) {
out, err := client.Exec(con, "cat /etc/os-release || uname -a || systeminfo") out, err := client.Exec("cat /etc/os-release || uname -a || systeminfo")
if err != nil { if err != nil {
return "", err return "", err
} }