agent/main.go - master - Virtual/Termina

agent/main.go
224 строки5 KB
Virtes
Fixes and Makefile addition
19 июн 2026, 12:47
19 июн 2026, 12:47776a0d3
package main

import (
	"database/sql"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"termina/pkg/utils"
	"time"

	"github.com/shirou/gopsutil/cpu"
	"github.com/shirou/gopsutil/disk"
	"github.com/shirou/gopsutil/mem"
	"github.com/shirou/gopsutil/net"
	"go.mongodb.org/mongo-driver/bson"

	_ "modernc.org/sqlite"
)

type Snapshot struct {
	Hostname    string      `bson:"hostname"`
	CollectedAt time.Time   `bson:"collected_at"`
	CPUModel    string      `bson:"cpu_model"`
	CPUCores    int         `bson:"cpu_cores"`
	RAMTotal    uint64      `bson:"ram_total"`
	RAMUsed     uint64      `bson:"ram_used"`
	RAMPercent  float64     `bson:"ram_percent"`
	Network     []NetInfo   `bson:"network"`
	Partitions  []Partition `bson:"partitions"`
}

type NetInfo struct {
	Name  string   `bson:"name"`
	MAC   string   `bson:"mac"`
	Addrs []string `bson:"addrs"`
}

type Partition struct {
	Mountpoint string `bson:"mountpoint"`
	TotalMB    uint64 `bson:"total_mb"`
	FreeMB     uint64 `bson:"free_mb"`
	UsedMB     uint64 `bson:"used_mb"`
	Fstype     string `bson:"fstype"`
}

func main() {
	fmt.Println("Init")

	logFiles := setupLogs()
	defer func(logFiles *os.File) {
		err := logFiles.Close()
		if err != nil {
			log.Fatal(err)
		}
	}(logFiles)

	db := setupDB()
	defer func(db *sql.DB) {
		err := db.Close()
		if err != nil {
			log.Fatal(err)
		}
	}(db)

	collectAndSave(db)

	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for range ticker.C {
		collectAndSave(db)
	}
}

func setupLogs() *os.File {
	fmt.Println("Setup")
	dir := utils.GetDataDir()
	logPath := filepath.Join(dir, "agent.log")
	logFile, err := os.OpenFile(logPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
	if err != nil {
		log.Printf("error while creating log, continuing > %v", err)
	}
	log.SetOutput(logFile)
	fmt.Println("Setup logs complete")
	return logFile
}

func setupDB() *sql.DB {
	dir := utils.GetDataDir()
	dbPath := filepath.Join(dir, "agent.db")

	db, err := sql.Open("sqlite", dbPath)
	if err != nil {
		log.Fatalf("database start error > %v\n", err)
	}
	_, err = db.Exec(`
		CREATE TABLE IF NOT EXISTS snapshots (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			device_id TEXT NOT NULL,
			collected_at DATETIME NOT NULL,
			hostname TEXT,
			raw_data BLOB NOT NULL	
		)
	`)
	if err != nil {
		log.Fatalf("Error while creating sqlite for BSON [EXITING] > %v", err)
	}
	fmt.Println("Got DB")
	return db
}
func collectAndSave(db *sql.DB) {
	//HOST & OS
	host, err := os.Hostname()
	if err != nil {
		log.Printf("host %s", err)
		return
	}

	//CPU
	cpuInfos, err := cpu.Info()
	if err != nil {
		log.Printf("cpu %s", err)
		return
	}
	cpuModelName := cpuInfos[0].ModelName
	cpuCores := cpuInfos[0].Cores
	if len(cpuInfos) == 0 {
		log.Printf("cpu info is empty")
		return
	}
	//RAM
	ramInfo, err := mem.VirtualMemory()
	if err != nil {
		log.Printf("ram %s", err)
		return
	}
	ramTotal := ramInfo.Total / (1024 * 1024)
	ramUsed := ramInfo.Used / (1024 * 1024)
	ramUsedPercent := ramInfo.UsedPercent

	//NET
	var netInfos []NetInfo
	netInterfaces, err := net.Interfaces()
	if err != nil {
		log.Printf("net %s", err)
	}

	for _, i := range netInterfaces {
		var addrs []string
		for _, ip := range i.Addrs {
			addrs = append(addrs, ip.Addr)
		}
		netInfos = append(netInfos, NetInfo{
			Name:  i.Name,
			MAC:   i.HardwareAddr,
			Addrs: addrs,
		})
	}

	//DISKS
	var partitions []Partition
	partitionsInfo, err := disk.Partitions(false)
	if err != nil {
		log.Printf("disks %s", err)
		return
	}
	for _, p := range partitionsInfo {
		partitionUsage, err := disk.Usage(p.Mountpoint)
		if err != nil {
			log.Printf("disk %s %s", p.Mountpoint, err)
			continue
		}
		partitionUsageTotal := partitionUsage.Total / (1024 * 1024)
		partitionUsageFree := partitionUsage.Free / (1024 * 1024)
		partitionUsageUsed := partitionUsage.Used / (1024 * 1024)
		partitions = append(partitions, Partition{
			Mountpoint: p.Mountpoint,
			TotalMB:    partitionUsageTotal,
			FreeMB:     partitionUsageFree,
			UsedMB:     partitionUsageUsed,
			Fstype:     p.Fstype,
		})
	}

	snapshot := Snapshot{
		Hostname:    host,
		CollectedAt: time.Now(),
		CPUModel:    cpuModelName,
		CPUCores:    int(cpuCores),
		RAMTotal:    ramTotal,
		RAMUsed:     ramUsed,
		RAMPercent:  ramUsedPercent,
		Network:     netInfos,
		Partitions:  partitions,
	}

	//	TO BSON
	bsonBytes, err := bson.Marshal(snapshot)
	if err != nil {
		log.Printf("BSON encode error%v", err)
		return
	}
	log.Printf("Length of BSON stream is > %d\n", len(bsonBytes))

	result, err := db.Exec(
		"INSERT INTO snapshots (device_id, collected_at, hostname, raw_data) VALUES (?, ?, ?, ?)",
		snapshot.Hostname,
		snapshot.CollectedAt,
		snapshot.Hostname,
		bsonBytes,
	)
	if err != nil {
		log.Printf("Insert error > %v", err)
		return
	}
	lastID, err := result.LastInsertId()
	if err != nil {
		log.Printf("Error fetching LASTID, maybe we are missing the agent.db somehow? %v", err)
		return
	}
	log.Printf("Snapshot saved under ID `%v`", lastID)
	fmt.Println("Got info and sent")
}