From 1c6c7ce69af2cf19acb7cada6ee02c66cec4e478 Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Tue, 8 Apr 2025 10:13:45 +0200 Subject: [PATCH 1/9] Remove unused redis instance --- cmd/admin/main.go | 2 +- cmd/api/main.go | 2 +- cmd/cli/main.go | 2 +- cmd/tls/main.go | 2 +- pkg/nodes/cache.go | 139 ------------------------------------ pkg/nodes/nodes.go | 9 +-- pkg/nodes/utils.go | 31 -------- pkg/nodes/utils_test.go | 30 -------- pkg/queries/queries_test.go | 2 +- 9 files changed, 8 insertions(+), 211 deletions(-) delete mode 100644 pkg/nodes/cache.go 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/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/nodes.go b/pkg/nodes/nodes.go index 0c633f24..faf94041 100644 --- a/pkg/nodes/nodes.go +++ b/pkg/nodes/nodes.go @@ -5,7 +5,6 @@ import ( "strings" "time" - redis "github.com/go-redis/redis/v8" "github.com/rs/zerolog/log" "gorm.io/gorm" ) @@ -83,15 +82,13 @@ type StatsData struct { // NodeManager to handle all nodes of the system type NodeManager struct { - DB *gorm.DB - Cache *redis.Client + DB *gorm.DB } // 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 { diff --git a/pkg/nodes/utils.go b/pkg/nodes/utils.go index c783469c..fff0297b 100644 --- a/pkg/nodes/utils.go +++ b/pkg/nodes/utils.go @@ -1,7 +1,6 @@ package nodes import ( - "fmt" "math" "time" ) @@ -17,33 +16,3 @@ func IsActive(n OsqueryNode, inactive int64) bool { } 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) -} - -// 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) -} - -// 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) -} - -// 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) -} - -// 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) -} - -// 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) -} diff --git a/pkg/nodes/utils_test.go b/pkg/nodes/utils_test.go index 8762face..b4133b8f 100644 --- a/pkg/nodes/utils_test.go +++ b/pkg/nodes/utils_test.go @@ -3,8 +3,6 @@ package nodes import ( "testing" "time" - - "github.com/stretchr/testify/assert" ) func TestIsActive(t *testing.T) { @@ -70,31 +68,3 @@ func TestIsActive(t *testing.T) { }) } } - -func TestCacheFullKey(t *testing.T) { - node := OsqueryNode{ - EnvironmentID: 123, - UUID: "uuid", - NodeKey: "node_key", - } - 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") -} - -func TestCachePartialKey(t *testing.T) { - node := OsqueryNode{ - EnvironmentID: 123, - UUID: "uuid", - NodeKey: "node_key", - } - 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") -} 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 From 44bdec50e96dd2c7797c0a9cee1339be175de84f Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Tue, 8 Apr 2025 10:29:19 +0200 Subject: [PATCH 2/9] Remove unused function --- pkg/nodes/nodes.go | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/pkg/nodes/nodes.go b/pkg/nodes/nodes.go index faf94041..088e0a72 100644 --- a/pkg/nodes/nodes.go +++ b/pkg/nodes/nodes.go @@ -101,14 +101,6 @@ func CreateNodes(backend *gorm.DB) *NodeManager { 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 { @@ -125,14 +117,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 From 5a96792f81f674a21f16d19143140f4a3a058438 Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Tue, 8 Apr 2025 13:40:46 +0200 Subject: [PATCH 3/9] fix last_seen bug --- pkg/nodes/metadata.go | 16 --------- pkg/nodes/models.go | 77 +++++++++++++++++++++++++++++++++++++++++++ pkg/nodes/nodes.go | 77 ++++++------------------------------------- 3 files changed, 87 insertions(+), 83 deletions(-) delete mode 100644 pkg/nodes/metadata.go create mode 100644 pkg/nodes/models.go 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/nodes.go b/pkg/nodes/nodes.go index 088e0a72..5a77c615 100644 --- a/pkg/nodes/nodes.go +++ b/pkg/nodes/nodes.go @@ -18,61 +18,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"` @@ -185,16 +130,16 @@ func (n *NodeManager) GetBySelector(stype, selector, target string, hours int64) 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 { + if err := n.DB.Where(s+" = ?", selector).Where("last_seen > ?", 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 { + if err := n.DB.Where(s+" = ?", selector).Where("last_seen < ?", time.Now().Add(time.Duration(hours)*time.Hour)).Find(&nodes).Error; err != nil { return nodes, err } } + // TODO: This function fetches all nodes with the same selector, but we should + // consider if we need all data or just the UUID and node_key return nodes, nil } @@ -207,13 +152,11 @@ func (n *NodeManager) Gets(target string, hours int64) ([]OsqueryNode, error) { 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 { + if err := n.DB.Where("last_seen > ?", 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 { + if err := n.DB.Where("last_seen < ?", time.Now().Add(time.Duration(hours)*time.Hour)).Find(&nodes).Error; err != nil { return nodes, err } } @@ -267,10 +210,10 @@ func (n *NodeManager) GetStatsByEnv(environment string, hours int64) (StatsData, 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 { + if err := n.DB.Model(&OsqueryNode{}).Where("environment = ?", environment).Where("last_seen > ?", 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 { + if err := n.DB.Model(&OsqueryNode{}).Where("environment = ?", environment).Where("last_seen < ?", tHours).Count(&stats.Inactive).Error; err != nil { return stats, err } return stats, nil @@ -283,10 +226,10 @@ func (n *NodeManager) GetStatsByPlatform(platform string, hours int64) (StatsDat 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 { + if err := n.DB.Model(&OsqueryNode{}).Where("platform = ?", platform).Where("last_seen > ?", 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 { + if err := n.DB.Model(&OsqueryNode{}).Where("platform = ?", platform).Where("last_seen < ?", tHours).Count(&stats.Inactive).Error; err != nil { return stats, err } return stats, nil From 04f73eded6a42b4d1f2a2e3fbc7bc4d0380db597 Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Thu, 10 Apr 2025 13:08:53 +0200 Subject: [PATCH 4/9] refactor utils functions --- pkg/nodes/nodes.go | 93 +++++++++++++++++----------------------------- pkg/nodes/utils.go | 71 ++++++++++++++++++++++++++++++----- 2 files changed, 96 insertions(+), 68 deletions(-) diff --git a/pkg/nodes/nodes.go b/pkg/nodes/nodes.go index 5a77c615..3f0cb433 100644 --- a/pkg/nodes/nodes.go +++ b/pkg/nodes/nodes.go @@ -117,49 +117,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("last_seen > ?", 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("last_seen < ?", time.Now().Add(time.Duration(hours)*time.Hour)).Find(&nodes).Error; err != nil { - return nodes, err - } - } - // TODO: This function fetches all nodes with the same selector, but we should - // consider if we need all data or just the UUID and node_key + 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("last_seen > ?", time.Now().Add(time.Duration(hours)*time.Hour)).Find(&nodes).Error; err != nil { - return nodes, err - } - case InactiveNodes: - if err := n.DB.Where("last_seen < ?", 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 } @@ -203,36 +200,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("last_seen > ?", tHours).Count(&stats.Active).Error; err != nil { - return stats, err - } - if err := n.DB.Model(&OsqueryNode{}).Where("environment = ?", environment).Where("last_seen < ?", 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("last_seen > ?", tHours).Count(&stats.Active).Error; err != nil { - return stats, err - } - if err := n.DB.Model(&OsqueryNode{}).Where("platform = ?", platform).Where("last_seen < ?", 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 diff --git a/pkg/nodes/utils.go b/pkg/nodes/utils.go index fff0297b..cd6e60e9 100644 --- a/pkg/nodes/utils.go +++ b/pkg/nodes/utils.go @@ -1,18 +1,71 @@ package nodes import ( - "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 +// 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 + + // A node is active if it was seen more recently than the inactive threshold + cutoffTime := ActiveTimeCutoff(inactive) + return n.LastSeen.After(cutoffTime) +} + +// ActiveTimeCutoff returns the cutoff time for active nodes +// based on the specified number of hours +func ActiveTimeCutoff(hours int64) time.Time { + return time.Now().Add(-time.Duration(hours) * time.Hour) +} + +// 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 + } +} + +// 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 + + // 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 + if err := baseQuery.Where("last_seen < ?", cutoff).Count(&stats.Inactive).Error; err != nil { + return stats, err + } + + return stats, nil } From d6a19658ad91b27ef0f8eadf5a103c00fef4b36f Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Thu, 10 Apr 2025 14:18:09 +0200 Subject: [PATCH 5/9] add test for new functions --- pkg/nodes/utils.go | 12 +- pkg/nodes/utils_test.go | 346 ++++++++++++++++++++++++++++++++++++++-- 2 files changed, 342 insertions(+), 16 deletions(-) diff --git a/pkg/nodes/utils.go b/pkg/nodes/utils.go index cd6e60e9..dd3f15ac 100644 --- a/pkg/nodes/utils.go +++ b/pkg/nodes/utils.go @@ -6,6 +6,9 @@ import ( "gorm.io/gorm" ) +// 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. @@ -24,7 +27,7 @@ func IsActive(n OsqueryNode, inactive int64) bool { // ActiveTimeCutoff returns the cutoff time for active nodes // based on the specified number of hours func ActiveTimeCutoff(hours int64) time.Time { - return time.Now().Add(-time.Duration(hours) * time.Hour) + return timeNow().Add(-time.Duration(hours) * time.Hour) } // ApplyNodeTarget adds the appropriate query constraints for the target node status @@ -38,7 +41,7 @@ func ApplyNodeTarget(query *gorm.DB, target string, hours int64) *gorm.DB { return query.Where("last_seen > ?", cutoff) case InactiveNodes: cutoff := ActiveTimeCutoff(hours) - return query.Where("last_seen < ?", cutoff) + return query.Where("last_seen <= ?", cutoff) default: return query } @@ -63,9 +66,8 @@ func GetStats(db *gorm.DB, column, value string, hours int64) (StatsData, error) } // Get inactive count - if err := baseQuery.Where("last_seen < ?", cutoff).Count(&stats.Inactive).Error; err != nil { - return stats, err - } + // Calculate inactive count as total - active to be consistent + stats.Inactive = stats.Total - stats.Active return stats, nil } diff --git a/pkg/nodes/utils_test.go b/pkg/nodes/utils_test.go index b4133b8f..b3d75693 100644 --- a/pkg/nodes/utils_test.go +++ b/pkg/nodes/utils_test.go @@ -3,11 +3,70 @@ package nodes import ( "testing" "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 { + 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(db *gorm.DB, params testNodeParams) OsqueryNode { + 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, + } + + db.Create(&node) + 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 { @@ -19,7 +78,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, @@ -27,7 +86,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, @@ -43,28 +102,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 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), + }, + } + + 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 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(db, testNodeParams{ + LastSeen: refTime.Add(-1 * time.Hour), // 1 hour ago + }) + inactiveNode := createMockNode(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") + }) + } +} + +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(db, testNodeParams{ + UUID: "ACTIVE-DARWIN-1", + Platform: "darwin", + LastSeen: refTime.Add(-1 * time.Hour), // 1 hour ago (active) + }) + + createMockNode(db, testNodeParams{ + UUID: "ACTIVE-DARWIN-2", + Platform: "darwin", + Environment: "prod", + LastSeen: refTime.Add(-2 * time.Hour), // 2 hours ago (active) + }) + + createMockNode(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(db, testNodeParams{ + UUID: "ACTIVE-WINDOWS-1", + Platform: "windows", + Environment: "dev", + LastSeen: refTime.Add(-6 * time.Hour), // 6 hours ago (active) + }) + + createMockNode(db, testNodeParams{ + UUID: "INACTIVE-WINDOWS-1", + Platform: "windows", + Environment: "dev", + LastSeen: refTime.Add(-48 * time.Hour), // 48 hours ago (inactive) + }) + + createMockNode(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") }) } } From 556e3c1aa3be2b7bb4db9b6831eb15f3d3f66523 Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Thu, 10 Apr 2025 14:29:28 +0200 Subject: [PATCH 6/9] Add check for creation --- pkg/nodes/utils_test.go | 29 ++++++++++++++++------------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/pkg/nodes/utils_test.go b/pkg/nodes/utils_test.go index b3d75693..a4e6e122 100644 --- a/pkg/nodes/utils_test.go +++ b/pkg/nodes/utils_test.go @@ -12,6 +12,7 @@ import ( // 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") @@ -36,15 +37,14 @@ type testNodeParams struct { // createMockNode creates a test node with the given parameters // It accepts optional parameters to customize the node -func createMockNode(db *gorm.DB, params testNodeParams) OsqueryNode { +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" } @@ -60,7 +60,10 @@ func createMockNode(db *gorm.DB, params testNodeParams) OsqueryNode { PlatformVersion: params.Version, } - db.Create(&node) + // 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 } @@ -187,10 +190,10 @@ func TestApplyNodeTarget(t *testing.T) { defer func() { timeNow = originalTimeNow }() // restore the original function // Create sample nodes with different last seen times - activeNode := createMockNode(db, testNodeParams{ + activeNode := createMockNode(t, db, testNodeParams{ LastSeen: refTime.Add(-1 * time.Hour), // 1 hour ago }) - inactiveNode := createMockNode(db, testNodeParams{ + inactiveNode := createMockNode(t, db, testNodeParams{ UUID: "INACTIVE-UUID", LastSeen: refTime.Add(-48 * time.Hour), // 48 hours ago }) @@ -268,20 +271,20 @@ func TestGetStats(t *testing.T) { // Create test nodes with different environments and platforms using the enhanced createMockNode // Platform: darwin - 2 active, 1 inactive - createMockNode(db, testNodeParams{ + createMockNode(t, db, testNodeParams{ UUID: "ACTIVE-DARWIN-1", Platform: "darwin", LastSeen: refTime.Add(-1 * time.Hour), // 1 hour ago (active) }) - createMockNode(db, testNodeParams{ + createMockNode(t, db, testNodeParams{ UUID: "ACTIVE-DARWIN-2", Platform: "darwin", Environment: "prod", LastSeen: refTime.Add(-2 * time.Hour), // 2 hours ago (active) }) - createMockNode(db, testNodeParams{ + createMockNode(t, db, testNodeParams{ UUID: "INACTIVE-DARWIN-1", Platform: "darwin", Environment: "prod", @@ -289,21 +292,21 @@ func TestGetStats(t *testing.T) { }) // Platform: windows - 1 active, 2 inactive - createMockNode(db, testNodeParams{ + createMockNode(t, db, testNodeParams{ UUID: "ACTIVE-WINDOWS-1", Platform: "windows", Environment: "dev", LastSeen: refTime.Add(-6 * time.Hour), // 6 hours ago (active) }) - createMockNode(db, testNodeParams{ - UUID: "INACTIVE-WINDOWS-1", + createMockNode(t, db, testNodeParams{ + UUID: "INACTIVE-WINDOWs-1", Platform: "windows", Environment: "dev", LastSeen: refTime.Add(-48 * time.Hour), // 48 hours ago (inactive) }) - createMockNode(db, testNodeParams{ + createMockNode(t, db, testNodeParams{ UUID: "INACTIVE-WINDOWS-2", Platform: "windows", Environment: "dev", From d4f75ccb3c1a5564006bdebbe654571b4679321a Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Thu, 10 Apr 2025 15:09:31 +0200 Subject: [PATCH 7/9] Remove unused types and import dependency --- pkg/cache/cache.go | 8 -------- pkg/types/osquery.go | 12 +++++------- 2 files changed, 5 insertions(+), 15 deletions(-) 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/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 From 0d93ed3bc4d1f142af11e6e4012f4de30f34b98a Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Thu, 10 Apr 2025 15:24:35 +0200 Subject: [PATCH 8/9] Add the node cache --- pkg/nodes/node-cache.go | 80 +++++++++++++++++++++++++++++++++++++++++ pkg/nodes/nodes.go | 21 +++++++++-- 2 files changed, 98 insertions(+), 3 deletions(-) create mode 100644 pkg/nodes/node-cache.go 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 3f0cb433..0cff8958 100644 --- a/pkg/nodes/nodes.go +++ b/pkg/nodes/nodes.go @@ -1,6 +1,7 @@ package nodes import ( + "context" "fmt" "strings" "time" @@ -27,7 +28,8 @@ type StatsData struct { // NodeManager to handle all nodes of the system type NodeManager struct { - DB *gorm.DB + DB *gorm.DB + Cache *NodeCache } // CreateNodes to initialize the nodes struct and its tables @@ -43,6 +45,8 @@ func CreateNodes(backend *gorm.DB) *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 } @@ -69,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 @@ -79,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) { @@ -356,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 From 3de03d42d6d881acef99db822d22bac7a97a864f Mon Sep 17 00:00:00 2001 From: Javier Marcos <1271349+javuto@users.noreply.github.com> Date: Fri, 11 Apr 2025 16:48:51 +0200 Subject: [PATCH 9/9] Update pkg/nodes/utils_test.go Fix for inconsistent casing Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- pkg/nodes/utils_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/nodes/utils_test.go b/pkg/nodes/utils_test.go index a4e6e122..3b212c34 100644 --- a/pkg/nodes/utils_test.go +++ b/pkg/nodes/utils_test.go @@ -300,7 +300,7 @@ func TestGetStats(t *testing.T) { }) createMockNode(t, db, testNodeParams{ - UUID: "INACTIVE-WINDOWs-1", + UUID: "INACTIVE-WINDOWS-1", Platform: "windows", Environment: "dev", LastSeen: refTime.Add(-48 * time.Hour), // 48 hours ago (inactive)