From 8cc7af0148aec59982d395d837c5e798392222de Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Mon, 4 Nov 2024 15:09:18 +0100 Subject: [PATCH 01/11] add a new table for node query --- queries/queries.go | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/queries/queries.go b/queries/queries.go index f0e720db..a325cab7 100644 --- a/queries/queries.go +++ b/queries/queries.go @@ -79,6 +79,16 @@ type DistributedQuery struct { Expiration time.Time } +// NodeQuery links a node to a query +type NodeQuery struct { + NodeQueryID uint `gorm:"primaryKey;autoIncrement"` + NodeID uint `gorm:"not null;index"` + QueryID uint `gorm:"not null;index"` + Status string `gorm:"type:varchar(50);default:'pending'"` // Indexed for fast lookups + AssignedAt time.Time `gorm:"default:CURRENT_TIMESTAMP"` + CompletedAt time.Time +} + // DistributedQueryTarget to keep target logic for queries type DistributedQueryTarget struct { gorm.Model @@ -107,6 +117,7 @@ type Queries struct { func CreateQueries(backend *gorm.DB) *Queries { var q *Queries q = &Queries{DB: backend} + // table distributed_queries if err := backend.AutoMigrate(&DistributedQuery{}); err != nil { log.Fatal().Msgf("Failed to AutoMigrate table (distributed_queries): %v", err) From ad4957341bae293a328591c1cc45b68330640658 Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Tue, 5 Nov 2024 13:41:38 +0100 Subject: [PATCH 02/11] update the behaviour when creating query --- api/handlers/queries.go | 8 +++++ logging/process.go | 4 +++ queries/queries.go | 78 ++++++++++++++++++++++++++++++++++++----- 3 files changed, 82 insertions(+), 8 deletions(-) diff --git a/api/handlers/queries.go b/api/handlers/queries.go index 6a932f03..2f182de2 100644 --- a/api/handlers/queries.go +++ b/api/handlers/queries.go @@ -218,6 +218,14 @@ func (h *HandlersApi) QueriesRunHandler(w http.ResponseWriter, r *http.Request) // Remove duplicates from expected expectedClear := removeStringDuplicates(expected) + + // Create new record for query list + for _, id := range expectedClear { + if err := h.Queries.CreateNodeQuery(newQuery.ID, id); err != nil { + log.Err(err).Msgf("error creating node query for query %s and node %s", newQuery.Name, id) + } + } + // Update value for expected if err := h.Queries.SetExpected(queryName, len(expectedClear), env.ID); err != nil { apiErrorResponse(w, "error setting expected", http.StatusInternalServerError, err) diff --git a/logging/process.go b/logging/process.go index 06ffe666..03ace40d 100644 --- a/logging/process.go +++ b/logging/process.go @@ -84,6 +84,10 @@ func (l *LoggerTLS) ProcessLogQueryResult(queriesWrite types.QueryWriteRequest, if err := l.Queries.TrackExecution(q, node.UUID, queriesWrite.Statuses[q]); err != nil { log.Err(err).Msg("error adding query execution") } + // Instead of creating a new record in a separate table, we can just update the query status + if err := l.Queries.UpdateQueryStatus(q, node.ID, queriesWrite.Statuses[q]); err != nil { + log.Err(err).Msg("error updating query status") + } // Check if query is completed if err := l.Queries.VerifyComplete(q, envid); err != nil { log.Err(err).Msg("error verifying and completing query") diff --git a/queries/queries.go b/queries/queries.go index a325cab7..2de517bb 100644 --- a/queries/queries.go +++ b/queries/queries.go @@ -81,12 +81,12 @@ type DistributedQuery struct { // NodeQuery links a node to a query type NodeQuery struct { - NodeQueryID uint `gorm:"primaryKey;autoIncrement"` - NodeID uint `gorm:"not null;index"` - QueryID uint `gorm:"not null;index"` - Status string `gorm:"type:varchar(50);default:'pending'"` // Indexed for fast lookups - AssignedAt time.Time `gorm:"default:CURRENT_TIMESTAMP"` - CompletedAt time.Time + ID uint `gorm:"primaryKey;autoIncrement"` + NodeID uint `gorm:"not null;index"` + QueryID uint `gorm:"not null;index"` + Status string `gorm:"type:varchar(10);default:'pending'"` + CreatedAt time.Time `gorm:"default:CURRENT_TIMESTAMP"` + UpdatedAt time.Time } // DistributedQueryTarget to keep target logic for queries @@ -115,9 +115,13 @@ type Queries struct { // CreateQueries to initialize the queries struct func CreateQueries(backend *gorm.DB) *Queries { - var q *Queries - q = &Queries{DB: backend} + //var q *Queries + q := &Queries{DB: backend} + // table node_queries + if err := backend.AutoMigrate(&NodeQuery{}); err != nil { + log.Fatal().Msgf("Failed to AutoMigrate table (node_queries): %v", err) + } // table distributed_queries if err := backend.AutoMigrate(&DistributedQuery{}); err != nil { log.Fatal().Msgf("Failed to AutoMigrate table (distributed_queries): %v", err) @@ -137,6 +141,31 @@ func CreateQueries(backend *gorm.DB) *Queries { return q } +func (q *Queries) NodeQueries_V2(node nodes.OsqueryNode) (QueryReadQueries, bool, error) { + + var results []struct { + QueryName string + Query string + } + + q.DB.Table("distributed_queries dq"). + Select("dq.name, dq.query"). + Joins("JOIN node_queries nq ON dq.id = nq.query_id"). + Where("nq.node_id = ? AND nq.status = ?", node.ID, "pending"). + Scan(&results) + + if len(results) == 0 { + return QueryReadQueries{}, false, nil + } + + qs := make(QueryReadQueries) + for _, _q := range results { + qs[_q.QueryName] = _q.Query + } + + return qs, false, nil +} + // NodeQueries to get all queries that belong to the provided node // FIXME this will impact the performance of the TLS endpoint due to being CPU and I/O hungry // FIMXE potential mitigation can be add a cache (Redis?) layer to store queries per node_key @@ -395,6 +424,18 @@ func (q *Queries) Create(query DistributedQuery) error { return nil } +// CreateNodeQuery to link a node to a query +func (q *Queries) CreateNodeQuery(nodeID, queryID uint) error { + nodeQuery := NodeQuery{ + NodeID: nodeID, + QueryID: queryID, + } + if err := q.DB.Create(&nodeQuery).Error; err != nil { + return err + } + return nil +} + // CreateTarget to create target entry for a given query func (q *Queries) CreateTarget(name, targetType, targetValue string) error { queryTarget := DistributedQueryTarget{ @@ -460,6 +501,27 @@ func (q *Queries) SetExpected(name string, expected int, envid uint) error { return nil } +// UpdateQueryStatus to update the status of each query +func (q *Queries) UpdateQueryStatus(queryName string, nodeID uint, statusCode int) error { + + var result string + if statusCode == 0 { + result = "completed" // TODO: need be replaced with a constant + } else { + result = "error" + } + var nodeQuery NodeQuery + // For the current setup, we need a joint query to update the status, + // I am wondering if we can put an extra field in the query so that we also get the query id back from the osquery + if err := q.DB.Where("node_id = ? AND query_id = ?", nodeID, queryName).Find(&nodeQuery).Error; err != nil { + return err + } + if err := q.DB.Model(&nodeQuery).Updates(map[string]interface{}{"status": result}).Error; err != nil { + return err + } + return nil +} + // TrackExecution to keep track of where queries have already ran func (q *Queries) TrackExecution(name, uuid string, result int) error { queryExecution := DistributedQueryExecution{ From e5bddf37a4a04d666a79315df94f51e23ccd8f2c Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Wed, 13 Nov 2024 16:07:56 +0100 Subject: [PATCH 03/11] add query --- .gitignore | 3 +++ api/handlers/queries.go | 17 ++++++++++++++--- tools/bruno/collection.bru | 4 ++++ tools/bruno/nodes/get-nodes.bru | 11 +++++++++++ tools/bruno/queries/post - create queries.bru | 19 +++++++++++++++++++ 5 files changed, 51 insertions(+), 3 deletions(-) create mode 100644 tools/bruno/nodes/get-nodes.bru create mode 100644 tools/bruno/queries/post - create queries.bru diff --git a/.gitignore b/.gitignore index df914592..113d8649 100644 --- a/.gitignore +++ b/.gitignore @@ -72,3 +72,6 @@ go.work.sum deploy/docker/conf/tls/* .env !deploy/docker/conf/tls/openssl.cnf.example + +# bruno +tools/bruno/collection.bru diff --git a/api/handlers/queries.go b/api/handlers/queries.go index 2f182de2..8a632720 100644 --- a/api/handlers/queries.go +++ b/api/handlers/queries.go @@ -143,6 +143,12 @@ func (h *HandlersApi) QueriesRunHandler(w http.ResponseWriter, r *http.Request) h.Inc(metricAPIQueriesErr) return } + // Get the query id + newQuery, err = h.Queries.Get(queryName, env.ID) + if err != nil { + apiErrorResponse(w, "error creating query", http.StatusInternalServerError, err) + return + } // Temporary list of UUIDs to calculate Expected var expected []string @@ -220,9 +226,14 @@ func (h *HandlersApi) QueriesRunHandler(w http.ResponseWriter, r *http.Request) expectedClear := removeStringDuplicates(expected) // Create new record for query list - for _, id := range expectedClear { - if err := h.Queries.CreateNodeQuery(newQuery.ID, id); err != nil { - log.Err(err).Msgf("error creating node query for query %s and node %s", newQuery.Name, id) + for _, nodeUUID := range expectedClear { + node, err := h.Nodes.GetByUUID(nodeUUID) + if err != nil { + log.Err(err).Msgf("error getting node %s and failed to create node query for it", nodeUUID) + continue + } + if err := h.Queries.CreateNodeQuery(newQuery.ID, node.ID); err != nil { + log.Err(err).Msgf("error creating node query for query %s and node %s", newQuery.Name, nodeUUID) } } diff --git a/tools/bruno/collection.bru b/tools/bruno/collection.bru index e69de29b..4c9c660e 100644 --- a/tools/bruno/collection.bru +++ b/tools/bruno/collection.bru @@ -0,0 +1,4 @@ +vars:pre-request { + env: 1a026f60-edc1-4189-ab70-be99d541a473 + baseUrl: http://localhost:9002 +} diff --git a/tools/bruno/nodes/get-nodes.bru b/tools/bruno/nodes/get-nodes.bru new file mode 100644 index 00000000..fe403055 --- /dev/null +++ b/tools/bruno/nodes/get-nodes.bru @@ -0,0 +1,11 @@ +meta { + name: get-nodes + type: http + seq: 6 +} + +get { + url: {{baseUrl}} /api/v1/nodes/{{env}}/all + body: none + auth: none +} diff --git a/tools/bruno/queries/post - create queries.bru b/tools/bruno/queries/post - create queries.bru new file mode 100644 index 00000000..cfbf4989 --- /dev/null +++ b/tools/bruno/queries/post - create queries.bru @@ -0,0 +1,19 @@ +meta { + name: post - create queries + type: http + seq: 9 +} + +post { + url: {{baseUrl}}/api/v1/queries/{{env}} + body: json + auth: none +} + +body:json { + { + "environment_list": ["dev"], + "query": "SELECT * FROM system_info;", + "exp_hours": 1 + } +} From 5a759a4804ad88b54d22c002524e137693113ba5 Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Wed, 13 Nov 2024 16:47:07 +0100 Subject: [PATCH 04/11] bug fix --- .gitignore | 1 + api/handlers/queries.go | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 113d8649..b4b489d3 100644 --- a/.gitignore +++ b/.gitignore @@ -72,6 +72,7 @@ go.work.sum deploy/docker/conf/tls/* .env !deploy/docker/conf/tls/openssl.cnf.example +tls.env # bruno tools/bruno/collection.bru diff --git a/api/handlers/queries.go b/api/handlers/queries.go index 8a632720..37b7c9e4 100644 --- a/api/handlers/queries.go +++ b/api/handlers/queries.go @@ -232,7 +232,7 @@ func (h *HandlersApi) QueriesRunHandler(w http.ResponseWriter, r *http.Request) log.Err(err).Msgf("error getting node %s and failed to create node query for it", nodeUUID) continue } - if err := h.Queries.CreateNodeQuery(newQuery.ID, node.ID); err != nil { + if err := h.Queries.CreateNodeQuery(node.ID, newQuery.ID); err != nil { log.Err(err).Msgf("error creating node query for query %s and node %s", newQuery.Name, nodeUUID) } } From 2fd500c8b7f7b18c6949e3e7ea45f29088e8247a Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Thu, 14 Nov 2024 11:07:30 +0100 Subject: [PATCH 05/11] fix --- logging/process.go | 12 ++---------- queries/queries.go | 47 +++++++++++++--------------------------------- 2 files changed, 15 insertions(+), 44 deletions(-) diff --git a/logging/process.go b/logging/process.go index 03ace40d..22cfe8f7 100644 --- a/logging/process.go +++ b/logging/process.go @@ -70,16 +70,8 @@ func (l *LoggerTLS) ProcessLogQueryResult(queriesWrite types.QueryWriteRequest, Message: queriesWrite.Messages[q], } go l.DispatchQueries(d, node, debug) - // Update internal metrics per query - var err error - if queriesWrite.Statuses[q] != 0 { - err = l.Queries.IncError(q, envid) - } else { - err = l.Queries.IncExecution(q, envid) - } - if err != nil { - log.Err(err).Msg("error updating query") - } + + // TODO: This TrackExeuction need be removed // Add a record for this query if err := l.Queries.TrackExecution(q, node.UUID, queriesWrite.Statuses[q]); err != nil { log.Err(err).Msg("error adding query execution") diff --git a/queries/queries.go b/queries/queries.go index 2de517bb..15f5a140 100644 --- a/queries/queries.go +++ b/queries/queries.go @@ -1,6 +1,7 @@ package queries import ( + "fmt" "time" "github.com/jmpsec/osctrl/nodes" @@ -141,7 +142,7 @@ func CreateQueries(backend *gorm.DB) *Queries { return q } -func (q *Queries) NodeQueries_V2(node nodes.OsqueryNode) (QueryReadQueries, bool, error) { +func (q *Queries) NodeQueries(node nodes.OsqueryNode) (QueryReadQueries, bool, error) { var results []struct { QueryName string @@ -166,36 +167,6 @@ func (q *Queries) NodeQueries_V2(node nodes.OsqueryNode) (QueryReadQueries, bool return qs, false, nil } -// NodeQueries to get all queries that belong to the provided node -// FIXME this will impact the performance of the TLS endpoint due to being CPU and I/O hungry -// FIMXE potential mitigation can be add a cache (Redis?) layer to store queries per node_key -func (q *Queries) NodeQueries(node nodes.OsqueryNode) (QueryReadQueries, bool, error) { - acelerate := false - // Get all current active queries and carves - queries, err := q.GetActive(node.EnvironmentID) - if err != nil { - return QueryReadQueries{}, false, err - } - // Iterate through active queries, see if they target this node and prepare data in the same loop - qs := make(QueryReadQueries) - for _, _q := range queries { - targets, err := q.GetTargets(_q.Name) - if err != nil { - return QueryReadQueries{}, false, err - } - // FIXME disable acceleration until figure out edge cases where it would trigger by mistake - /* - if len(targets) == 1 { - acelerate = true - } - */ - if isQueryTarget(node, targets) && q.NotYetExecuted(_q.Name, node.UUID) { - qs[_q.Name] = _q.Query - } - } - return qs, acelerate, nil -} - // Gets all queries by target (active/completed/all/all-full/deleted/hidden/expired) func (q *Queries) Gets(target, qtype string, envid uint) ([]DistributedQuery, error) { var queries []DistributedQuery @@ -510,10 +481,18 @@ func (q *Queries) UpdateQueryStatus(queryName string, nodeID uint, statusCode in } else { result = "error" } + + var query DistributedQuery + // TODO: Get the query id + // I think we can put an extra field in the query so that we also get the query id back from the osquery + // This way we can avoid this query to get the query id + if err := q.DB.Where("name = ?", queryName).Find(&query).Error; err != nil { + return fmt.Errorf("error getting query id: %v", err) + } + var nodeQuery NodeQuery - // For the current setup, we need a joint query to update the status, - // I am wondering if we can put an extra field in the query so that we also get the query id back from the osquery - if err := q.DB.Where("node_id = ? AND query_id = ?", nodeID, queryName).Find(&nodeQuery).Error; err != nil { + + if err := q.DB.Where("node_id = ? AND query_id = ?", nodeID, query.ID).Find(&nodeQuery).Error; err != nil { return err } if err := q.DB.Model(&nodeQuery).Updates(map[string]interface{}{"status": result}).Error; err != nil { From e69b42f54b222b9a48f08ebd97afcda4dafe9ae8 Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Thu, 14 Nov 2024 11:26:09 +0100 Subject: [PATCH 06/11] update for osctrl-admin --- admin/handlers/post.go | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/admin/handlers/post.go b/admin/handlers/post.go index 3e070f8a..0c153a7e 100644 --- a/admin/handlers/post.go +++ b/admin/handlers/post.go @@ -223,6 +223,18 @@ func (h *HandlersAdmin) QueryRunPOSTHandler(w http.ResponseWriter, r *http.Reque } // Remove duplicates from expected expectedClear := removeStringDuplicates(expected) + + // Create new record for query list + for _, nodeUUID := range expectedClear { + node, err := h.Nodes.GetByUUID(nodeUUID) + if err != nil { + log.Err(err).Msgf("error getting node %s and failed to create node query for it", nodeUUID) + continue + } + if err := h.Queries.CreateNodeQuery(node.ID, newQuery.ID); err != nil { + log.Err(err).Msgf("error creating node query for query %s and node %s", newQuery.Name, nodeUUID) + } + } // Update value for expected if err := h.Queries.SetExpected(newQuery.Name, len(expectedClear), env.ID); err != nil { adminErrorResponse(w, "error setting expected", http.StatusInternalServerError, err) From 5b7e55908df3e953a839d3b7a436da23fa5c0b18 Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Thu, 14 Nov 2024 13:00:38 +0100 Subject: [PATCH 07/11] Add test for node query --- queries/go.mod | 6 ++ queries/go.sum | 20 ++---- queries/queries.go | 6 +- queries/queries_test.go | 133 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 148 insertions(+), 17 deletions(-) create mode 100644 queries/queries_test.go diff --git a/queries/go.mod b/queries/go.mod index b62e1b53..09e38f52 100644 --- a/queries/go.mod +++ b/queries/go.mod @@ -9,13 +9,19 @@ replace github.com/jmpsec/osctrl/utils => ../utils require ( github.com/jmpsec/osctrl/nodes v0.0.0-20241107152746-1f093f5e8faf github.com/jmpsec/osctrl/utils v0.0.0-20241107150205-621ec8aafdae + github.com/stretchr/testify v1.9.0 + gorm.io/driver/sqlite v1.5.6 gorm.io/gorm v1.25.12 ) require ( + github.com/davecgh/go-spew v1.1.1 // indirect github.com/mattn/go-colorable v0.1.13 // indirect github.com/mattn/go-isatty v0.0.20 // indirect + github.com/mattn/go-sqlite3 v1.14.22 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect golang.org/x/sys v0.26.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect ) require ( diff --git a/queries/go.sum b/queries/go.sum index be037f90..bdfbe367 100644 --- a/queries/go.sum +++ b/queries/go.sum @@ -11,10 +11,11 @@ github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/ github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA= github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg= github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM= -github.com/mattn/go-isatty v0.0.19 h1:JITubQf0MOLdlGRuRq+jtsDlekdYPia9ZFsB8h/APPA= github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/mattn/go-sqlite3 v1.14.22 h1:2gZY6PC6kBnID23Tichd1K+Z0oS6nE/XwU+Vz/5o4kU= +github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= @@ -27,25 +28,16 @@ github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsT github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.12.0 h1:CM0HF96J0hcLAwsHPJZjfdNzs0gftsLfgKt57wWHJ0o= golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.25.0 h1:r+8e+loiHxRqhXVl6ML1nO3l1+oFoWbnlu2Ehimmi34= -golang.org/x/sys v0.25.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/sys v0.26.0 h1:KHjCJyddX0LoSTb3J+vWpupP9p0oznkqVk/IfjymZbo= golang.org/x/sys v0.26.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= -golang.org/x/text v0.16.0 h1:a94ExnEXNtEwYLGJSIUxnWoxoRz/ZcCsV63ROupILh4= -golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI= -golang.org/x/text v0.17.0 h1:XtiM5bkSOt+ewxlOE/aE/AKEHibwj/6gvWMl9Rsh0Qc= -golang.org/x/text v0.17.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY= -golang.org/x/text v0.18.0 h1:XvMDiNzPAl0jr17s6W9lcaIhGUfUORdGCNsuLmPG224= -golang.org/x/text v0.18.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY= golang.org/x/text v0.19.0 h1:kTxAhCbGbxhK0IwgSKiMO5awPoDQ0RpfiVYBfK860YM= golang.org/x/text v0.19.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= -gorm.io/gorm v1.25.10 h1:dQpO+33KalOA+aFYGlK+EfxcI5MbO7EP2yYygwh9h+s= -gorm.io/gorm v1.25.10/go.mod h1:hbnx/Oo0ChWMn1BIhpy1oYozzpM15i4YPuHDmfYtwg8= -gorm.io/gorm v1.25.11 h1:/Wfyg1B/je1hnDx3sMkX+gAlxrlZpn6X0BXRlwXlvHg= -gorm.io/gorm v1.25.11/go.mod h1:xh7N7RHfYlNc5EmcI/El95gXusucDrQnHXe0+CgWcLQ= +gorm.io/driver/sqlite v1.5.6 h1:fO/X46qn5NUEEOZtnjJRWRzZMe8nqJiQ9E+0hi+hKQE= +gorm.io/driver/sqlite v1.5.6/go.mod h1:U+J8craQU6Fzkcvu8oLeAQmi50TkwPEhHDEjQZXDah4= gorm.io/gorm v1.25.12 h1:I0u8i2hWQItBq1WfE0o2+WuL9+8L21K9e2HHSTE/0f8= gorm.io/gorm v1.25.12/go.mod h1:xh7N7RHfYlNc5EmcI/El95gXusucDrQnHXe0+CgWcLQ= diff --git a/queries/queries.go b/queries/queries.go index 15f5a140..3d4a814a 100644 --- a/queries/queries.go +++ b/queries/queries.go @@ -145,8 +145,8 @@ func CreateQueries(backend *gorm.DB) *Queries { func (q *Queries) NodeQueries(node nodes.OsqueryNode) (QueryReadQueries, bool, error) { var results []struct { - QueryName string - Query string + Name string + Query string } q.DB.Table("distributed_queries dq"). @@ -161,7 +161,7 @@ func (q *Queries) NodeQueries(node nodes.OsqueryNode) (QueryReadQueries, bool, e qs := make(QueryReadQueries) for _, _q := range results { - qs[_q.QueryName] = _q.Query + qs[_q.Name] = _q.Query } return qs, false, nil diff --git a/queries/queries_test.go b/queries/queries_test.go new file mode 100644 index 00000000..845e26e3 --- /dev/null +++ b/queries/queries_test.go @@ -0,0 +1,133 @@ +package queries_test + +import ( + "fmt" + "log" + "testing" + "time" + + "github.com/jmpsec/osctrl/nodes" + "github.com/jmpsec/osctrl/queries" + "github.com/stretchr/testify/assert" + "gorm.io/driver/sqlite" + "gorm.io/gorm" +) + +func setupTestDB() (*gorm.DB, error) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + if err != nil { + return nil, err + } + return db, nil +} + +func TestNodeQueries(t *testing.T) { + db, err := setupTestDB() + if err != nil { + t.Fatalf("Failed to setup test database: %v", err) + } + + // Create tables + q := queries.CreateQueries(db) + nodes.CreateNodes(db) + + // Create test data + node := nodes.OsqueryNode{ + Model: gorm.Model{ID: 1}, + } + distributedQuery := queries.DistributedQuery{ + Model: gorm.Model{ID: 1}, + Name: "test_query", + Query: "SELECT * FROM osquery_info;", + EnvironmentID: 1, + Expiration: time.Now().Add(24 * time.Hour), + } + nodeQuery := queries.NodeQuery{ + NodeID: 1, + QueryID: 1, + Status: "pending", + } + + // Query sqlite_master to list all tables + var tables []string + err = db.Raw("SELECT name FROM sqlite_master WHERE type='table'").Scan(&tables).Error + if err != nil { + log.Fatalf("failed to list tables: %v", err) + } + + fmt.Println("Tables in the database:") + for _, table := range tables { + fmt.Println(table) + } + + if err := db.Create(&node).Error; err != nil { + t.Fatalf("Failed to create test node: %v", err) + } + if err := db.Create(&distributedQuery).Error; err != nil { + t.Fatalf("Failed to create test distributed query: %v", err) + } + if err := db.Create(&nodeQuery).Error; err != nil { + t.Fatalf("Failed to create test node query: %v", err) + } + + // Test NodeQueries function + queries, _, err := q.NodeQueries(node) + if err != nil { + t.Fatalf("NodeQueries returned an error: %v", err) + } + // Print queries + fmt.Println(queries) + + assert.NotEmpty(t, queries, "Expected non-empty queries") + assert.Equal(t, "SELECT * FROM osquery_info;", queries["test_query"], "Query does not match expected value") +} +func TestUpdateQueryStatus(t *testing.T) { + db, err := setupTestDB() + if err != nil { + t.Fatalf("Failed to setup test database: %v", err) + } + + // Create tables + q := queries.CreateQueries(db) + nodes.CreateNodes(db) + + // Create test data + node := nodes.OsqueryNode{ + Model: gorm.Model{ID: 1}, + } + distributedQuery := queries.DistributedQuery{ + Model: gorm.Model{ID: 1}, + Name: "test_query", + Query: "SELECT * FROM osquery_info;", + EnvironmentID: 1, + Expiration: time.Now().Add(24 * time.Hour), + } + nodeQuery := queries.NodeQuery{ + NodeID: 1, + QueryID: 1, + Status: "pending", + } + + if err := db.Create(&node).Error; err != nil { + t.Fatalf("Failed to create test node: %v", err) + } + if err := db.Create(&distributedQuery).Error; err != nil { + t.Fatalf("Failed to create test distributed query: %v", err) + } + if err := db.Create(&nodeQuery).Error; err != nil { + t.Fatalf("Failed to create test node query: %v", err) + } + + // Test UpdateQueryStatus function + err = q.UpdateQueryStatus("test_query", 1, 0) + if err != nil { + t.Fatalf("UpdateQueryStatus returned an error: %v", err) + } + + var updatedNodeQuery queries.NodeQuery + if err := db.Where("node_id = ? AND query_id = ?", 1, 1).Find(&updatedNodeQuery).Error; err != nil { + t.Fatalf("Failed to find updated node query: %v", err) + } + + assert.Equal(t, "completed", updatedNodeQuery.Status, "Status does not match expected value") +} From 5e366551d25d4baf3724a2f89828a69718da8f44 Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Thu, 14 Nov 2024 13:27:02 +0100 Subject: [PATCH 08/11] Get the query id for osctrl-admin --- admin/handlers/post.go | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/admin/handlers/post.go b/admin/handlers/post.go index 0c153a7e..44bb151d 100644 --- a/admin/handlers/post.go +++ b/admin/handlers/post.go @@ -150,6 +150,12 @@ func (h *HandlersAdmin) QueryRunPOSTHandler(w http.ResponseWriter, r *http.Reque h.Inc(metricAdminErr) return } + // Get the query id + newQuery, err = h.Queries.Get(newQuery.Name, env.ID) + if err != nil { + adminErrorResponse(w, "error creating query", http.StatusInternalServerError, err) + return + } // Temporary list of UUIDs to calculate Expected var expected []string // Create environment target From 92f6f81a177bf37e5c86ec871c24180be951ac2f Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Thu, 14 Nov 2024 13:33:31 +0100 Subject: [PATCH 09/11] Add test for CreateNodeQuery --- queries/queries_test.go | 43 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/queries/queries_test.go b/queries/queries_test.go index 845e26e3..a8389c97 100644 --- a/queries/queries_test.go +++ b/queries/queries_test.go @@ -131,3 +131,46 @@ func TestUpdateQueryStatus(t *testing.T) { assert.Equal(t, "completed", updatedNodeQuery.Status, "Status does not match expected value") } +func TestCreateNodeQuery(t *testing.T) { + db, err := setupTestDB() + if err != nil { + t.Fatalf("Failed to setup test database: %v", err) + } + + // Create tables + q := queries.CreateQueries(db) + nodes.CreateNodes(db) + + // Create test data + node := nodes.OsqueryNode{ + Model: gorm.Model{ID: 1}, + } + distributedQuery := queries.DistributedQuery{ + Model: gorm.Model{ID: 1}, + Name: "test_query", + Query: "SELECT * FROM osquery_info;", + EnvironmentID: 1, + Expiration: time.Now().Add(24 * time.Hour), + } + + if err := db.Create(&node).Error; err != nil { + t.Fatalf("Failed to create test node: %v", err) + } + if err := db.Create(&distributedQuery).Error; err != nil { + t.Fatalf("Failed to create test distributed query: %v", err) + } + + // Test CreateNodeQuery function + err = q.CreateNodeQuery(1, 1) + if err != nil { + t.Fatalf("CreateNodeQuery returned an error: %v", err) + } + + var nodeQuery queries.NodeQuery + if err := db.Where("node_id = ? AND query_id = ?", 1, 1).Find(&nodeQuery).Error; err != nil { + t.Fatalf("Failed to find created node query: %v", err) + } + + assert.Equal(t, uint(1), nodeQuery.NodeID, "NodeID does not match expected value") + assert.Equal(t, uint(1), nodeQuery.QueryID, "QueryID does not match expected value") +} From 81909d93f937b74ffc85da82f0f65baf51381a61 Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Thu, 14 Nov 2024 13:45:35 +0100 Subject: [PATCH 10/11] add real-time process stats back --- logging/process.go | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/logging/process.go b/logging/process.go index 22cfe8f7..82eebd3e 100644 --- a/logging/process.go +++ b/logging/process.go @@ -70,7 +70,17 @@ func (l *LoggerTLS) ProcessLogQueryResult(queriesWrite types.QueryWriteRequest, Message: queriesWrite.Messages[q], } go l.DispatchQueries(d, node, debug) - + // TODO: need be refactored + // Update internal metrics per query + var err error + if queriesWrite.Statuses[q] != 0 { + err = l.Queries.IncError(q, envid) + } else { + err = l.Queries.IncExecution(q, envid) + } + if err != nil { + log.Err(err).Msg("error updating query") + } // TODO: This TrackExeuction need be removed // Add a record for this query if err := l.Queries.TrackExecution(q, node.UUID, queriesWrite.Statuses[q]); err != nil { From fb2ec8ccc19f2419e56d6e158b93ec02cc2438d5 Mon Sep 17 00:00:00 2001 From: Zhuoyuan Liu Date: Tue, 3 Dec 2024 11:07:21 +0100 Subject: [PATCH 11/11] Use const for status --- queries/queries.go | 22 +++++++++++++--------- queries/queries_test.go | 6 +++--- 2 files changed, 16 insertions(+), 12 deletions(-) diff --git a/queries/queries.go b/queries/queries.go index 3d4a814a..3d613962 100644 --- a/queries/queries.go +++ b/queries/queries.go @@ -58,6 +58,12 @@ const ( TargetHidden string = "hidden" ) +const ( + DistributedQueryStatusPending string = "pending" + DistributedQueryStatusCompleted string = "completed" + DistributedQueryStatusError string = "error" +) + // DistributedQuery as abstraction of a distributed query type DistributedQuery struct { gorm.Model @@ -82,12 +88,10 @@ type DistributedQuery struct { // NodeQuery links a node to a query type NodeQuery struct { - ID uint `gorm:"primaryKey;autoIncrement"` - NodeID uint `gorm:"not null;index"` - QueryID uint `gorm:"not null;index"` - Status string `gorm:"type:varchar(10);default:'pending'"` - CreatedAt time.Time `gorm:"default:CURRENT_TIMESTAMP"` - UpdatedAt time.Time + gorm.Model + NodeID uint `gorm:"not null;index"` + QueryID uint `gorm:"not null;index"` + Status string `gorm:"type:varchar(8);default:'pending'"` } // DistributedQueryTarget to keep target logic for queries @@ -152,7 +156,7 @@ func (q *Queries) NodeQueries(node nodes.OsqueryNode) (QueryReadQueries, bool, e q.DB.Table("distributed_queries dq"). Select("dq.name, dq.query"). Joins("JOIN node_queries nq ON dq.id = nq.query_id"). - Where("nq.node_id = ? AND nq.status = ?", node.ID, "pending"). + Where("nq.node_id = ? AND nq.status = ?", node.ID, DistributedQueryStatusPending). Scan(&results) if len(results) == 0 { @@ -477,9 +481,9 @@ func (q *Queries) UpdateQueryStatus(queryName string, nodeID uint, statusCode in var result string if statusCode == 0 { - result = "completed" // TODO: need be replaced with a constant + result = DistributedQueryStatusCompleted } else { - result = "error" + result = DistributedQueryStatusError } var query DistributedQuery diff --git a/queries/queries_test.go b/queries/queries_test.go index a8389c97..cb1e642a 100644 --- a/queries/queries_test.go +++ b/queries/queries_test.go @@ -45,7 +45,7 @@ func TestNodeQueries(t *testing.T) { nodeQuery := queries.NodeQuery{ NodeID: 1, QueryID: 1, - Status: "pending", + Status: queries.DistributedQueryStatusPending, } // Query sqlite_master to list all tables @@ -105,7 +105,7 @@ func TestUpdateQueryStatus(t *testing.T) { nodeQuery := queries.NodeQuery{ NodeID: 1, QueryID: 1, - Status: "pending", + Status: queries.DistributedQueryStatusPending, } if err := db.Create(&node).Error; err != nil { @@ -129,7 +129,7 @@ func TestUpdateQueryStatus(t *testing.T) { t.Fatalf("Failed to find updated node query: %v", err) } - assert.Equal(t, "completed", updatedNodeQuery.Status, "Status does not match expected value") + assert.Equal(t, queries.DistributedQueryStatusCompleted, updatedNodeQuery.Status, "Status does not match expected value") } func TestCreateNodeQuery(t *testing.T) { db, err := setupTestDB()