This commit is contained in:
Akizon77
2025-04-11 17:26:34 +08:00
commit c2a9148d4c
14 changed files with 705 additions and 0 deletions

2
.gitignore vendored Normal file
View File

@@ -0,0 +1,2 @@
agent.json
.vscode/

68
config/local.go Normal file
View File

@@ -0,0 +1,68 @@
package config
import (
"encoding/json"
"flag"
"os"
)
type LocalConfig struct {
Endpoint string `json:"endpoint"`
Token string `json:"token"`
Terminal bool `json:"terminal"`
MaxRetries int `json:"maxRetries"`
ReconnectInterval int `json:"reconnectInterval"`
IgnoreUnsafeCert bool `json:"ignoreUnsafeCert"`
}
func LoadConfig() (LocalConfig, error) {
var (
endpoint string
token string
terminal bool
path string
maxRetries int
reconnectInterval int
ignoreUnsafeCert bool
)
flag.StringVar(&endpoint, "e", "", "The endpoint URL")
flag.StringVar(&token, "token", "", "The authentication token")
flag.BoolVar(&terminal, "terminal", false, "Enable or disable terminal (default: false)")
flag.StringVar(&path, "c", "agent.json", "Path to the configuration file")
flag.IntVar(&maxRetries, "maxRetries", 10, "Maximum number of retries for WebSocket connection")
flag.IntVar(&reconnectInterval, "reconnectInterval", 5, "Reconnect interval in seconds")
flag.BoolVar(&ignoreUnsafeCert,"ignoreUnsafeCert", false, "Ignore unsafe certificate errors")
flag.Parse()
// Ensure -c cannot coexist with other flags
if path != "agent.json" && (endpoint != "" || token != "" || !terminal) {
return LocalConfig{}, flag.ErrHelp
}
// 必填项 Endpoint、Token 没有读取配置文件
if endpoint == "" || token == "" {
file, err := os.Open(path)
if err != nil {
return LocalConfig{}, err
}
defer file.Close()
var localConfig LocalConfig
if err := json.NewDecoder(file).Decode(&localConfig); err != nil {
return LocalConfig{}, err
}
return localConfig, nil
}
return LocalConfig{
Endpoint: endpoint,
Token: token,
Terminal: terminal,
MaxRetries: maxRetries,
ReconnectInterval: reconnectInterval,
IgnoreUnsafeCert: ignoreUnsafeCert,
}, nil
}

70
config/remote.go Normal file
View File

@@ -0,0 +1,70 @@
package config
import (
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
"time"
)
type RemoteConfig struct {
Cpu bool `json:"cpu"`
Gpu bool `json:"gpu"`
Ram bool `json:"ram"`
Swap bool `json:"swap"`
Load bool `json:"load"`
Uptime bool `json:"uptime"`
Temperature bool `json:"temperature"`
Os bool `json:"os"`
Disk bool `json:"disk"`
Network bool `json:"network"`
Process bool `json:"process"`
Interval int `json:"interval"`
Connections bool `json:"connections"`
}
// 使用HTTP GET请求远程配置
//
// GET /api/getRemoteConfig
//
// Request the remote configuration
func LoadRemoteConfig(endpoint string, token string) (RemoteConfig, error) {
const maxRetry = 3
endpoint = strings.TrimSuffix(endpoint, "/") + "/api/getRemoteConfig" + "?token=" + token
var resp *http.Response
var err error
for attempt := 1; attempt <= maxRetry; attempt++ {
resp, err = http.Get(endpoint,)
if err == nil && resp.StatusCode == http.StatusOK {
break
}
if resp != nil {
resp.Body.Close()
}
if attempt == maxRetry {
if err != nil {
return RemoteConfig{}, fmt.Errorf("failed to fetchafter %d attempts: %v", maxRetry, err)
}
return RemoteConfig{}, fmt.Errorf("failed to fetch after %d attempts: %s", maxRetry, resp.Status)
}
time.Sleep(time.Second * time.Duration(attempt)) // Exponential backoff
}
defer resp.Body.Close()
response, err := io.ReadAll(resp.Body)
if err != nil {
return RemoteConfig{}, err
}
var remoteConfig RemoteConfig
if err := json.Unmarshal(response, &remoteConfig); err != nil {
return RemoteConfig{}, err
}
return remoteConfig, nil
}

17
go.mod Normal file
View File

@@ -0,0 +1,17 @@
module komari
go 1.23.2
require (
github.com/gorilla/websocket v1.5.3
github.com/shirou/gopsutil v3.21.11+incompatible
golang.org/x/sys v0.32.0
)
require (
github.com/go-ole/go-ole v1.2.6 // indirect
github.com/stretchr/testify v1.10.0 // indirect
github.com/tklauser/go-sysconf v0.3.15 // indirect
github.com/tklauser/numcpus v0.10.0 // indirect
github.com/yusufpapurcu/wmi v1.2.4 // indirect
)

23
go.sum Normal file
View File

@@ -0,0 +1,23 @@
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY=
github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0=
github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/shirou/gopsutil v3.21.11+incompatible h1:+1+c1VGhc88SSonWP6foOcLhvnKlUeu/erjjvaPEYiI=
github.com/shirou/gopsutil v3.21.11+incompatible/go.mod h1:5b4v6he4MtMOwMlS0TUMTu2PcXUg8+E1lC7eC3UO/RA=
github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
github.com/tklauser/go-sysconf v0.3.15 h1:VE89k0criAymJ/Os65CSn1IXaol+1wrsFHEB8Ol49K4=
github.com/tklauser/go-sysconf v0.3.15/go.mod h1:Dmjwr6tYFIseJw7a3dRLJfsHAMXZ3nEnL/aZY+0IuI4=
github.com/tklauser/numcpus v0.10.0 h1:18njr6LDBk1zuna922MgdjQuJFjrdppsZG60sHGfjso=
github.com/tklauser/numcpus v0.10.0/go.mod h1:BiTKazU708GQTYF4mB+cmlpT2Is1gLk7XVuEeem8LsQ=
github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0=
github.com/yusufpapurcu/wmi v1.2.4/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0=
golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
golang.org/x/sys v0.32.0 h1:s77OFDvIQeibCmezSnk/q6iAfkdiQaJi4VzroCFrN20=
golang.org/x/sys v0.32.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

265
main.go Normal file
View File

@@ -0,0 +1,265 @@
package main
import (
"crypto/tls"
"encoding/json"
"fmt"
"komari/config"
"komari/monitoring"
"log"
"net/http"
"strings"
"time"
"github.com/gorilla/websocket"
)
func main() {
localConfig, err := config.LoadConfig()
if err != nil {
log.Fatalln("Failed to load local config:", err)
}
if localConfig.IgnoreUnsafeCert {
http.DefaultTransport.(*http.Transport).TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
}
remoteConfig, err := config.LoadRemoteConfig(localConfig.Endpoint, localConfig.Token)
if err != nil {
log.Fatalln("Failed to load remote config:", err)
}
//log.Println("Remote Config:", remoteConfig)
err = uploadBasicInfo(localConfig.Endpoint, localConfig.Token)
if err != nil {
log.Fatalln("Failed to upload basic info:", err)
}
websocketEndpoint := strings.TrimSuffix(localConfig.Endpoint, "/") + "/ws/report"
websocketEndpoint = "ws" + strings.TrimPrefix(websocketEndpoint, "http")
var conn *websocket.Conn
defer func() {
if conn != nil {
conn.Close()
}
}()
ticker := time.NewTicker(time.Duration(remoteConfig.Interval * int(time.Second)))
defer ticker.Stop()
for range ticker.C {
// If no connection, attempt to connect
if conn == nil {
log.Println("Attempting to connect to WebSocket...")
retry := 0
for retry < localConfig.MaxRetries {
conn, err = connectWebSocket(websocketEndpoint, localConfig.Endpoint, localConfig.Token)
if err == nil {
log.Println("WebSocket connected")
go handleWebSocketMessages(localConfig, remoteConfig, conn, make(chan struct{}))
break
}
retry++
time.Sleep(time.Duration(localConfig.ReconnectInterval) * time.Second)
}
if retry >= localConfig.MaxRetries {
log.Println("Max retries reached, falling back to POST")
// Send report via POST and continue
data := report(localConfig, remoteConfig)
if err := reportWithPOST(localConfig.Endpoint, data); err != nil {
log.Println("Failed to send POST report:", err)
}
continue
}
}
// Send report via WebSocket
data := report(localConfig, remoteConfig)
err = conn.WriteMessage(websocket.TextMessage, data)
if err != nil {
log.Println("Failed to send WebSocket message:", err)
conn.Close()
conn = nil // Mark connection as dead
continue
}
}
}
// connectWebSocket attempts to establish a WebSocket connection and upload basic info
func connectWebSocket(websocketEndpoint, endpoint, token string) (*websocket.Conn, error) {
dialer := &websocket.Dialer{
HandshakeTimeout: 5 * time.Second,
}
conn, _, err := dialer.Dial(websocketEndpoint, nil)
if err != nil {
return nil, err
}
// Upload basic info after successful connection
if err := uploadBasicInfo(endpoint, token); err != nil {
log.Println("Failed to upload basic info:", err)
// Note: We don't return error here to allow the connection to proceed
}
return conn, nil
}
func handleWebSocketMessages(localConfig config.LocalConfig, remoteConfig config.RemoteConfig, conn *websocket.Conn, done chan<- struct{}) {
defer close(done)
for {
_, message_raw, err := conn.ReadMessage()
if err != nil {
log.Println("WebSocket read error:", err)
return
}
// TODO: Remote config update
// TODO: Handle incoming messages
log.Println("Received message:", string(message_raw))
message := make(map[string]interface{})
err = json.Unmarshal(message_raw, &message)
if err != nil {
log.Println("Bad ws message:", err)
continue
}
}
}
func reportWithPOST(endpoint string, data []byte) error {
url := strings.TrimSuffix(endpoint, "/") + "/api/report"
req, err := http.NewRequest("POST", url, strings.NewReader(string(data)))
if err != nil {
return err
}
req.Header.Set("Content-Type", "application/json")
client := &http.Client{}
resp, err := client.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return err
}
return nil
}
func uploadBasicInfo(endpoint string, token string) error {
cpu := monitoring.Cpu()
osname := monitoring.OSName()
data := map[string]interface{}{
"token": token,
"cpu": map[string]interface{}{
"name": cpu.CPUName,
"cores": cpu.CPUCores,
"arch": cpu.CPUArchitecture,
},
"os": osname,
}
endpoint = strings.TrimSuffix(endpoint, "/") + "/api/nodeBasicInfo"
payload, err := json.Marshal(data)
if err != nil {
return err
}
req, err := http.NewRequest("POST", endpoint, strings.NewReader(string(payload)))
if err != nil {
return err
}
req.Header.Set("Content-Type", "application/json")
client := &http.Client{}
resp, err := client.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return fmt.Errorf("status code: %d", resp.StatusCode)
}
return nil
}
func report(localConfig config.LocalConfig, remoteConfig config.RemoteConfig) []byte {
message := ""
data := map[string]interface{}{
"token": localConfig.Token,
}
if remoteConfig.Cpu {
cpu := monitoring.Cpu()
data["cpu"] = map[string]interface{}{
"usage": cpu.CPUUsage,
}
}
if remoteConfig.Ram {
ram := monitoring.Ram()
data["ram"] = map[string]interface{}{
"total": ram.Total,
"used": ram.Used,
}
}
if remoteConfig.Swap {
swap := monitoring.Swap()
data["swap"] = map[string]interface{}{
"total": swap.Total,
"used": swap.Used,
}
}
if remoteConfig.Load {
load := monitoring.Load()
data["load"] = map[string]interface{}{
"load1": load.Load1,
"load5": load.Load5,
"load15": load.Load15,
}
}
if remoteConfig.Disk {
disk := monitoring.Disk()
data["disk"] = map[string]interface{}{
"total": disk.Total,
"used": disk.Used,
}
}
if remoteConfig.Network {
networkUp, networkDown, err := monitoring.NetworkSpeed()
if err != nil {
message += fmt.Sprintf("failed to get network speed: %v\n", err)
}
data["network"] = map[string]interface{}{
"up": networkUp,
"down": networkDown,
}
}
if remoteConfig.Connections {
tcpCount, udpCount, err := monitoring.ConnectionsCount()
if err != nil {
message += fmt.Sprintf("failed to get connections: %v\n", err)
}
data["network"] = map[string]interface{}{
"tcp": tcpCount,
"udp": udpCount,
}
}
if remoteConfig.Uptime {
uptime, err := monitoring.Uptime()
if err != nil {
message += fmt.Sprintf("failed to get uptime: %v\n", err)
}
data["uptime"] = uptime
}
data["message"] = message
s, err := json.Marshal(data)
if err != nil {
log.Println("Failed to marshal data:", err)
}
return s
}

53
monitoring/cpu.go Normal file
View File

@@ -0,0 +1,53 @@
package monitoring
import (
"runtime"
"strings"
"time"
"github.com/shirou/gopsutil/cpu"
)
type CpuInfo struct {
CPUName string `json:"cpu_name"`
CPUArchitecture string `json:"cpu_architecture"`
CPUCores int `json:"cpu_cores"`
CPUUsage float64 `json:"cpu_usage"`
}
func Cpu() CpuInfo {
cpuinfo := CpuInfo{}
info, err := cpu.Info()
if err != nil {
cpuinfo.CPUName = "Unknown"
}
// multiple CPU
// 多个 CPU
if len(info) > 1 {
for _, cpu := range info {
cpuinfo.CPUName += cpu.ModelName + ", "
}
cpuinfo.CPUName = cpuinfo.CPUName[:len(cpuinfo.CPUName)-2] // Remove trailing comma and space
} else if len(info) == 1 {
cpuinfo.CPUName = info[0].ModelName
}
cpuinfo.CPUName = strings.TrimSpace(cpuinfo.CPUName)
cpuinfo.CPUArchitecture = runtime.GOARCH
cores, err := cpu.Counts(true)
if err != nil {
cpuinfo.CPUCores = 1 // Error case
}
cpuinfo.CPUCores = cores
// Get CPU Usage
percentages, err := cpu.Percent(1*time.Second, false)
if err != nil {
cpuinfo.CPUUsage = 0.0 // Error case
} else {
cpuinfo.CPUUsage = percentages[0]
}
return cpuinfo
}

35
monitoring/disk.go Normal file
View File

@@ -0,0 +1,35 @@
package monitoring
import (
"github.com/shirou/gopsutil/disk"
)
type DiskInfo struct {
Total uint64 `json:"total"`
Used uint64 `json:"used"`
}
func Disk() DiskInfo {
diskinfo := DiskInfo{}
usage, err := disk.Partitions(true)
if err != nil {
diskinfo.Total = 0
diskinfo.Used = 0
} else {
for _, part := range usage {
if part.Mountpoint != "/tmp" && part.Mountpoint != "/var/tmp" && part.Mountpoint != "/dev/shm" {
// Skip /tmp, /var/tmp, and /dev/shm
// 获取磁盘使用情况
u, err := disk.Usage(part.Mountpoint)
if err != nil {
diskinfo.Total = 0
diskinfo.Used = 0
} else {
diskinfo.Total += u.Total
diskinfo.Used += u.Used
}
}
}
}
return diskinfo
}

1
monitoring/gpu.go Normal file
View File

@@ -0,0 +1 @@
package monitoring

25
monitoring/load.go Normal file
View File

@@ -0,0 +1,25 @@
package monitoring
import (
"github.com/shirou/gopsutil/load"
)
type LoadInfo struct {
Load1 float64 `json:"load_1"`
Load5 float64 `json:"load_5"`
Load15 float64 `json:"load_15"`
}
func Load() LoadInfo {
avg, err := load.Avg()
if err != nil {
return LoadInfo{Load1: 0, Load5: 0, Load15: 0}
}
return LoadInfo{
Load1: avg.Load1,
Load5: avg.Load5,
Load15: avg.Load15,
}
}

35
monitoring/mem.go Normal file
View File

@@ -0,0 +1,35 @@
package monitoring
import (
"github.com/shirou/gopsutil/mem"
)
type RamInfo struct {
Total uint64 `json:"total"`
Used uint64 `json:"used"`
}
func Ram() RamInfo {
raminfo := RamInfo{}
v, err := mem.VirtualMemory()
if err != nil {
raminfo.Total = 0
raminfo.Used = 0
} else {
raminfo.Total = v.Total
raminfo.Used = v.Used
}
return raminfo
}
func Swap() RamInfo {
swapinfo := RamInfo{}
s, err := mem.SwapMemory()
if err != nil {
swapinfo.Total = 0
swapinfo.Used = 0
} else {
swapinfo.Total = s.Total
swapinfo.Used = s.Used
}
return swapinfo
}

51
monitoring/net.go Normal file
View File

@@ -0,0 +1,51 @@
package monitoring
import (
"fmt"
"github.com/shirou/gopsutil/net"
)
func ConnectionsCount() (tcpCount, udpCount int, err error) {
tcps, err := net.Connections("tcp")
if err != nil {
return 0, 0, fmt.Errorf("failed to get TCP connections: %w", err)
}
udps, err := net.Connections("udp")
if err != nil {
return 0, 0, fmt.Errorf("failed to get UDP connections: %w", err)
}
return len(tcps), len(udps), nil
}
func NetworkSpeed() (upSpeed, downSpeed float64, err error) {
// Get the network IO counters
ioCounters, err := net.IOCounters(false)
if err != nil {
return 0, 0, fmt.Errorf("failed to get network IO counters: %w", err)
}
if len(ioCounters) == 0 {
return 0, 0, fmt.Errorf("no network interfaces found")
}
for _, interfaceStats := range ioCounters {
loopbackNames := []string{"lo", "lo0", "localhost", "brd0", "docker0", "docker1", "veth0", "veth1", "veth2", "veth3", "veth4", "veth5", "veth6", "veth7"}
isLoopback := false
for _, name := range loopbackNames {
if interfaceStats.Name == name {
isLoopback = true
break
}
}
if isLoopback {
continue // Skip loopback interface
}
upSpeed += float64(interfaceStats.BytesSent) / float64(interfaceStats.PacketsSent)
downSpeed += float64(interfaceStats.BytesRecv) / float64(interfaceStats.PacketsRecv)
}
return upSpeed, downSpeed, nil
}

49
monitoring/os.go Normal file
View File

@@ -0,0 +1,49 @@
package monitoring
import (
"bufio"
"os"
"runtime"
"strings"
"golang.org/x/sys/windows/registry"
)
func OSName() string {
if runtime.GOOS == "windows" {
key, err := registry.OpenKey(registry.LOCAL_MACHINE, `SOFTWARE\Microsoft\Windows NT\CurrentVersion`, registry.QUERY_VALUE)
if err != nil {
return "Microsoft Windows"
}
defer key.Close()
productName, _, err := key.GetStringValue("ProductName")
if err != nil {
return "Microsoft Windows"
}
return productName
} else if runtime.GOOS == "linux" {
file, err := os.Open("/etc/os-release")
if err != nil {
return "Linux"
}
defer file.Close()
scanner := bufio.NewScanner(file)
for scanner.Scan() {
line := scanner.Text()
if strings.HasPrefix(line, "PRETTY_NAME=") {
return strings.Trim(line[len("PRETTY_NAME="):], `"`)
}
}
if err := scanner.Err(); err != nil {
return "Linux"
}
return "Linux"
}
return "Unknown"
}

11
monitoring/uptime.go Normal file
View File

@@ -0,0 +1,11 @@
package monitoring
import (
"github.com/shirou/gopsutil/host"
)
func Uptime() (uint64, error) {
return host.Uptime()
}