545 lines
14 KiB
Go
545 lines
14 KiB
Go
package agent
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"os/signal"
|
|
"sync"
|
|
"syscall"
|
|
"time"
|
|
|
|
"github.com/fleetforge/forge-agent/internal/client"
|
|
"github.com/fleetforge/forge-agent/internal/config"
|
|
"github.com/fleetforge/forge-agent/internal/inventory"
|
|
"github.com/fleetforge/forge-agent/internal/metrics"
|
|
"github.com/fleetforge/forge-agent/internal/pkgmanager"
|
|
"github.com/fleetforge/forge-agent/internal/svcmgr"
|
|
"github.com/fleetforge/forge-agent/internal/tunnel"
|
|
"github.com/fleetforge/forge-agent/internal/vnc"
|
|
)
|
|
|
|
const AgentVersion = "0.1.6"
|
|
|
|
var (
|
|
tunnelMu sync.Mutex
|
|
tunnels = map[string]*tunnel.Tunnel{}
|
|
)
|
|
|
|
func storeTunnel(sessionID string, t *tunnel.Tunnel) {
|
|
tunnelMu.Lock()
|
|
defer tunnelMu.Unlock()
|
|
tunnels[sessionID] = t
|
|
}
|
|
|
|
func removeTunnel(sessionID string) {
|
|
tunnelMu.Lock()
|
|
defer tunnelMu.Unlock()
|
|
delete(tunnels, sessionID)
|
|
}
|
|
|
|
func stopTunnel(sessionID string) {
|
|
tunnelMu.Lock()
|
|
t := tunnels[sessionID]
|
|
tunnelMu.Unlock()
|
|
if t != nil {
|
|
t.Stop()
|
|
}
|
|
}
|
|
|
|
type Agent struct {
|
|
cfg *config.Config
|
|
client *client.Client
|
|
}
|
|
|
|
func New(cfg *config.Config) *Agent {
|
|
return &Agent{
|
|
cfg: cfg,
|
|
client: client.New(cfg.ServerURL, cfg.Credential),
|
|
}
|
|
}
|
|
|
|
func (a *Agent) Run() error {
|
|
log.Println("Forge Agent starting...")
|
|
log.Printf("Server: %s", a.cfg.ServerURL)
|
|
log.Printf("Device ID: %s", a.cfg.DeviceID)
|
|
|
|
ticker := time.NewTicker(time.Duration(a.cfg.HeartbeatSec) * time.Second)
|
|
defer ticker.Stop()
|
|
|
|
sigCh := make(chan os.Signal, 1)
|
|
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
|
|
|
|
if a.cfg.DeviceID == "" {
|
|
log.Println("No DeviceID configured - waiting for enrollment")
|
|
} else {
|
|
if err := pkgmanager.BootstrapManagers(); err != nil {
|
|
log.Printf("WARNING: Manager bootstrap failed: %v", err)
|
|
}
|
|
a.sendHeartbeat()
|
|
}
|
|
|
|
for {
|
|
select {
|
|
case <-ticker.C:
|
|
if a.cfg.DeviceID == "" || a.cfg.Credential == "" {
|
|
a.tryEnroll()
|
|
} else {
|
|
a.sendHeartbeat()
|
|
a.checkJobs()
|
|
}
|
|
case sig := <-sigCh:
|
|
log.Printf("Received signal %v, shutting down...", sig)
|
|
return nil
|
|
}
|
|
}
|
|
}
|
|
|
|
func (a *Agent) tryEnroll() {
|
|
if a.cfg.EnrollmentToken == "" {
|
|
return
|
|
}
|
|
log.Println("Attempting auto-enrollment...")
|
|
|
|
info, err := inventory.CollectSystemInfo()
|
|
if err != nil {
|
|
log.Printf("Enrollment failed (inventory): %v", err)
|
|
return
|
|
}
|
|
|
|
c := client.New(a.cfg.ServerURL, "")
|
|
resp, err := c.Enroll(client.EnrollRequest{
|
|
Token: a.cfg.EnrollmentToken,
|
|
Hostname: info.Hostname,
|
|
Platform: info.Platform,
|
|
OSName: info.OSName,
|
|
OSVersion: info.OSVersion,
|
|
Architecture: info.Architecture,
|
|
KernelVersion: info.KernelVersion,
|
|
CPUModel: info.CPUModel,
|
|
CPUCores: info.CPUCores,
|
|
TotalMemoryBytes: info.TotalMemory,
|
|
TotalDiskBytes: info.TotalDisk,
|
|
IPAddress: info.IPAddress,
|
|
AgentVersion: AgentVersion,
|
|
})
|
|
if err != nil {
|
|
log.Printf("Enrollment failed: %v", err)
|
|
return
|
|
}
|
|
|
|
a.cfg.DeviceID = resp.DeviceID
|
|
a.cfg.Credential = resp.Credential
|
|
a.cfg.EnrollmentToken = ""
|
|
if err := config.Save(a.cfg); err != nil {
|
|
log.Printf("Failed to save config after enrollment: %v", err)
|
|
return
|
|
}
|
|
a.client = client.New(a.cfg.ServerURL, a.cfg.Credential)
|
|
log.Printf("Enrollment successful! Device ID: %s", resp.DeviceID)
|
|
}
|
|
|
|
func (a *Agent) sendHeartbeat() {
|
|
info, err := inventory.CollectSystemInfo()
|
|
if err != nil {
|
|
log.Printf("Failed to collect system info: %v", err)
|
|
return
|
|
}
|
|
|
|
met, err := metrics.Collect()
|
|
if err != nil {
|
|
log.Printf("Failed to collect metrics: %v", err)
|
|
return
|
|
}
|
|
|
|
req := client.HeartbeatRequest{
|
|
Hostname: info.Hostname,
|
|
OSVersion: info.OSVersion,
|
|
KernelVersion: info.KernelVersion,
|
|
AgentVersion: AgentVersion,
|
|
IPAddress: info.IPAddress,
|
|
Metrics: &client.MetricsData{
|
|
CPUPercent: met.CPUPercent,
|
|
MemoryPercent: met.MemoryPercent,
|
|
MemoryUsed: met.MemoryUsed,
|
|
MemoryTotal: met.MemoryTotal,
|
|
DiskPercent: met.DiskPercent,
|
|
DiskUsed: met.DiskUsed,
|
|
DiskTotal: met.DiskTotal,
|
|
Load1: met.Load1,
|
|
Load5: met.Load5,
|
|
Load15: met.Load15,
|
|
UptimeSeconds: met.UptimeSeconds,
|
|
},
|
|
}
|
|
|
|
if len(info.Packages) > 0 || len(info.Services) > 0 || len(info.NetworkInterfaces) > 0 {
|
|
inv := &client.InventoryData{}
|
|
for _, p := range info.Packages {
|
|
inv.Packages = append(inv.Packages, client.PackageData{
|
|
Name: p.Name,
|
|
Version: p.Version,
|
|
Manager: p.Manager,
|
|
})
|
|
}
|
|
for _, s := range info.Services {
|
|
inv.Services = append(inv.Services, client.ServiceData{
|
|
Name: s.Name,
|
|
Status: s.Status,
|
|
Enabled: s.Enabled,
|
|
})
|
|
}
|
|
for _, ni := range info.NetworkInterfaces {
|
|
inv.NetworkInterfaces = append(inv.NetworkInterfaces, client.NetworkInterfaceData{
|
|
Name: ni.Name,
|
|
Addresses: ni.Addresses,
|
|
MacAddress: ni.MAC,
|
|
})
|
|
}
|
|
req.Inventory = inv
|
|
}
|
|
|
|
if err := a.client.Heartbeat(req); err != nil {
|
|
log.Printf("Heartbeat failed: %v", err)
|
|
} else {
|
|
log.Printf("Heartbeat sent successfully")
|
|
}
|
|
}
|
|
|
|
func (a *Agent) checkJobs() {
|
|
jobs, err := a.client.FetchJobs()
|
|
if err != nil {
|
|
log.Printf("Failed to fetch jobs: %v", err)
|
|
return
|
|
}
|
|
|
|
for _, job := range jobs {
|
|
a.executeJob(job)
|
|
}
|
|
}
|
|
|
|
func (a *Agent) executeJob(job client.JobInfo) {
|
|
log.Printf("Executing job %s (type: %s)", job.JobID, job.JobType)
|
|
|
|
result := client.JobResult{
|
|
Status: "succeeded",
|
|
Result: map[string]interface{}{},
|
|
}
|
|
result.TargetID = job.TargetID
|
|
|
|
if err := a.client.ReportStart(job.JobID, job.TargetID); err != nil {
|
|
log.Printf("Failed to report start for job %s: %v", job.JobID, err)
|
|
}
|
|
|
|
switch job.JobType {
|
|
case "update_inventory":
|
|
a.sendHeartbeat()
|
|
result.Result["inventory_refreshed"] = true
|
|
|
|
case "update_system_info":
|
|
met, err := metrics.Collect()
|
|
if err != nil {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = err.Error()
|
|
} else {
|
|
result.Result["cpu_percent"] = met.CPUPercent
|
|
result.Result["memory_percent"] = met.MemoryPercent
|
|
result.Result["disk_percent"] = met.DiskPercent
|
|
result.Result["uptime"] = met.UptimeSeconds
|
|
}
|
|
|
|
case "install_package":
|
|
pkgName, _ := job.Parameters["package_name"].(string)
|
|
if pkgName == "" {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = "missing parameter: package_name"
|
|
break
|
|
}
|
|
output, err := pkgmanager.InstallPackage(pkgName)
|
|
if err != nil {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = err.Error()
|
|
}
|
|
result.Result["package"] = pkgName
|
|
result.Result["output"] = truncateOutput(output)
|
|
|
|
case "update_package":
|
|
pkgName, _ := job.Parameters["package_name"].(string)
|
|
if pkgName == "" {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = "missing parameter: package_name"
|
|
break
|
|
}
|
|
output, err := pkgmanager.UpdatePackage(pkgName)
|
|
if err != nil {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = err.Error()
|
|
}
|
|
result.Result["package"] = pkgName
|
|
result.Result["output"] = truncateOutput(output)
|
|
|
|
case "remove_package":
|
|
pkgName, _ := job.Parameters["package_name"].(string)
|
|
if pkgName == "" {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = "missing parameter: package_name"
|
|
break
|
|
}
|
|
source, _ := job.Parameters["source"].(string)
|
|
var manager string
|
|
switch source {
|
|
case "winget":
|
|
manager = "winget"
|
|
case "chocolatey":
|
|
manager = "choco"
|
|
}
|
|
var output string
|
|
var err error
|
|
if manager != "" {
|
|
output, err = pkgmanager.RemovePackageWithManager(pkgName, manager)
|
|
} else {
|
|
output, err = pkgmanager.RemovePackage(pkgName)
|
|
}
|
|
if err != nil {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = err.Error()
|
|
}
|
|
result.Result["package"] = pkgName
|
|
result.Result["output"] = truncateOutput(output)
|
|
|
|
case "start_service":
|
|
svcName, _ := job.Parameters["service_name"].(string)
|
|
if svcName == "" {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = "missing parameter: service_name"
|
|
break
|
|
}
|
|
output, err := svcmgr.StartService(svcName)
|
|
if err != nil {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = err.Error()
|
|
}
|
|
result.Result["service"] = svcName
|
|
result.Result["output"] = truncateOutput(output)
|
|
|
|
case "stop_service":
|
|
svcName, _ := job.Parameters["service_name"].(string)
|
|
if svcName == "" {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = "missing parameter: service_name"
|
|
break
|
|
}
|
|
output, err := svcmgr.StopService(svcName)
|
|
if err != nil {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = err.Error()
|
|
}
|
|
result.Result["service"] = svcName
|
|
result.Result["output"] = truncateOutput(output)
|
|
|
|
case "restart_service":
|
|
svcName, _ := job.Parameters["service_name"].(string)
|
|
if svcName == "" {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = "missing parameter: service_name"
|
|
break
|
|
}
|
|
output, err := svcmgr.RestartService(svcName)
|
|
if err != nil {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = err.Error()
|
|
}
|
|
result.Result["service"] = svcName
|
|
result.Result["output"] = truncateOutput(output)
|
|
|
|
case "check_service_status":
|
|
svcName, _ := job.Parameters["service_name"].(string)
|
|
if svcName == "" {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = "missing parameter: service_name"
|
|
break
|
|
}
|
|
status, enabled, err := svcmgr.GetServiceStatus(svcName)
|
|
if err != nil {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = err.Error()
|
|
}
|
|
result.Result["service"] = svcName
|
|
result.Result["status"] = status
|
|
result.Result["enabled"] = enabled
|
|
|
|
case "reboot_system":
|
|
log.Println("System reboot requested, initiating reboot...")
|
|
result.Result["message"] = "reboot initiated"
|
|
if err := a.client.SubmitResult(job.JobID, result); err != nil {
|
|
log.Printf("Failed to submit result for reboot job: %v", err)
|
|
return
|
|
}
|
|
log.Println("Rebooting system in 5 seconds...")
|
|
time.Sleep(5 * time.Second)
|
|
svcmgr.Reboot()
|
|
return
|
|
|
|
case "shutdown_system":
|
|
log.Println("System shutdown requested, initiating shutdown...")
|
|
result.Result["message"] = "shutdown initiated"
|
|
if err := a.client.SubmitResult(job.JobID, result); err != nil {
|
|
log.Printf("Failed to submit result for shutdown job: %v", err)
|
|
return
|
|
}
|
|
log.Println("Shutting down system in 5 seconds...")
|
|
time.Sleep(5 * time.Second)
|
|
svcmgr.Shutdown()
|
|
return
|
|
|
|
case "install_registry_package":
|
|
source, _ := job.Parameters["source"].(string)
|
|
pkgID, _ := job.Parameters["package_id"].(string)
|
|
if source == "" || pkgID == "" {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = "missing parameter: source/package_id"
|
|
break
|
|
}
|
|
var manager string
|
|
switch source {
|
|
case "winget":
|
|
manager = "winget"
|
|
case "chocolatey":
|
|
manager = "choco"
|
|
default:
|
|
result.Status = "failed"
|
|
result.ErrorMessage = fmt.Sprintf("unsupported source: %s", source)
|
|
break
|
|
}
|
|
if result.Status == "failed" {
|
|
break
|
|
}
|
|
version, _ := job.Parameters["package_version"].(string)
|
|
output, err := pkgmanager.InstallPackageWithManager(pkgID, manager, version)
|
|
if err != nil {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = err.Error()
|
|
}
|
|
result.Result["package"] = pkgID
|
|
result.Result["source"] = source
|
|
result.Result["output"] = truncateOutput(output)
|
|
|
|
case "start_remote_session":
|
|
sessionID, _ := job.Parameters["session_id"].(string)
|
|
if sessionID == "" {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = "missing parameter: session_id"
|
|
break
|
|
}
|
|
port := 5900
|
|
if p, ok := job.Parameters["vnc_port"].(float64); ok && p > 0 {
|
|
port = int(p)
|
|
}
|
|
password, _ := job.Parameters["vnc_password"].(string)
|
|
|
|
vncServer, err := vnc.Ensure(port, password)
|
|
if err != nil {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = err.Error()
|
|
break
|
|
}
|
|
log.Printf("VNC server started on port %d for session %s", port, sessionID)
|
|
|
|
result.Result["session_id"] = sessionID
|
|
result.Result["vnc_port"] = port
|
|
result.Result["vnc_started"] = true
|
|
|
|
// Keep the session alive in a background goroutine; VNC is stopped when
|
|
// the tunnel closes.
|
|
go func(sid string, p int, vs *vnc.Server) {
|
|
t := tunnel.New(a.cfg.ServerURL, a.cfg.Credential, sid, p)
|
|
storeTunnel(sid, t)
|
|
defer func() {
|
|
removeTunnel(sid)
|
|
vs.Stop()
|
|
log.Printf("VNC server stopped after tunnel close for session %s", sid)
|
|
}()
|
|
if err := t.Run(); err != nil {
|
|
log.Printf("Tunnel for session %s closed: %v", sid, err)
|
|
}
|
|
}(sessionID, port, vncServer)
|
|
|
|
case "stop_remote_session":
|
|
sessionID, _ := job.Parameters["session_id"].(string)
|
|
if sessionID == "" {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = "missing parameter: session_id"
|
|
break
|
|
}
|
|
stopTunnel(sessionID)
|
|
result.Result["session_id"] = sessionID
|
|
result.Result["vnc_stopped"] = true
|
|
|
|
case "install_from_file":
|
|
filePath, _ := job.Parameters["file_path"].(string)
|
|
if filePath == "" {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = "missing parameter: file_path"
|
|
break
|
|
}
|
|
output, err := pkgmanager.InstallPackageFromFile(filePath)
|
|
if err != nil {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = err.Error()
|
|
}
|
|
result.Result["file_path"] = filePath
|
|
result.Result["output"] = truncateOutput(output)
|
|
|
|
case "download_and_install":
|
|
downloadURL, _ := job.Parameters["download_url"].(string)
|
|
packageId, _ := job.Parameters["package_id"].(string)
|
|
if downloadURL == "" && packageId == "" {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = "missing parameter: download_url or package_id"
|
|
break
|
|
}
|
|
if downloadURL == "" && packageId != "" {
|
|
downloadURL = a.cfg.ServerURL + "/api/v1/agent/software/" + packageId + "/download"
|
|
}
|
|
tmpDir := os.TempDir()
|
|
destPath := tmpDir + "/forge-download-" + packageId + ".pkg"
|
|
log.Printf("Downloading package from %s to %s", downloadURL, destPath)
|
|
if err := pkgmanager.DownloadFile(downloadURL, destPath); err != nil {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = fmt.Sprintf("download failed: %v", err)
|
|
break
|
|
}
|
|
output, err := pkgmanager.InstallPackageFromFile(destPath)
|
|
if err != nil {
|
|
result.Status = "failed"
|
|
result.ErrorMessage = err.Error()
|
|
}
|
|
os.Remove(destPath)
|
|
result.Result["download_url"] = downloadURL
|
|
result.Result["output"] = truncateOutput(output)
|
|
|
|
default:
|
|
result.Status = "failed"
|
|
result.ErrorMessage = fmt.Sprintf("unsupported job type: %s", job.JobType)
|
|
}
|
|
|
|
if a.client.IsJobCancelled(job.JobID) {
|
|
log.Printf("Job %s cancelled on server; discarding result", job.JobID)
|
|
result.Status = "cancelled"
|
|
}
|
|
|
|
if err := a.client.SubmitResult(job.JobID, result); err != nil {
|
|
log.Printf("Failed to submit result for job %s: %v", job.JobID, err)
|
|
} else {
|
|
log.Printf("Job %s completed with status: %s", job.JobID, result.Status)
|
|
}
|
|
}
|
|
|
|
func truncateOutput(output string) string {
|
|
const maxLen = 4096
|
|
if len(output) > maxLen {
|
|
return output[:maxLen] + "... (truncated)"
|
|
}
|
|
return output
|
|
}
|