diff --git a/cmd/admin/main.go b/cmd/admin/main.go index a7647037..6bf84bd6 100644 --- a/cmd/admin/main.go +++ b/cmd/admin/main.go @@ -195,7 +195,7 @@ func osctrlAdminService() { log.Info().Msg("Initialize settings") settingsmgr = settings.NewSettings(db.Conn) log.Info().Msg("Initialize nodes") - nodesmgr = nodes.CreateNodes(db.Conn, redis.Client) + nodesmgr = nodes.CreateNodes(db.Conn) log.Info().Msg("Initialize queries") queriesmgr = queries.CreateQueries(db.Conn) log.Info().Msg("Initialize carves") diff --git a/cmd/api/main.go b/cmd/api/main.go index 5af7ef61..f9a53435 100644 --- a/cmd/api/main.go +++ b/cmd/api/main.go @@ -179,7 +179,7 @@ func osctrlAPIService() { log.Info().Msg("Initialize settings") settingsmgr = settings.NewSettings(db.Conn) log.Info().Msg("Initialize nodes") - nodesmgr = nodes.CreateNodes(db.Conn, redis.Client) + nodesmgr = nodes.CreateNodes(db.Conn) log.Info().Msg("Initialize queries") queriesmgr = queries.CreateQueries(db.Conn) log.Info().Msg("Initialize carves") diff --git a/cmd/cli/main.go b/cmd/cli/main.go index 226b78a2..eaab5610 100644 --- a/cmd/cli/main.go +++ b/cmd/cli/main.go @@ -1800,7 +1800,7 @@ func cliWrapper(action func(*cli.Context) error) func(*cli.Context) error { // Initialize settings settingsmgr = settings.NewSettings(db.Conn) // Initialize nodes - nodesmgr = nodes.CreateNodes(db.Conn, nil) + nodesmgr = nodes.CreateNodes(db.Conn) // Initialize queries queriesmgr = queries.CreateQueries(db.Conn) // Initialize carves diff --git a/cmd/tls/main.go b/cmd/tls/main.go index 782ea91e..0db85314 100644 --- a/cmd/tls/main.go +++ b/cmd/tls/main.go @@ -181,7 +181,7 @@ func osctrlService() { log.Info().Msg("Initialize settings") settingsmgr = settings.NewSettings(db.Conn) log.Info().Msg("Initialize nodes") - nodesmgr = nodes.CreateNodes(db.Conn, redis.Client) + nodesmgr = nodes.CreateNodes(db.Conn) log.Info().Msg("Initialize tags") tagsmgr = tags.CreateTagManager(db.Conn) log.Info().Msg("Initialize queries") diff --git a/pkg/cache/cache.go b/pkg/cache/cache.go index 78b2a277..884e38f9 100644 --- a/pkg/cache/cache.go +++ b/pkg/cache/cache.go @@ -5,7 +5,6 @@ import ( "fmt" redis "github.com/go-redis/redis/v8" - "github.com/jmpsec/osctrl/pkg/types" "github.com/spf13/viper" ) @@ -30,13 +29,6 @@ type JSONConfigurationRedis struct { ConnRetry int `json:"connRetry"` } -// CachedQueryWriteData to store in cache query logs -type CachedQueryWriteData struct { - UnixTime int `json:"unixTime"` - HostIdentifier string `json:"hostIdentifier"` - QueryData types.QueryWriteData -} - // LoadConfiguration to load the redis configuration file and assign to variables func LoadConfiguration(file, key string) (JSONConfigurationRedis, error) { var config JSONConfigurationRedis diff --git a/pkg/nodes/cache.go b/pkg/nodes/cache.go deleted file mode 100644 index b9a24ced..00000000 --- a/pkg/nodes/cache.go +++ /dev/null @@ -1,139 +0,0 @@ -package nodes - -import ( - "context" - "math" - "strconv" - "time" -) - -// Keys to identify the node fields in the cache -const ( - kID = "id" - kCreated = "created_at" - kNodeKey = "node_key" - kUUID = "uuid" - kPlatform = "platform" - kPlatformVersion = "platform_version" - kOsqueryVersion = "osquery_version" - kHostname = "hostname" - kIPAddress = "ip_address" - kLocalname = "localname" - kUsername = "username" - kOsqueryUser = "osquery_user" - kEnvironment = "environment" - kEnvironmentID = "environment_id" - kHardwareSerial = "hardware_serial" - kCPU = "cpu" - kMemory = "memory" - kLastSeen = "last_seen" - kBytesReceived = "bytes_received" -) - -// SetFullCached sets the provided full node in the cache, default expiration is 1 hour -func (n *NodeManager) SetFullCached(node OsqueryNode, ctx context.Context) error { - k := CacheFullKeyByUUID(node) - if err := n.Cache.HSet(ctx, k, map[string]interface{}{ - kID: node.ID, - kCreated: node.CreatedAt.Format(time.RFC3339), - kNodeKey: node.NodeKey, - kUUID: node.UUID, - kPlatform: node.Platform, - kPlatformVersion: node.PlatformVersion, - kOsqueryVersion: node.OsqueryVersion, - kHostname: node.Hostname, - kIPAddress: node.IPAddress, - kLocalname: node.Localname, - kUsername: node.Username, - kOsqueryUser: node.OsqueryUser, - kEnvironment: node.Environment, - kEnvironmentID: node.EnvironmentID, - kHardwareSerial: node.HardwareSerial, - kCPU: node.CPU, - kMemory: node.Memory, - kLastSeen: time.Now().Format(time.RFC3339), - kBytesReceived: node.BytesReceived, - }).Err(); err != nil { - return err - } - return n.Cache.ExpireAt(ctx, k, time.Now().Add(1*time.Hour)).Err() -} - -// SetPartialCached sets the provided partial node in the cache, default expiration is 1 hour -func (n *NodeManager) SetPartialCached(node OsqueryNode, ctx context.Context) error { - k := CachePartialKeyByNodeKey(node) - if err := n.Cache.HSet(ctx, k, map[string]interface{}{ - kID: node.ID, - kNodeKey: node.NodeKey, - kUUID: node.UUID, - kIPAddress: node.IPAddress, - }).Err(); err != nil { - return err - } - return n.Cache.ExpireAt(ctx, k, time.Now().Add(1*time.Hour)).Err() -} - -// GetFullFromCache returns the full node from the cache by node UUID and environment ID -func (n *NodeManager) GetFullFromCache(uuid string, envID uint, ctx context.Context) (OsqueryNode, error) { - var node OsqueryNode - res, err := n.Cache.HGetAll(ctx, CacheFullKeyRaw(uuid, envID)).Result() - if err != nil { - return node, err - } - node.CreatedAt, err = time.Parse(time.RFC3339, res[kCreated]) - if err != nil { - node.CreatedAt = time.Time{} - } - node.LastSeen, err = time.Parse(time.RFC3339, res[kLastSeen]) - if err != nil { - node.LastSeen = time.Time{} - } - resID, err := strconv.ParseUint(res[kID], 10, 32) - if err != nil || resID > math.MaxInt { - node.ID = 0 - } - node.ID = uint(resID) - node.NodeKey = res[kNodeKey] - node.UUID = res[kUUID] - node.Platform = res[kPlatform] - node.PlatformVersion = res[kPlatformVersion] - node.OsqueryVersion = res[kOsqueryVersion] - node.Hostname = res[kHostname] - node.IPAddress = res[kIPAddress] - node.Localname = res[kLocalname] - node.Username = res[kUsername] - node.OsqueryUser = res[kOsqueryUser] - node.Environment = res[kEnvironment] - resEnvID, err := strconv.ParseUint(res[kEnvironmentID], 10, 32) - if err != nil || resEnvID > math.MaxInt { - resEnvID = 0 - } - node.EnvironmentID = uint(resEnvID) - node.HardwareSerial = res[kHardwareSerial] - node.CPU = res[kCPU] - node.Memory = res[kMemory] - resBytes, err := strconv.ParseInt(res[kBytesReceived], 10, 32) - if err != nil || resBytes > math.MaxInt { - resBytes = 0 - } - node.BytesReceived = int(resBytes) - return node, nil -} - -// GetPartialFromCache returns the partial node from the cache by node key and environment ID -func (n *NodeManager) GetPartialFromCache(nodeKey string, envID uint, ctx context.Context) (OsqueryNode, error) { - var node OsqueryNode - res, err := n.Cache.HGetAll(ctx, CachePartialKeyRaw(nodeKey, envID)).Result() - if err != nil { - return node, err - } - resID, err := strconv.ParseUint(res[kID], 10, 32) - if err != nil || resID > math.MaxInt { - node.ID = 0 - } - node.ID = uint(resID) - node.NodeKey = res[kNodeKey] - node.UUID = res[kUUID] - node.IPAddress = res[kIPAddress] - return node, nil -} diff --git a/pkg/nodes/metadata.go b/pkg/nodes/metadata.go deleted file mode 100644 index 77d55400..00000000 --- a/pkg/nodes/metadata.go +++ /dev/null @@ -1,16 +0,0 @@ -package nodes - -// NodeMetadata to hold metadata for a node -type NodeMetadata struct { - IPAddress string - Username string - OsqueryUser string - Hostname string - Localname string - ConfigHash string - DaemonHash string - OsqueryVersion string - Platform string - PlatformVersion string - BytesReceived int -} diff --git a/pkg/nodes/models.go b/pkg/nodes/models.go new file mode 100644 index 00000000..e1192fad --- /dev/null +++ b/pkg/nodes/models.go @@ -0,0 +1,77 @@ +package nodes + +import ( + "time" + + "gorm.io/gorm" +) + +// OsqueryNode as abstraction of a node +type OsqueryNode struct { + gorm.Model + NodeKey string `gorm:"index"` + UUID string `gorm:"index"` + Platform string + PlatformVersion string + OsqueryVersion string + Hostname string + Localname string + IPAddress string + Username string + OsqueryUser string + Environment string + CPU string + Memory string + HardwareSerial string + DaemonHash string + ConfigHash string + BytesReceived int + RawEnrollment string + LastSeen time.Time + UserID uint + EnvironmentID uint + ExtraData string +} + +// ArchiveOsqueryNode as abstraction of an archived node +type ArchiveOsqueryNode struct { + gorm.Model + NodeKey string `gorm:"index"` + UUID string `gorm:"index"` + Trigger string + Platform string + PlatformVersion string + OsqueryVersion string + Hostname string + Localname string + IPAddress string + Username string + OsqueryUser string + Environment string + CPU string + Memory string + HardwareSerial string + ConfigHash string + DaemonHash string + BytesReceived int + RawEnrollment string + LastSeen time.Time + UserID uint + EnvironmentID uint + ExtraData string +} + +// NodeMetadata to hold metadata for a node +type NodeMetadata struct { + IPAddress string + Username string + OsqueryUser string + Hostname string + Localname string + ConfigHash string + DaemonHash string + OsqueryVersion string + Platform string + PlatformVersion string + BytesReceived int +} diff --git a/pkg/nodes/node-cache.go b/pkg/nodes/node-cache.go new file mode 100644 index 00000000..aaddcf78 --- /dev/null +++ b/pkg/nodes/node-cache.go @@ -0,0 +1,80 @@ +package nodes + +import ( + "context" + "time" + + "github.com/jmpsec/osctrl/pkg/cache" +) + +const ( + cacheName = "nodes" + // Default time-to-live for cached nodes + defaultTTL = 60 * time.Minute + // Default cleanup interval for the cache + defaultCleanupInterval = 30 * time.Minute +) + +// NodeCache provides cached access to OsqueryNode objects +type NodeCache struct { + // The cache itself, storing OsqueryNode objects + cache *cache.MemoryCache[OsqueryNode] + + // Reference to the node manager for cache misses + nodes *NodeManager +} + +// NewNodeCache creates a new node cache +func NewNodeCache(nodes *NodeManager) *NodeCache { + // Create a new cache with appropriate cleanup interval + nodeCache := cache.NewMemoryCache( + cache.WithCleanupInterval[OsqueryNode](defaultCleanupInterval), + cache.WithName[OsqueryNode](cacheName), + ) + + return &NodeCache{ + cache: nodeCache, + nodes: nodes, + } +} + +// GetByKey retrieves a node by node_key, using cache when available +func (nc *NodeCache) GetByKey(ctx context.Context, nodeKey string) (OsqueryNode, error) { + // Try to get from cache first + if node, found := nc.cache.Get(ctx, nodeKey); found { + return node, nil + } + + // Not in cache, fetch from database + node, err := nc.nodes.getByKeyFromDB(nodeKey) + if err != nil { + return OsqueryNode{}, err + } + + // Store in cache for future requests + nc.cache.Set(ctx, nodeKey, node, defaultTTL) + + return node, nil +} + +// InvalidateNode removes a specific node from the cache +func (nc *NodeCache) InvalidateNode(ctx context.Context, nodeKey string) { + nc.cache.Delete(ctx, nodeKey) +} + +// InvalidateAll clears the entire cache +func (nc *NodeCache) InvalidateAll(ctx context.Context) { + nc.cache.Clear(ctx) +} + +// UpdateNodeInCache updates a node in the cache +func (nc *NodeCache) UpdateNodeInCache(ctx context.Context, node OsqueryNode) { + nc.cache.Set(ctx, node.NodeKey, node, defaultTTL) +} + +// Close stops the cleanup goroutine and releases resources +func (nc *NodeCache) Close() { + if nc.cache != nil { + nc.cache.Stop() + } +} diff --git a/pkg/nodes/nodes.go b/pkg/nodes/nodes.go index 0c633f24..0cff8958 100644 --- a/pkg/nodes/nodes.go +++ b/pkg/nodes/nodes.go @@ -1,11 +1,11 @@ package nodes import ( + "context" "fmt" "strings" "time" - redis "github.com/go-redis/redis/v8" "github.com/rs/zerolog/log" "gorm.io/gorm" ) @@ -19,61 +19,6 @@ const ( AllNodes = "all" ) -// OsqueryNode as abstraction of a node -type OsqueryNode struct { - gorm.Model - NodeKey string `gorm:"index"` - UUID string `gorm:"index"` - Platform string - PlatformVersion string - OsqueryVersion string - Hostname string - Localname string - IPAddress string - Username string - OsqueryUser string - Environment string - CPU string - Memory string - HardwareSerial string - DaemonHash string - ConfigHash string - BytesReceived int - RawEnrollment string - LastSeen time.Time - UserID uint - EnvironmentID uint - ExtraData string -} - -// ArchiveOsqueryNode as abstraction of an archived node -type ArchiveOsqueryNode struct { - gorm.Model - NodeKey string `gorm:"index"` - UUID string `gorm:"index"` - Trigger string - Platform string - PlatformVersion string - OsqueryVersion string - Hostname string - Localname string - IPAddress string - Username string - OsqueryUser string - Environment string - CPU string - Memory string - HardwareSerial string - ConfigHash string - DaemonHash string - BytesReceived int - RawEnrollment string - LastSeen time.Time - UserID uint - EnvironmentID uint - ExtraData string -} - // StatsData to display node stats type StatsData struct { Total int64 `json:"total"` @@ -84,14 +29,13 @@ type StatsData struct { // NodeManager to handle all nodes of the system type NodeManager struct { DB *gorm.DB - Cache *redis.Client + Cache *NodeCache } // CreateNodes to initialize the nodes struct and its tables -func CreateNodes(backend *gorm.DB, cache *redis.Client) *NodeManager { +func CreateNodes(backend *gorm.DB) *NodeManager { var n *NodeManager = &NodeManager{ - DB: backend, - Cache: cache, + DB: backend, } // table osquery_nodes if err := backend.AutoMigrate(&OsqueryNode{}); err != nil { @@ -101,17 +45,11 @@ func CreateNodes(backend *gorm.DB, cache *redis.Client) *NodeManager { if err := backend.AutoMigrate(&ArchiveOsqueryNode{}); err != nil { log.Fatal().Msgf("Failed to AutoMigrate table (archive_osquery_nodes): %v", err) } + // Create and initialize the cache + n.Cache = NewNodeCache(n) return n } -// CheckByKey to check if node exists by node_key -// node_key is expected lowercase -func (n *NodeManager) CheckByKey(nodeKey string) bool { - var results int64 - n.DB.Model(&OsqueryNode{}).Where("node_key = ?", strings.ToLower(nodeKey)).Count(&results) - return (results > 0) -} - // CheckByUUID to check if node exists by UUID // UUID is expected uppercase func (n *NodeManager) CheckByUUID(uuid string) bool { @@ -128,14 +66,6 @@ func (n *NodeManager) CheckByUUIDEnv(uuid, environment string) bool { return (results > 0) } -// CheckByUUIDEnvID to check if node exists by UUID in a specific environment -// UUID is expected uppercase -func (n *NodeManager) CheckByUUIDEnvID(uuid string, envID int) bool { - var results int64 - n.DB.Model(&OsqueryNode{}).Where("uuid = ? AND environment_id = ?", strings.ToUpper(uuid), envID).Count(&results) - return (results > 0) -} - // CheckByHost to check if node exists by Hostname func (n *NodeManager) CheckByHost(host string) bool { var results int64 @@ -143,9 +73,10 @@ func (n *NodeManager) CheckByHost(host string) bool { return (results > 0) } -// GetByKey to retrieve full node object from DB, by node_key +// getByKeyFromDB to retrieve full node object directly from DB, by node_key +// This is used by the cache system on cache misses // node_key is expected lowercase -func (n *NodeManager) GetByKey(nodekey string) (OsqueryNode, error) { +func (n *NodeManager) getByKeyFromDB(nodekey string) (OsqueryNode, error) { var node OsqueryNode if err := n.DB.Where("node_key = ?", strings.ToLower(nodekey)).First(&node).Error; err != nil { return node, err @@ -153,6 +84,14 @@ func (n *NodeManager) GetByKey(nodekey string) (OsqueryNode, error) { return node, nil } +// GetByKey to retrieve full node object from DB or cache, by node_key +// node_key is expected lowercase +func (n *NodeManager) GetByKey(nodekey string) (OsqueryNode, error) { + // Currently, the the cache would not be updated frequently + // It should only be used for fetching the node object that is rarely updated + return n.Cache.GetByKey(context.Background(), strings.ToLower(nodekey)) +} + // GetByIdentifier to retrieve full node object from DB, by uuid or hostname or localname // UUID is expected uppercase func (n *NodeManager) GetByIdentifier(identifier string) (OsqueryNode, error) { @@ -191,51 +130,46 @@ func (n *NodeManager) GetByUUIDEnv(uuid string, envid uint) (OsqueryNode, error) // GetBySelector to retrieve target nodes by selector func (n *NodeManager) GetBySelector(stype, selector, target string, hours int64) ([]OsqueryNode, error) { var nodes []OsqueryNode - var s string + var column string + switch stype { case "environment": - s = "environment" + column = "environment" case "platform": - s = "platform" - } - switch target { - case AllNodes: - if err := n.DB.Where(s+" = ?", selector).Find(&nodes).Error; err != nil { - return nodes, err - } - case ActiveNodes: - // if err := n.DB.Where(s+" = ?", selector).Where("updated_at > ?", time.Now().AddDate(0, 0, -3)).Find(&nodes).Error; err != nil { - if err := n.DB.Where(s+" = ?", selector).Where("updated_at > ?", time.Now().Add(time.Duration(hours)*time.Hour)).Find(&nodes).Error; err != nil { - return nodes, err - } - case InactiveNodes: - // if err := n.DB.Where(s+" = ?", selector).Where("updated_at < ?", time.Now().AddDate(0, 0, -3)).Find(&nodes).Error; err != nil { - if err := n.DB.Where(s+" = ?", selector).Where("updated_at < ?", time.Now().Add(time.Duration(hours)*time.Hour)).Find(&nodes).Error; err != nil { - return nodes, err - } + column = "platform" + default: + return nodes, fmt.Errorf("invalid selector type: %s", stype) } + + // Build query with base condition + query := n.DB.Where(column+" = ?", selector) + + // Apply active/inactive filtering + query = ApplyNodeTarget(query, target, hours) + + // Execute query + if err := query.Find(&nodes).Error; err != nil { + return nodes, err + } + return nodes, nil } // Gets to retrieve all/active/inactive nodes func (n *NodeManager) Gets(target string, hours int64) ([]OsqueryNode, error) { var nodes []OsqueryNode - switch target { - case AllNodes: - if err := n.DB.Find(&nodes).Error; err != nil { - return nodes, err - } - case ActiveNodes: - // if err := n.DB.Where("updated_at > ?", time.Now().AddDate(0, 0, -3)).Find(&nodes).Error; err != nil { - if err := n.DB.Where("updated_at > ?", time.Now().Add(time.Duration(hours)*time.Hour)).Find(&nodes).Error; err != nil { - return nodes, err - } - case InactiveNodes: - // if err := n.DB.Where("updated_at < ?", time.Now().AddDate(0, 0, -3)).Find(&nodes).Error; err != nil { - if err := n.DB.Where("updated_at < ?", time.Now().Add(time.Duration(hours)*time.Hour)).Find(&nodes).Error; err != nil { - return nodes, err - } + + // Start with base query + query := n.DB + + // Apply active/inactive filtering + query = ApplyNodeTarget(query, target, hours) + + // Execute query + if err := query.Find(&nodes).Error; err != nil { + return nodes, err } + return nodes, nil } @@ -279,36 +213,14 @@ func (n *NodeManager) GetEnvPlatforms(environment string) ([]string, error) { return platforms, nil } -// GetStatsByEnv to populate table stats about nodes by environment. Active machine is < 3 days +// GetStatsByEnv to populate table stats about nodes by environment func (n *NodeManager) GetStatsByEnv(environment string, hours int64) (StatsData, error) { - var stats StatsData - if err := n.DB.Model(&OsqueryNode{}).Where("environment = ?", environment).Count(&stats.Total).Error; err != nil { - return stats, err - } - tHours := time.Now().Add(time.Duration(hours) * time.Hour) - if err := n.DB.Model(&OsqueryNode{}).Where("environment = ?", environment).Where("updated_at > ?", tHours).Count(&stats.Active).Error; err != nil { - return stats, err - } - if err := n.DB.Model(&OsqueryNode{}).Where("environment = ?", environment).Where("updated_at < ?", tHours).Count(&stats.Inactive).Error; err != nil { - return stats, err - } - return stats, nil + return GetStats(n.DB, "environment", environment, hours) } -// GetStatsByPlatform to populate table stats about nodes by platform. Active machine is < 3 days +// GetStatsByPlatform to populate table stats about nodes by platform func (n *NodeManager) GetStatsByPlatform(platform string, hours int64) (StatsData, error) { - var stats StatsData - if err := n.DB.Model(&OsqueryNode{}).Where("platform = ?", platform).Count(&stats.Total).Error; err != nil { - return stats, err - } - tHours := time.Now().Add(time.Duration(hours) * time.Hour) - if err := n.DB.Model(&OsqueryNode{}).Where("platform = ?", platform).Where("updated_at > ?", tHours).Count(&stats.Active).Error; err != nil { - return stats, err - } - if err := n.DB.Model(&OsqueryNode{}).Where("platform = ?", platform).Where("updated_at < ?", tHours).Count(&stats.Inactive).Error; err != nil { - return stats, err - } - return stats, nil + return GetStats(n.DB, "platform", platform, hours) } // UpdateMetadataByUUID to update node metadata by UUID @@ -457,7 +369,9 @@ func (n *NodeManager) RefreshLastSeenBatch(nodeID []uint) error { } func (n *NodeManager) UpdateIP(nodeID uint, ip string) error { + // Update the IP address in the database return n.DB.Model(&OsqueryNode{}).Where("id = ?", nodeID).UpdateColumn("ip_address", ip).Error + } // MetadataRefresh to perform all needed update operations per node to keep metadata refreshed diff --git a/pkg/nodes/utils.go b/pkg/nodes/utils.go index c783469c..dd3f15ac 100644 --- a/pkg/nodes/utils.go +++ b/pkg/nodes/utils.go @@ -1,49 +1,73 @@ package nodes import ( - "fmt" - "math" "time" + + "gorm.io/gorm" ) -// Helper to get what is the last seen time for a node, inactive should be negative to check for past activity +// For testing - allows us to mock time.Now() +var timeNow = time.Now + +// IsActive determines if a node is active based on when it was last seen. +// The inactive parameter specifies the number of hours a node can be without +// checking in before it's considered inactive. +// Returns true if the node has checked in within the specified timeframe. func IsActive(n OsqueryNode, inactive int64) bool { - now := time.Now() - // Check config if not empty/zero - if !n.LastSeen.IsZero() { - if math.Abs(n.LastSeen.Sub(now).Hours()) < math.Abs(float64(inactive)) { - return true - } + // If LastSeen is zero (never seen), node is not active + if n.LastSeen.IsZero() { + return false } - return false -} -// Helper to generate the key to identify a full node in the cache, by UUID -func CacheFullKeyByUUID(n OsqueryNode) string { - return CacheFullKeyRaw(n.UUID, n.EnvironmentID) + // A node is active if it was seen more recently than the inactive threshold + cutoffTime := ActiveTimeCutoff(inactive) + return n.LastSeen.After(cutoffTime) } -// Helper to generate the key to identify a full node in the cache, by node_key -func CacheFullKeyByNodeKey(n OsqueryNode) string { - return CacheFullKeyRaw(n.NodeKey, n.EnvironmentID) +// ActiveTimeCutoff returns the cutoff time for active nodes +// based on the specified number of hours +func ActiveTimeCutoff(hours int64) time.Time { + return timeNow().Add(-time.Duration(hours) * time.Hour) } -// Helper to generate the key to identify a full node in the cache -func CacheFullKeyRaw(identifier string, envID uint) string { - return fmt.Sprintf("fullnode:%d:%s", envID, identifier) +// ApplyNodeTarget adds the appropriate query constraints for the target node status +// (active, inactive, all) to the provided gorm query +func ApplyNodeTarget(query *gorm.DB, target string, hours int64) *gorm.DB { + switch target { + case AllNodes: + return query + case ActiveNodes: + cutoff := ActiveTimeCutoff(hours) + return query.Where("last_seen > ?", cutoff) + case InactiveNodes: + cutoff := ActiveTimeCutoff(hours) + return query.Where("last_seen <= ?", cutoff) + default: + return query + } } -// Helper to generate the key to identify partially node in the cache, by UUID -func CachePartialKeyByUUID(n OsqueryNode) string { - return CachePartialKeyRaw(n.UUID, n.EnvironmentID) -} +// GetStats retrieves node statistics (total, active, inactive) for the given filter condition +func GetStats(db *gorm.DB, column, value string, hours int64) (StatsData, error) { + var stats StatsData -// Helper to generate the key to identify partially node in the cache, by node_key -func CachePartialKeyByNodeKey(n OsqueryNode) string { - return CachePartialKeyRaw(n.NodeKey, n.EnvironmentID) -} + // Base query with the filter condition + baseQuery := db.Model(&OsqueryNode{}).Where(column+" = ?", value) + + // Get total count + if err := baseQuery.Count(&stats.Total).Error; err != nil { + return stats, err + } + + // Get active count + cutoff := ActiveTimeCutoff(hours) + if err := baseQuery.Where("last_seen > ?", cutoff).Count(&stats.Active).Error; err != nil { + return stats, err + } + + // Get inactive count + // Calculate inactive count as total - active to be consistent + stats.Inactive = stats.Total - stats.Active -// Helper to generate the key to identify partially node in the cache -func CachePartialKeyRaw(identifier string, envID uint) string { - return fmt.Sprintf("partialnode:%d:%s", envID, identifier) + return stats, nil } diff --git a/pkg/nodes/utils_test.go b/pkg/nodes/utils_test.go index 8762face..3b212c34 100644 --- a/pkg/nodes/utils_test.go +++ b/pkg/nodes/utils_test.go @@ -5,11 +5,71 @@ import ( "time" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" ) +// setupTestDB creates an in-memory SQLite database for testing +func setupTestDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{}) + require.NoError(t, err, "Failed to open in-memory database") + + // Migrate the schema + err = db.AutoMigrate(&OsqueryNode{}) + require.NoError(t, err, "Failed to migrate schema") + + return db +} + +// testNodeParams defines optional parameters for creating mock nodes +type testNodeParams struct { + UUID string + Hostname string + Platform string + LastSeen time.Time + Environment string + Username string + OsqueryUser string + Version string +} + +// createMockNode creates a test node with the given parameters +// It accepts optional parameters to customize the node +func createMockNode(t *testing.T, db *gorm.DB, params testNodeParams) OsqueryNode { + t.Helper() + if params.UUID == "" { + params.UUID = "TEST-UUID" + } + if params.Hostname == "" { + params.Hostname = "test-host" + } + if params.Platform == "" { + params.Platform = "darwin" + } + + node := OsqueryNode{ + UUID: params.UUID, + Hostname: params.Hostname, + Platform: params.Platform, + LastSeen: params.LastSeen, + Environment: params.Environment, + Username: params.Username, + OsqueryUser: params.OsqueryUser, + PlatformVersion: params.Version, + } + + // Create the node in the database and check for errors + err := db.Create(&node).Error + require.NoError(t, err, "Failed to create mock node with UUID %s", params.UUID) + + return node +} + func TestIsActive(t *testing.T) { - // Setup current time for reference - now := time.Now() + // Use a fixed reference time for deterministic tests + refTime := time.Date(2025, 4, 10, 12, 0, 0, 0, time.UTC) // Test cases tests := []struct { @@ -21,7 +81,7 @@ func TestIsActive(t *testing.T) { { name: "Active node - recently seen", node: OsqueryNode{ - LastSeen: now.Add(-1 * time.Hour), // 1 hour ago + LastSeen: refTime.Add(-1 * time.Hour), // 1 hour ago }, inactivity: 24, // 24 hours expected: true, @@ -29,7 +89,7 @@ func TestIsActive(t *testing.T) { { name: "Inactive node - seen too long ago", node: OsqueryNode{ - LastSeen: now.Add(-48 * time.Hour), // 48 hours ago + LastSeen: refTime.Add(-48 * time.Hour), // 48 hours ago }, inactivity: 24, // 24 hours expected: false, @@ -45,56 +105,293 @@ func TestIsActive(t *testing.T) { { name: "Edge case - seen exactly at inactivity threshold", node: OsqueryNode{ - LastSeen: now.Add(time.Duration(-24) * time.Hour), // 24 hours ago + LastSeen: refTime.Add(-24 * time.Hour), // 24 hours ago }, inactivity: 24, // 24 hours - expected: false, // Expected false due to the implementation checking for "less than" not "less than or equal" + expected: false, // Should be inactive because it's not after the cutoff }, { name: "Negative inactivity parameter", node: OsqueryNode{ - LastSeen: now.Add(-12 * time.Hour), // 12 hours ago + LastSeen: refTime.Add(-12 * time.Hour), // 12 hours ago }, - inactivity: -24, // -24 hours - expected: true, + inactivity: -24, // -24 hours (moves cutoff into the future) + expected: false, // Should be inactive with negative inactivity }, } + // Mock time.Now() to return our reference time + originalTimeNow := timeNow + timeNow = func() time.Time { return refTime } + defer func() { timeNow = originalTimeNow }() // restore the original function + // Run test cases for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { result := IsActive(tt.node, tt.inactivity) - if result != tt.expected { - t.Errorf("IsActive() = %v, want %v", result, tt.expected) - } + assert.Equal(t, tt.expected, result, "IsActive() returned unexpected result") }) } } -func TestCacheFullKey(t *testing.T) { - node := OsqueryNode{ - EnvironmentID: 123, - UUID: "uuid", - NodeKey: "node_key", +func TestActiveTimeCutoff(t *testing.T) { + // Use a fixed reference time for deterministic tests + refTime := time.Date(2025, 4, 10, 12, 0, 0, 0, time.UTC) + + // Mock time.Now() to return our reference time + originalTimeNow := timeNow + timeNow = func() time.Time { return refTime } + defer func() { timeNow = originalTimeNow }() // restore the original function + + tests := []struct { + name string + hours int64 + expected time.Time + }{ + { + name: "24 hours cutoff", + hours: 24, + expected: refTime.Add(-24 * time.Hour), + }, + { + name: "Zero hours cutoff", + hours: 0, + expected: refTime, + }, + { + name: "Negative hours cutoff", + hours: -10, + expected: refTime.Add(10 * time.Hour), // Future time + }, + { + name: "Large hours cutoff", + hours: 720, // 30 days + expected: refTime.Add(-720 * time.Hour), + }, } - assert.Equal(t, CacheFullKeyByUUID(node), "fullnode:123:uuid") - assert.Equal(t, CacheFullKeyByNodeKey(node), "fullnode:123:node_key") -} -func TestCacheFullKeyRaw(t *testing.T) { - assert.Equal(t, CacheFullKeyRaw("uuid", 123), "fullnode:123:uuid") + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + cutoff := ActiveTimeCutoff(tt.hours) + assert.Equal(t, tt.expected, cutoff, "ActiveTimeCutoff() returned unexpected time") + }) + } } -func TestCachePartialKey(t *testing.T) { - node := OsqueryNode{ - EnvironmentID: 123, - UUID: "uuid", - NodeKey: "node_key", +func TestApplyNodeTarget(t *testing.T) { + db := setupTestDB(t) + + // Use a fixed reference time for deterministic tests + refTime := time.Date(2025, 4, 10, 12, 0, 0, 0, time.UTC) + + // Mock time.Now() to return our reference time + originalTimeNow := timeNow + timeNow = func() time.Time { return refTime } + defer func() { timeNow = originalTimeNow }() // restore the original function + + // Create sample nodes with different last seen times + activeNode := createMockNode(t, db, testNodeParams{ + LastSeen: refTime.Add(-1 * time.Hour), // 1 hour ago + }) + inactiveNode := createMockNode(t, db, testNodeParams{ + UUID: "INACTIVE-UUID", + LastSeen: refTime.Add(-48 * time.Hour), // 48 hours ago + }) + + // Define test cases + tests := []struct { + name string + target string + hours int64 + expectedNodeIDs []uint + }{ + { + name: "All nodes target", + target: AllNodes, + hours: 24, + expectedNodeIDs: []uint{activeNode.ID, inactiveNode.ID}, + }, + { + name: "Active nodes target", + target: ActiveNodes, + hours: 24, + expectedNodeIDs: []uint{activeNode.ID}, + }, + { + name: "Inactive nodes target", + target: InactiveNodes, + hours: 24, + expectedNodeIDs: []uint{inactiveNode.ID}, + }, + { + name: "Invalid target defaults to all nodes", + target: "invalid-target", + hours: 24, + expectedNodeIDs: []uint{activeNode.ID, inactiveNode.ID}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var nodes []OsqueryNode + + // Apply the node target filter + query := db.Model(&OsqueryNode{}) + result := ApplyNodeTarget(query, tt.target, tt.hours) + + // Execute the query + err := result.Find(&nodes).Error + require.NoError(t, err, "Failed to execute the filtered query") + + // Extract the IDs for comparison + var actualNodeIDs []uint + for _, n := range nodes { + actualNodeIDs = append(actualNodeIDs, n.ID) + } + + assert.ElementsMatch(t, tt.expectedNodeIDs, actualNodeIDs, "Filtered nodes don't match expected IDs") + }) } - assert.Equal(t, CachePartialKeyByUUID(node), "partialnode:123:uuid") - assert.Equal(t, CachePartialKeyByNodeKey(node), "partialnode:123:node_key") } -func TestCachePartialKeyRaw(t *testing.T) { - assert.Equal(t, CachePartialKeyRaw("uuid", 123), "partialnode:123:uuid") +func TestGetStats(t *testing.T) { + db := setupTestDB(t) + + // Use a fixed reference time for deterministic tests + refTime := time.Date(2025, 4, 10, 12, 0, 0, 0, time.UTC) + + // Mock time.Now() to return our reference time + originalTimeNow := timeNow + timeNow = func() time.Time { return refTime } + defer func() { timeNow = originalTimeNow }() // restore the original function + + // Clear database to ensure clean state + db.Exec("DELETE FROM osquery_nodes") + + // Create test nodes with different environments and platforms using the enhanced createMockNode + + // Platform: darwin - 2 active, 1 inactive + createMockNode(t, db, testNodeParams{ + UUID: "ACTIVE-DARWIN-1", + Platform: "darwin", + LastSeen: refTime.Add(-1 * time.Hour), // 1 hour ago (active) + }) + + createMockNode(t, db, testNodeParams{ + UUID: "ACTIVE-DARWIN-2", + Platform: "darwin", + Environment: "prod", + LastSeen: refTime.Add(-2 * time.Hour), // 2 hours ago (active) + }) + + createMockNode(t, db, testNodeParams{ + UUID: "INACTIVE-DARWIN-1", + Platform: "darwin", + Environment: "prod", + LastSeen: refTime.Add(-48 * time.Hour), // 48 hours ago (inactive) + }) + + // Platform: windows - 1 active, 2 inactive + createMockNode(t, db, testNodeParams{ + UUID: "ACTIVE-WINDOWS-1", + Platform: "windows", + Environment: "dev", + LastSeen: refTime.Add(-6 * time.Hour), // 6 hours ago (active) + }) + + createMockNode(t, db, testNodeParams{ + UUID: "INACTIVE-WINDOWS-1", + Platform: "windows", + Environment: "dev", + LastSeen: refTime.Add(-48 * time.Hour), // 48 hours ago (inactive) + }) + + createMockNode(t, db, testNodeParams{ + UUID: "INACTIVE-WINDOWS-2", + Platform: "windows", + Environment: "dev", + LastSeen: refTime.Add(-72 * time.Hour), // 72 hours ago (inactive) + Username: "test-user", + }) + + // Verify our test setup + var count int64 + db.Model(&OsqueryNode{}).Count(&count) + require.Equal(t, int64(6), count, "Should have 6 test nodes in database") + + // Test cases + tests := []struct { + name string + column string + value string + hours int64 + expectedStats StatsData + }{ + { + name: "Stats for darwin platform", + column: "platform", + value: "darwin", + hours: 24, + expectedStats: StatsData{ + Total: 3, + Active: 2, + Inactive: 1, + }, + }, + { + name: "Stats for windows platform", + column: "platform", + value: "windows", + hours: 24, + expectedStats: StatsData{ + Total: 3, + Active: 1, + Inactive: 2, + }, + }, + { + name: "Stats for prod environment", + column: "environment", + value: "prod", + hours: 24, + expectedStats: StatsData{ + Total: 2, + Active: 1, + Inactive: 1, + }, + }, + { + name: "Stats for dev environment", + column: "environment", + value: "dev", + hours: 24, + expectedStats: StatsData{ + Total: 3, + Active: 1, + Inactive: 2, + }, + }, + { + name: "Stats for non-existent value", + column: "platform", + value: "linux", + hours: 24, + expectedStats: StatsData{ + Total: 0, + Active: 0, + Inactive: 0, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + stats, err := GetStats(db, tt.column, tt.value, tt.hours) + require.NoError(t, err, "GetStats should not return an error") + + assert.Equal(t, tt.expectedStats.Total, stats.Total, "Total count mismatch") + assert.Equal(t, tt.expectedStats.Active, stats.Active, "Active count mismatch") + assert.Equal(t, tt.expectedStats.Inactive, stats.Inactive, "Inactive count mismatch") + }) + } } diff --git a/pkg/queries/queries_test.go b/pkg/queries/queries_test.go index 080bf61b..284a200f 100644 --- a/pkg/queries/queries_test.go +++ b/pkg/queries/queries_test.go @@ -22,7 +22,7 @@ func testDB(t *testing.T) *gorm.DB { q := queries.CreateQueries(db) require.NotNil(t, q, "Failed to create queries") - n := nodes.CreateNodes(db, nil) + n := nodes.CreateNodes(db) require.NotNil(t, n, "Failed to create nodes") return db diff --git a/pkg/types/osquery.go b/pkg/types/osquery.go index 8986dd55..c4ea8166 100644 --- a/pkg/types/osquery.go +++ b/pkg/types/osquery.go @@ -3,8 +3,6 @@ package types import ( "encoding/json" "strconv" - - "github.com/jmpsec/osctrl/pkg/queries" ) // Types of log types @@ -187,16 +185,16 @@ type QueryReadRequest GenericRequest // QueryReadResponse for on-demand queries from nodes type QueryReadResponse struct { - Queries queries.QueryReadQueries `json:"queries"` - NodeInvalid bool `json:"node_invalid"` + Queries map[string]string `json:"queries"` + NodeInvalid bool `json:"node_invalid"` } // AcceleratedQueryReadResponse for accelerated on-demand queries from nodes // https://github.com/osquery/osquery/blob/master/osquery/distributed/distributed.cpp#L219-L231 type AcceleratedQueryReadResponse struct { - Queries queries.QueryReadQueries `json:"queries"` - NodeInvalid bool `json:"node_invalid"` - Accelerate int `json:"accelerate"` + Queries map[string]string `json:"queries"` + NodeInvalid bool `json:"node_invalid"` + Accelerate int `json:"accelerate"` } // QueryWriteQueries to hold the on-demand queries results