From 87330a9753c1ba8573fd24948dce3bcab41f5afe Mon Sep 17 00:00:00 2001 From: Zach Wasserman Date: Mon, 25 Jan 2021 16:26:14 -0800 Subject: [PATCH] Add support for denylist parameter in scheduled queries (#209) Closes #202 --- server/datastore/datastore_packs_test.go | 2 ++ .../datastore_scheduled_queries_test.go | 12 +++++++++ ...119174155_AddDenylistToScheduledQueries.go | 25 +++++++++++++++++++ server/datastore/mysql/packs.go | 12 ++++----- server/datastore/mysql/scheduled_queries.go | 17 +++++++------ server/kolide/packs.go | 1 + server/kolide/scheduled_queries.go | 2 ++ server/service/endpoint_packs.go | 2 +- 8 files changed, 59 insertions(+), 14 deletions(-) create mode 100644 server/datastore/mysql/migrations/tables/20210119174155_AddDenylistToScheduledQueries.go diff --git a/server/datastore/datastore_packs_test.go b/server/datastore/datastore_packs_test.go index 2d82be5073..c1ba78a6d4 100644 --- a/server/datastore/datastore_packs_test.go +++ b/server/datastore/datastore_packs_test.go @@ -217,6 +217,7 @@ func setupPackSpecsTest(t *testing.T, ds kolide.Datastore) []*kolide.PackSpec { Name: "foo_snapshot", Interval: 600, Snapshot: boolPtr(true), + Denylist: boolPtr(false), }, kolide.PackSpecQuery{ Name: "q2", @@ -226,6 +227,7 @@ func setupPackSpecsTest(t *testing.T, ds kolide.Datastore) []*kolide.PackSpec { Shard: uintPtr(73), Platform: stringPtr("foobar"), Version: stringPtr("0.0.0.0.0.1"), + Denylist: boolPtr(true), }, }, }, diff --git a/server/datastore/datastore_scheduled_queries_test.go b/server/datastore/datastore_scheduled_queries_test.go index 1d719d81e2..4e70079f98 100644 --- a/server/datastore/datastore_scheduled_queries_test.go +++ b/server/datastore/datastore_scheduled_queries_test.go @@ -99,6 +99,18 @@ func testScheduledQuery(t *testing.T, ds kolide.Datastore) { query, err := ds.ScheduledQuery(sq1.ID) require.Nil(t, err) assert.Equal(t, uint(60), query.Interval) + assert.Nil(t, query.Denylist) + + denylist := false + query.Denylist = &denylist + + query, err = ds.SaveScheduledQuery(query) + + query, err = ds.ScheduledQuery(sq1.ID) + require.Nil(t, err) + assert.Equal(t, uint(60), query.Interval) + require.NotNil(t, query.Denylist) + assert.False(t, *query.Denylist) } func testDeleteScheduledQuery(t *testing.T, ds kolide.Datastore) { diff --git a/server/datastore/mysql/migrations/tables/20210119174155_AddDenylistToScheduledQueries.go b/server/datastore/mysql/migrations/tables/20210119174155_AddDenylistToScheduledQueries.go new file mode 100644 index 0000000000..1571017a85 --- /dev/null +++ b/server/datastore/mysql/migrations/tables/20210119174155_AddDenylistToScheduledQueries.go @@ -0,0 +1,25 @@ +package tables + +import ( + "database/sql" +) + +func init() { + MigrationClient.AddMigration(Up20210119174155, Down20210119174155) +} + +func Up20210119174155(tx *sql.Tx) error { + _, err := tx.Exec(` + ALTER TABLE scheduled_queries + ADD COLUMN denylist TINYINT(1) DEFAULT NULL + `) + return err +} + +func Down20210119174155(tx *sql.Tx) error { + _, err := tx.Exec(` + ALTER TABLE scheduled_queries + DROP COLUMN denylist + `) + return err +} diff --git a/server/datastore/mysql/packs.go b/server/datastore/mysql/packs.go index 445100b0d6..da3a794b31 100644 --- a/server/datastore/mysql/packs.go +++ b/server/datastore/mysql/packs.go @@ -4,8 +4,8 @@ import ( "database/sql" "fmt" - "github.com/jmoiron/sqlx" "github.com/fleetdm/fleet/server/kolide" + "github.com/jmoiron/sqlx" "github.com/pkg/errors" ) @@ -65,16 +65,16 @@ func applyPackSpec(tx *sqlx.Tx, spec *kolide.PackSpec) error { query = ` INSERT INTO scheduled_queries ( pack_id, query_name, name, description, ` + "`interval`" + `, - snapshot, removed, shard, platform, version + snapshot, removed, shard, platform, version, denylist ) VALUES ( ?, ?, ?, ?, ?, - ?, ?, ?, ?, ? + ?, ?, ?, ?, ?, ? ) ` _, err := tx.Exec(query, packID, q.QueryName, q.Name, q.Description, q.Interval, - q.Snapshot, q.Removed, q.Shard, q.Platform, q.Version, + q.Snapshot, q.Removed, q.Shard, q.Platform, q.Version, q.Denylist, ) switch { case isChildForeignKeyError(err): @@ -129,7 +129,7 @@ WHERE pack_id = ? AND pt.type = ? AND pt.target_id = l.id query = ` SELECT query_name, name, description, ` + "`interval`" + `, -snapshot, removed, shard, platform, version +snapshot, removed, shard, platform, version, denylist FROM scheduled_queries WHERE pack_id = ? ` @@ -179,7 +179,7 @@ WHERE pack_id = ? AND pt.type = ? AND pt.target_id = l.id query = ` SELECT query_name, name, description, ` + "`interval`" + `, -snapshot, removed, shard, platform, version +snapshot, removed, shard, platform, version, denylist FROM scheduled_queries WHERE pack_id = ? ` diff --git a/server/datastore/mysql/scheduled_queries.go b/server/datastore/mysql/scheduled_queries.go index 6e5d87903b..151c862ecd 100644 --- a/server/datastore/mysql/scheduled_queries.go +++ b/server/datastore/mysql/scheduled_queries.go @@ -21,6 +21,7 @@ func (d *Datastore) ListScheduledQueriesInPack(id uint, opts kolide.ListOptions) sq.platform, sq.version, sq.shard, + sq.denylist, q.query, q.id AS query_id FROM scheduled_queries sq @@ -53,15 +54,16 @@ func (d *Datastore) NewScheduledQuery(sq *kolide.ScheduledQuery, opts ...kolide. ` + "`interval`" + `, platform, version, - shard + shard, + denylist ) - SELECT name, ?, ?, ?, ?, ?, ?, ?, ? + SELECT name, ?, ?, ?, ?, ?, ?, ?, ?, ? FROM queries WHERE id = ? ` - result, err := db.Exec(query, sq.Name, sq.PackID, sq.Snapshot, sq.Removed, sq.Interval, sq.Platform, sq.Version, sq.Shard, sq.QueryID) + result, err := db.Exec(query, sq.Name, sq.PackID, sq.Snapshot, sq.Removed, sq.Interval, sq.Platform, sq.Version, sq.Shard, sq.Denylist, sq.QueryID) if err != nil { - return nil, errors.Wrap(err, "inserting scheduled query") + return nil, errors.Wrap(err, "insert scheduled query") } id, _ := result.LastInsertId() @@ -93,10 +95,10 @@ func (d *Datastore) NewScheduledQuery(sq *kolide.ScheduledQuery, opts ...kolide. func (d *Datastore) SaveScheduledQuery(sq *kolide.ScheduledQuery) (*kolide.ScheduledQuery, error) { query := ` UPDATE scheduled_queries - SET pack_id = ?, query_id = ?, ` + "`interval`" + ` = ?, snapshot = ?, removed = ?, platform = ?, version = ?, shard = ? + SET pack_id = ?, query_id = ?, ` + "`interval`" + ` = ?, snapshot = ?, removed = ?, platform = ?, version = ?, shard = ?, denylist = ? WHERE id = ? ` - result, err := d.db.Exec(query, sq.PackID, sq.QueryID, sq.Interval, sq.Snapshot, sq.Removed, sq.Platform, sq.Version, sq.Shard, sq.ID) + result, err := d.db.Exec(query, sq.PackID, sq.QueryID, sq.Interval, sq.Snapshot, sq.Removed, sq.Platform, sq.Version, sq.Shard, sq.Denylist, sq.ID) if err != nil { return nil, errors.Wrap(err, "saving a scheduled query") } @@ -129,6 +131,7 @@ func (d *Datastore) ScheduledQuery(id uint) (*kolide.ScheduledQuery, error) { sq.shard, sq.query_name, sq.description, + sq.denylist, q.query, q.name, q.id AS query_id @@ -139,7 +142,7 @@ func (d *Datastore) ScheduledQuery(id uint) (*kolide.ScheduledQuery, error) { ` sq := &kolide.ScheduledQuery{} if err := d.db.Get(sq, query, id); err != nil { - return nil, errors.Wrap(err, "selecting a scheduled query") + return nil, errors.Wrap(err, "select scheduled query") } return sq, nil diff --git a/server/kolide/packs.go b/server/kolide/packs.go index e9257efe7a..88475eee7d 100644 --- a/server/kolide/packs.go +++ b/server/kolide/packs.go @@ -167,6 +167,7 @@ type PackSpecQuery struct { Shard *uint `json:"shard,omitempty"` Platform *string `json:"platform,omitempty"` Version *string `json:"version,omitempty"` + Denylist *bool `json:"denylist,omitempty"` } // PackTarget associates a pack with either a host or a label diff --git a/server/kolide/scheduled_queries.go b/server/kolide/scheduled_queries.go index d51438ec42..e20dd9af0b 100644 --- a/server/kolide/scheduled_queries.go +++ b/server/kolide/scheduled_queries.go @@ -37,6 +37,7 @@ type ScheduledQuery struct { Platform *string `json:"platform,omitempty"` Version *string `json:"version,omitempty"` Shard *uint `json:"shard"` + Denylist *bool `json:"denylist"` } type ScheduledQueryPayload struct { @@ -48,4 +49,5 @@ type ScheduledQueryPayload struct { Platform *string `json:"platform"` Version *string `json:"version"` Shard *null.Int `json:"shard"` + Denylist *bool `json:"denylist"` } diff --git a/server/service/endpoint_packs.go b/server/service/endpoint_packs.go index d62e6d2823..01a60e26eb 100644 --- a/server/service/endpoint_packs.go +++ b/server/service/endpoint_packs.go @@ -3,8 +3,8 @@ package service import ( "context" - "github.com/go-kit/kit/endpoint" "github.com/fleetdm/fleet/server/kolide" + "github.com/go-kit/kit/endpoint" ) type packResponse struct {