From 7a00992ff90efa93ae453650d216d83518c4b0f0 Mon Sep 17 00:00:00 2001 From: Khairul Hidayat Date: Wed, 13 Nov 2024 22:45:03 +0000 Subject: [PATCH] feat: add server stats --- .../containers/interactive-session.tsx | 23 +- .../containers/server-stats-bar.tsx | 84 +++++++ frontend/hooks/useWebsocket.ts | 59 +++++ server/app/hosts/utils.go | 9 +- server/app/ws/router.go | 5 +- server/app/ws/stats/ssh.go | 209 ++++++++++++++++++ server/app/ws/stats/stats.go | 42 ++++ .../app/ws/{term_incus.go => term/incus.go} | 2 +- server/app/ws/{term_pve.go => term/pve.go} | 2 +- server/app/ws/{term_ssh.go => term/ssh.go} | 9 +- server/app/ws/{ => term}/term.go | 2 +- server/lib/ssh.go | 38 +++- 12 files changed, 452 insertions(+), 32 deletions(-) create mode 100644 frontend/components/containers/server-stats-bar.tsx create mode 100644 frontend/hooks/useWebsocket.ts create mode 100644 server/app/ws/stats/ssh.go create mode 100644 server/app/ws/stats/stats.go rename server/app/ws/{term_incus.go => term/incus.go} (99%) rename server/app/ws/{term_pve.go => term/pve.go} (99%) rename server/app/ws/{term_ssh.go => term/ssh.go} (94%) rename server/app/ws/{ => term}/term.go (99%) diff --git a/frontend/components/containers/interactive-session.tsx b/frontend/components/containers/interactive-session.tsx index 79d5df1..02d951d 100644 --- a/frontend/components/containers/interactive-session.tsx +++ b/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 ; + return ( + <> + + + + ); case "pve": case "incus": return params.client === "vnc" ? ( - + ) : ( - + ); default: @@ -51,8 +58,4 @@ const InteractiveSession = ({ type, params }: InteractiveSessionProps) => { } }; -function getBaseUrl(server?: AppServer | null) { - return server?.url.replace("http://", "ws://") || ""; -} - export default InteractiveSession; diff --git a/frontend/components/containers/server-stats-bar.tsx b/frontend/components/containers/server-stats-bar.tsx new file mode 100644 index 0000000..923503d --- /dev/null +++ b/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 ( + + + + {Math.round(cpu)}% + + + + + + {memory.used} MB / {memory.total} MB ( + {Math.round((memory.used / memory.total) * 100) || 0}%) + + + + + + {disk.used} / {disk.total} ({disk.percent}) + + + + + {network.rx} MB + + {network.tx} MB + + ); +}; + +export default ServerStatsBar; diff --git a/frontend/hooks/useWebsocket.ts b/frontend/hooks/useWebsocket.ts new file mode 100644 index 0000000..877ee69 --- /dev/null +++ b/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(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}`; + }; +}; diff --git a/server/app/hosts/utils.go b/server/app/hosts/utils.go index d953d7a..2a20139 100644 --- a/server/app/hosts/utils.go +++ b/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 } diff --git a/server/app/ws/router.go b/server/app/ws/router.go index d0753fa..a293eb8 100644 --- a/server/app/ws/router.go +++ b/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)) } diff --git a/server/app/ws/stats/ssh.go b/server/app/ws/stats/ssh.go new file mode 100644 index 0000000..6800adc --- /dev/null +++ b/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 +} diff --git a/server/app/ws/stats/stats.go b/server/app/ws/stats/stats.go new file mode 100644 index 0000000..4ff697c --- /dev/null +++ b/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())) + } +} diff --git a/server/app/ws/term_incus.go b/server/app/ws/term/incus.go similarity index 99% rename from server/app/ws/term_incus.go rename to server/app/ws/term/incus.go index ad7cd90..3cfdbcc 100644 --- a/server/app/ws/term_incus.go +++ b/server/app/ws/term/incus.go @@ -1,4 +1,4 @@ -package ws +package term import ( "crypto/tls" diff --git a/server/app/ws/term_pve.go b/server/app/ws/term/pve.go similarity index 99% rename from server/app/ws/term_pve.go rename to server/app/ws/term/pve.go index 80de70a..d6060df 100644 --- a/server/app/ws/term_pve.go +++ b/server/app/ws/term/pve.go @@ -1,4 +1,4 @@ -package ws +package term import ( "crypto/tls" diff --git a/server/app/ws/term_ssh.go b/server/app/ws/term/ssh.go similarity index 94% rename from server/app/ws/term_ssh.go rename to server/app/ws/term/ssh.go index 82dfadd..7cd6f7d 100644 --- a/server/app/ws/term_ssh.go +++ b/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 diff --git a/server/app/ws/term.go b/server/app/ws/term/term.go similarity index 99% rename from server/app/ws/term.go rename to server/app/ws/term/term.go index 958f2cb..776ee52 100644 --- a/server/app/ws/term.go +++ b/server/app/ws/term/term.go @@ -1,4 +1,4 @@ -package ws +package term import ( "log" diff --git a/server/lib/ssh.go b/server/lib/ssh.go index 1fcb39e..66c9b2c 100644 --- a/server/lib/ssh.go +++ b/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 }