Sigma-C2/server/agent_manager.go

package main

import (
	"log"
	"sync"
	"time"
)

type AgentInfo struct {
	Type           string
	ID             string
	Listener       string
	OSVersion      string
	Architecture   string
	Hostname       string
	Username       string
	LocalIP        string
	Procname       string
	PID            string
	RemoteIP       string
	LastConnection time.Time
}

var agents sync.Map

// Update or create agent information with new data
func UpdateAgent(agentID string, newInfo *AgentInfo) {
	if existingInfo, exists := GetAgent(agentID); exists {
		// Update existing agent's fields with new data if provided
		if newInfo.OSVersion != "" {
			existingInfo.OSVersion = newInfo.OSVersion
		}
		if newInfo.Architecture != "" {
			existingInfo.Architecture = newInfo.Architecture
		}
		if newInfo.Hostname != "" {
			existingInfo.Hostname = newInfo.Hostname
		}
		if newInfo.Username != "" {
			existingInfo.Username = newInfo.Username
		}
		if newInfo.LocalIP != "" {
			existingInfo.LocalIP = newInfo.LocalIP
		}
		if newInfo.Procname != "" {
			existingInfo.Procname = newInfo.Procname
		}
		if newInfo.PID != "" {
			existingInfo.PID = newInfo.PID
		}
		if newInfo.RemoteIP != "" {
			existingInfo.RemoteIP = newInfo.RemoteIP
		}
		// Always update the last connection time
		existingInfo.LastConnection = time.Now()

		agents.Store(agentID, existingInfo)
		log.Printf("Updated agent entry for %s", agentID)
	} else {
		// If agent doesn't exist, use the newInfo struct to create a new agent
		agents.Store(agentID, newInfo)
		log.Printf("Created new agent entry for %s", agentID)
		// Send new list of agents to operators to update tab completion
		UpdateAgentsTabCompletion()
	}
}

// Get agent info by name
func GetAgent(agentID string) (*AgentInfo, bool) {
	value, exists := agents.Load(agentID)
	if !exists {
		return nil, false
	}
	return value.(*AgentInfo), true
}

// Get agents list of agents' names
func GetAgentNames() []string {
	var names []string
	agents.Range(func(key, value any) bool {
		names = append(names, key.(string))
		return true
	})
	return names
}