86 lines
2.0 KiB
Go
86 lines
2.0 KiB
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
|
|
}
|