diff --git a/server/datastore/datastore_campaigns_test.go b/server/datastore/datastore_campaigns_test.go index f7fc0fa309..689a5dd297 100644 --- a/server/datastore/datastore_campaigns_test.go +++ b/server/datastore/datastore_campaigns_test.go @@ -25,9 +25,11 @@ func checkTargets(t *testing.T, ds kolide.Datastore, campaignID uint, expectedHo } func testDistributedQueryCampaign(t *testing.T, ds kolide.Datastore) { + user := newUser(t, ds, "Zach", "zwass", "zwass@kolide.co", true) + mockClock := clock.NewMockClock() - query := newQuery(t, ds, "test", "select * from time") + query := newQuery(t, ds, "test", "select * from time", user.ID) campaign := newCampaign(t, ds, query.ID, kolide.QueryRunning, mockClock.Now()) @@ -64,9 +66,11 @@ func testDistributedQueryCampaign(t *testing.T, ds kolide.Datastore) { } func testCleanupDistributedQueryCampaigns(t *testing.T, ds kolide.Datastore) { + user := newUser(t, ds, "Zach", "zwass", "zwass@kolide.co", true) + mockClock := clock.NewMockClock() - query := newQuery(t, ds, "test", "select * from time") + query := newQuery(t, ds, "test", "select * from time", user.ID) c1 := newCampaign(t, ds, query.ID, kolide.QueryWaiting, mockClock.Now()) c2 := newCampaign(t, ds, query.ID, kolide.QueryRunning, mockClock.Now()) diff --git a/server/datastore/datastore_hosts_test.go b/server/datastore/datastore_hosts_test.go index 50f3a3e845..c32871806a 100644 --- a/server/datastore/datastore_hosts_test.go +++ b/server/datastore/datastore_hosts_test.go @@ -37,8 +37,8 @@ var enrollTests = []struct { }, } -func testSaveHosts(t *testing.T, db kolide.Datastore) { - host, err := db.NewHost(&kolide.Host{ +func testSaveHosts(t *testing.T, ds kolide.Datastore) { + host, err := ds.NewHost(&kolide.Host{ DetailUpdateTime: time.Now(), NodeKey: "1", UUID: "1", @@ -48,10 +48,10 @@ func testSaveHosts(t *testing.T, db kolide.Datastore) { require.NotNil(t, host) host.HostName = "bar.local" - err = db.SaveHost(host) + err = ds.SaveHost(host) require.Nil(t, err) - host, err = db.Host(host.ID) + host, err = ds.Host(host.ID) require.Nil(t, err) assert.Equal(t, "bar.local", host.HostName) @@ -68,18 +68,18 @@ func testSaveHosts(t *testing.T, db kolide.Datastore) { }, } - err = db.SaveHost(host) + err = ds.SaveHost(host) require.Nil(t, err) - host, err = db.Host(host.ID) + host, err = ds.Host(host.ID) require.Nil(t, err) require.NotNil(t, host) require.Equal(t, 2, len(host.NetworkInterfaces)) primaryNicID := host.NetworkInterfaces[0].ID host.PrimaryNetworkInterfaceID = &primaryNicID - err = db.SaveHost(host) + err = ds.SaveHost(host) require.Nil(t, err) - host, err = db.Host(host.ID) + host, err = ds.Host(host.ID) require.Nil(t, err) require.NotNil(t, host) require.Equal(t, 2, len(host.NetworkInterfaces)) @@ -89,9 +89,9 @@ func testSaveHosts(t *testing.T, db kolide.Datastore) { host.NetworkInterfaces = []*kolide.NetworkInterface{ host.NetworkInterfaces[1], } - err = db.SaveHost(host) + err = ds.SaveHost(host) require.Nil(t, err) - host, err = db.Host(host.ID) + host, err = ds.Host(host.ID) require.Nil(t, err) require.NotNil(t, host) assert.Equal(t, host.NetworkInterfaces[0].ID, *host.PrimaryNetworkInterfaceID) @@ -99,24 +99,24 @@ func testSaveHosts(t *testing.T, db kolide.Datastore) { // remove all nics primary nic should be nil host.NetworkInterfaces = []*kolide.NetworkInterface{} - err = db.SaveHost(host) + err = ds.SaveHost(host) require.Nil(t, err) assert.Nil(t, host.PrimaryNetworkInterfaceID) - host, err = db.Host(host.ID) + host, err = ds.Host(host.ID) require.Nil(t, err) require.NotNil(t, host) assert.Nil(t, host.PrimaryNetworkInterfaceID) - err = db.DeleteHost(host) + err = ds.DeleteHost(host) assert.Nil(t, err) - host, err = db.Host(host.ID) + host, err = ds.Host(host.ID) assert.NotNil(t, err) assert.Nil(t, host) } -func testDeleteHost(t *testing.T, db kolide.Datastore) { - host, err := db.NewHost(&kolide.Host{ +func testDeleteHost(t *testing.T, ds kolide.Datastore) { + host, err := ds.NewHost(&kolide.Host{ DetailUpdateTime: time.Now(), NodeKey: "1", UUID: "1", @@ -125,17 +125,17 @@ func testDeleteHost(t *testing.T, db kolide.Datastore) { require.Nil(t, err) require.NotNil(t, host) - err = db.DeleteHost(host) + err = ds.DeleteHost(host) assert.Nil(t, err) - host, err = db.Host(host.ID) + host, err = ds.Host(host.ID) assert.NotNil(t, err) } -func testListHost(t *testing.T, db kolide.Datastore) { +func testListHost(t *testing.T, ds kolide.Datastore) { hosts := []*kolide.Host{} for i := 0; i < 10; i++ { - host, err := db.NewHost(&kolide.Host{ + host, err := ds.NewHost(&kolide.Host{ DetailUpdateTime: time.Now(), OsqueryHostID: strconv.Itoa(i), NodeKey: fmt.Sprintf("%d", i), @@ -160,7 +160,7 @@ func testListHost(t *testing.T, db kolide.Datastore) { }, } - err := db.SaveHost(hosts[1]) + err := ds.SaveHost(hosts[1]) require.Nil(t, err) hosts[3].NetworkInterfaces = []*kolide.NetworkInterface{ @@ -169,10 +169,10 @@ func testListHost(t *testing.T, db kolide.Datastore) { IPAddress: "99.100.101.104", }, } - err = db.SaveHost(hosts[3]) + err = ds.SaveHost(hosts[3]) require.Nil(t, err) - hosts2, err := db.ListHosts(kolide.ListOptions{}) + hosts2, err := ds.ListHosts(kolide.ListOptions{}) require.Nil(t, err) assert.Equal(t, len(hosts), len(hosts2)) @@ -182,13 +182,13 @@ func testListHost(t *testing.T, db kolide.Datastore) { assert.Equal(t, "en1", hosts2[1].NetworkInterfaces[1].Interface) assert.Equal(t, "en2", hosts2[3].NetworkInterfaces[0].Interface) - err = db.DeleteHost(hosts[0]) + err = ds.DeleteHost(hosts[0]) require.Nil(t, err) - hosts2, err = db.ListHosts(kolide.ListOptions{}) + hosts2, err = ds.ListHosts(kolide.ListOptions{}) require.Nil(t, err) assert.Equal(t, len(hosts)-1, len(hosts2)) - hosts, err = db.ListHosts(kolide.ListOptions{}) + hosts, err = ds.ListHosts(kolide.ListOptions{}) require.Nil(t, err) require.Equal(t, len(hosts2), len(hosts)) hosts[0].NetworkInterfaces = []*kolide.NetworkInterface{ @@ -202,9 +202,9 @@ func testListHost(t *testing.T, db kolide.Datastore) { }, } - err = db.SaveHost(hosts[0]) + err = ds.SaveHost(hosts[0]) require.Nil(t, err) - hosts2, err = db.ListHosts(kolide.ListOptions{}) + hosts2, err = ds.ListHosts(kolide.ListOptions{}) require.Nil(t, err) require.Equal(t, hosts[0].ID, hosts2[0].ID) assert.Equal(t, len(hosts[0].NetworkInterfaces), len(hosts2[0].NetworkInterfaces)) @@ -212,10 +212,10 @@ func testListHost(t *testing.T, db kolide.Datastore) { assert.Equal(t, hosts[0].ID, hosts2[0].NetworkInterfaces[0].HostID) } -func testEnrollHost(t *testing.T, db kolide.Datastore) { +func testEnrollHost(t *testing.T, ds kolide.Datastore) { var hosts []*kolide.Host for _, tt := range enrollTests { - h, err := db.EnrollHost(tt.uuid, tt.nodeKeySize) + h, err := ds.EnrollHost(tt.uuid, tt.nodeKeySize) require.Nil(t, err) hosts = append(hosts, h) @@ -225,25 +225,25 @@ func testEnrollHost(t *testing.T, db kolide.Datastore) { } -func testAuthenticateHost(t *testing.T, db kolide.Datastore) { +func testAuthenticateHost(t *testing.T, ds kolide.Datastore) { for _, tt := range enrollTests { - h, err := db.EnrollHost(tt.uuid, tt.nodeKeySize) + h, err := ds.EnrollHost(tt.uuid, tt.nodeKeySize) require.Nil(t, err) - returned, err := db.AuthenticateHost(h.NodeKey) + returned, err := ds.AuthenticateHost(h.NodeKey) require.Nil(t, err) assert.Equal(t, h.NodeKey, returned.NodeKey) } - _, err := db.AuthenticateHost("7B1A9DC9-B042-489F-8D5A-EEC2412C95AA") + _, err := ds.AuthenticateHost("7B1A9DC9-B042-489F-8D5A-EEC2412C95AA") assert.NotNil(t, err) - _, err = db.AuthenticateHost("") + _, err = ds.AuthenticateHost("") assert.NotNil(t, err) } -func testSearchHosts(t *testing.T, db kolide.Datastore) { - _, err := db.NewHost(&kolide.Host{ +func testSearchHosts(t *testing.T, ds kolide.Datastore) { + _, err := ds.NewHost(&kolide.Host{ OsqueryHostID: "1234", DetailUpdateTime: time.Now(), NodeKey: "1", @@ -252,7 +252,7 @@ func testSearchHosts(t *testing.T, db kolide.Datastore) { }) require.Nil(t, err) - h2, err := db.NewHost(&kolide.Host{ + h2, err := ds.NewHost(&kolide.Host{ OsqueryHostID: "5679", DetailUpdateTime: time.Now(), NodeKey: "2", @@ -261,7 +261,7 @@ func testSearchHosts(t *testing.T, db kolide.Datastore) { }) require.Nil(t, err) - h3, err := db.NewHost(&kolide.Host{ + h3, err := ds.NewHost(&kolide.Host{ OsqueryHostID: "99999", DetailUpdateTime: time.Now(), NodeKey: "3", @@ -272,24 +272,24 @@ func testSearchHosts(t *testing.T, db kolide.Datastore) { // We once threw errors when the search query was empty. Verify that we // don't error. - _, err = db.SearchHosts("") + _, err = ds.SearchHosts("") require.Nil(t, err) - hosts, err := db.SearchHosts("foo") + hosts, err := ds.SearchHosts("foo") assert.Nil(t, err) assert.Len(t, hosts, 2) - host, err := db.SearchHosts("foo", h3.ID) + host, err := ds.SearchHosts("foo", h3.ID) require.Nil(t, err) require.Len(t, host, 1) assert.Equal(t, "foo.local", host[0].HostName) - host, err = db.SearchHosts("foo", h3.ID, h2.ID) + host, err = ds.SearchHosts("foo", h3.ID, h2.ID) require.Nil(t, err) require.Len(t, host, 1) assert.Equal(t, "foo.local", host[0].HostName) - none, err := db.SearchHosts("xxx") + none, err := ds.SearchHosts("xxx") assert.Nil(t, err) assert.Len(t, none, 0) @@ -304,15 +304,15 @@ func testSearchHosts(t *testing.T, db kolide.Datastore) { IPAddress: "99.100.101.103", }, } - err = db.SaveHost(h2) + err = ds.SaveHost(h2) require.Nil(t, err) - hits, err := db.SearchHosts("99.100.101") + hits, err := ds.SearchHosts("99.100.101") require.Nil(t, err) require.Equal(t, 1, len(hits)) assert.Equal(t, 2, len(hits[0].NetworkInterfaces)) - hits, err = db.SearchHosts("99.100.111") + hits, err = ds.SearchHosts("99.100.111") require.Nil(t, err) assert.Equal(t, 0, len(hits)) @@ -322,23 +322,23 @@ func testSearchHosts(t *testing.T, db kolide.Datastore) { IPAddress: "99.100.101.104", }, } - err = db.SaveHost(h3) + err = ds.SaveHost(h3) require.Nil(t, err) - hits, err = db.SearchHosts("99.100.101") + hits, err = ds.SearchHosts("99.100.101") require.Nil(t, err) assert.Equal(t, 2, len(hits)) assert.Equal(t, 2, len(hits[0].NetworkInterfaces)) assert.Equal(t, "en0", hits[0].NetworkInterfaces[0].Interface) - hits, err = db.SearchHosts("99.100.101", h3.ID) + hits, err = ds.SearchHosts("99.100.101", h3.ID) require.Nil(t, err) assert.Equal(t, 1, len(hits)) } -func testSearchHostsLimit(t *testing.T, db kolide.Datastore) { +func testSearchHostsLimit(t *testing.T, ds kolide.Datastore) { for i := 0; i < 15; i++ { - _, err := db.NewHost(&kolide.Host{ + _, err := ds.NewHost(&kolide.Host{ DetailUpdateTime: time.Now(), OsqueryHostID: fmt.Sprintf("host%d", i), NodeKey: fmt.Sprintf("%d", i), @@ -348,13 +348,15 @@ func testSearchHostsLimit(t *testing.T, db kolide.Datastore) { require.Nil(t, err) } - hosts, err := db.SearchHosts("foo") + hosts, err := ds.SearchHosts("foo") require.Nil(t, err) assert.Len(t, hosts, 10) } -func testDistributedQueriesForHost(t *testing.T, db kolide.Datastore) { - h1, err := db.NewHost(&kolide.Host{ +func testDistributedQueriesForHost(t *testing.T, ds kolide.Datastore) { + user := newUser(t, ds, "Zach", "zwass", "zwass@kolide.co", true) + + h1, err := ds.NewHost(&kolide.Host{ OsqueryHostID: "1", DetailUpdateTime: time.Now(), NodeKey: "1", @@ -363,7 +365,7 @@ func testDistributedQueriesForHost(t *testing.T, db kolide.Datastore) { }) require.Nil(t, err) - h2, err := db.NewHost(&kolide.Host{ + h2, err := ds.NewHost(&kolide.Host{ OsqueryHostID: "2", DetailUpdateTime: time.Now(), NodeKey: "2", @@ -374,15 +376,15 @@ func testDistributedQueriesForHost(t *testing.T, db kolide.Datastore) { // All should have no queries var queries map[uint]string - queries, err = db.DistributedQueriesForHost(h1) + queries, err = ds.DistributedQueriesForHost(h1) require.Nil(t, err) assert.Empty(t, queries) - queries, err = db.DistributedQueriesForHost(h2) + queries, err = ds.DistributedQueriesForHost(h2) require.Nil(t, err) assert.Empty(t, queries) // Create a label - l1, err := db.NewLabel(&kolide.Label{ + l1, err := ds.NewLabel(&kolide.Label{ Name: "label foo", Query: "query1", }) @@ -391,16 +393,17 @@ func testDistributedQueriesForHost(t *testing.T, db kolide.Datastore) { // Add hosts to label for _, h := range []*kolide.Host{h1, h2} { - err = db.RecordLabelQueryExecutions(h, map[string]bool{l1ID: true}, time.Now()) + err = ds.RecordLabelQueryExecutions(h, map[string]bool{l1ID: true}, time.Now()) require.Nil(t, err) } // Create a query q1 := &kolide.Query{ - Name: "bar", - Query: "select * from bar", + Name: "bar", + Query: "select * from bar", + AuthorID: user.ID, } - q1, err = db.NewQuery(q1) + q1, err = ds.NewQuery(q1) require.Nil(t, err) // Create a query campaign @@ -408,7 +411,7 @@ func testDistributedQueriesForHost(t *testing.T, db kolide.Datastore) { QueryID: q1.ID, Status: kolide.QueryRunning, } - c1, err = db.NewDistributedQueryCampaign(c1) + c1, err = ds.NewDistributedQueryCampaign(c1) require.Nil(t, err) // Add a target to the campaign @@ -417,15 +420,15 @@ func testDistributedQueriesForHost(t *testing.T, db kolide.Datastore) { DistributedQueryCampaignID: c1.ID, TargetID: l1.ID, } - target, err = db.NewDistributedQueryCampaignTarget(target) + target, err = ds.NewDistributedQueryCampaignTarget(target) require.Nil(t, err) // All should have the query now - queries, err = db.DistributedQueriesForHost(h1) + queries, err = ds.DistributedQueriesForHost(h1) require.Nil(t, err) assert.Len(t, queries, 1) assert.Equal(t, "select * from bar", queries[c1.ID]) - queries, err = db.DistributedQueriesForHost(h2) + queries, err = ds.DistributedQueriesForHost(h2) require.Nil(t, err) assert.Len(t, queries, 1) assert.Equal(t, "select * from bar", queries[c1.ID]) @@ -436,22 +439,23 @@ func testDistributedQueriesForHost(t *testing.T, db kolide.Datastore) { DistributedQueryCampaignID: c1.ID, Status: kolide.ExecutionSucceeded, } - _, err = db.NewDistributedQueryExecution(exec) + _, err = ds.NewDistributedQueryExecution(exec) require.Nil(t, err) // Add another query/campaign q2 := &kolide.Query{ - Name: "foo", - Query: "select * from foo", + Name: "foo", + Query: "select * from foo", + AuthorID: user.ID, } - q2, err = db.NewQuery(q2) + q2, err = ds.NewQuery(q2) require.Nil(t, err) c2 := &kolide.DistributedQueryCampaign{ QueryID: q2.ID, Status: kolide.QueryRunning, } - c2, err = db.NewDistributedQueryCampaign(c2) + c2, err = ds.NewDistributedQueryCampaign(c2) require.Nil(t, err) // This one targets only h1 @@ -460,30 +464,30 @@ func testDistributedQueriesForHost(t *testing.T, db kolide.Datastore) { DistributedQueryCampaignID: c2.ID, TargetID: h1.ID, } - _, err = db.NewDistributedQueryCampaignTarget(target) + _, err = ds.NewDistributedQueryCampaignTarget(target) require.Nil(t, err) // Check for correct queries - queries, err = db.DistributedQueriesForHost(h1) + queries, err = ds.DistributedQueriesForHost(h1) require.Nil(t, err) assert.Len(t, queries, 1) assert.Equal(t, "select * from foo", queries[c2.ID]) - queries, err = db.DistributedQueriesForHost(h2) + queries, err = ds.DistributedQueriesForHost(h2) require.Nil(t, err) assert.Len(t, queries, 1) assert.Equal(t, "select * from bar", queries[c1.ID]) // End both of the campaigns c1.Status = kolide.QueryComplete - require.Nil(t, db.SaveDistributedQueryCampaign(c1)) + require.Nil(t, ds.SaveDistributedQueryCampaign(c1)) c2.Status = kolide.QueryComplete - require.Nil(t, db.SaveDistributedQueryCampaign(c2)) + require.Nil(t, ds.SaveDistributedQueryCampaign(c2)) // Now no queries should be returned - queries, err = db.DistributedQueriesForHost(h1) + queries, err = ds.DistributedQueriesForHost(h1) require.Nil(t, err) assert.Empty(t, queries) - queries, err = db.DistributedQueriesForHost(h2) + queries, err = ds.DistributedQueriesForHost(h2) require.Nil(t, err) assert.Empty(t, queries) diff --git a/server/datastore/datastore_packs_test.go b/server/datastore/datastore_packs_test.go index 61aac42c8d..55e8d4fbf2 100644 --- a/server/datastore/datastore_packs_test.go +++ b/server/datastore/datastore_packs_test.go @@ -30,6 +30,8 @@ func testDeletePack(t *testing.T, ds kolide.Datastore) { } func testAddAndRemoveQueryFromPack(t *testing.T, ds kolide.Datastore) { + user := newUser(t, ds, "Zach", "zwass", "zwass@kolide.co", false) + pack := &kolide.Pack{ Name: "foo", } @@ -38,42 +40,46 @@ func testAddAndRemoveQueryFromPack(t *testing.T, ds kolide.Datastore) { assert.NotEqual(t, uint(0), pack.ID) q1 := &kolide.Query{ - Name: "bar", - Query: "bar", + Name: "bar", + Query: "bar", + AuthorID: user.ID, } q1, err = ds.NewQuery(q1) - assert.Nil(t, err) + require.Nil(t, err) assert.NotEqual(t, uint(0), q1.ID) err = ds.AddQueryToPack(q1.ID, pack.ID) - assert.Nil(t, err) + require.Nil(t, err) q2 := &kolide.Query{ - Name: "baz", - Query: "baz", + Name: "baz", + Query: "baz", + AuthorID: user.ID, } q2, err = ds.NewQuery(q2) - assert.Nil(t, err) + require.Nil(t, err) assert.NotEqual(t, uint(0), q2.ID) assert.NotEqual(t, q1.ID, q2.ID) err = ds.AddQueryToPack(q2.ID, pack.ID) - assert.Nil(t, err) + require.Nil(t, err) queries, err := ds.ListQueriesInPack(pack) - assert.Nil(t, err) + require.Nil(t, err) assert.Len(t, queries, 2) err = ds.RemoveQueryFromPack(q1, pack) - assert.Nil(t, err) + require.Nil(t, err) queries, err = ds.ListQueriesInPack(pack) - assert.Nil(t, err) + require.Nil(t, err) assert.Len(t, queries, 1) } func testGetHostsInPack(t *testing.T, ds kolide.Datastore) { + user := newUser(t, ds, "Zach", "zwass", "zwass@kolide.co", true) + mockClock := clock.NewMockClock() p1, err := ds.NewPack(&kolide.Pack{ @@ -82,14 +88,16 @@ func testGetHostsInPack(t *testing.T, ds kolide.Datastore) { require.Nil(t, err) q1, err := ds.NewQuery(&kolide.Query{ - Name: "foo", - Query: "foo", + Name: "foo", + Query: "foo", + AuthorID: user.ID, }) require.Nil(t, err) q2, err := ds.NewQuery(&kolide.Query{ - Name: "bar", - Query: "bar", + Name: "bar", + Query: "bar", + AuthorID: user.ID, }) require.Nil(t, err) diff --git a/server/datastore/datastore_queries_test.go b/server/datastore/datastore_queries_test.go index 177711d33e..500b564ff4 100644 --- a/server/datastore/datastore_queries_test.go +++ b/server/datastore/datastore_queries_test.go @@ -11,17 +11,20 @@ import ( ) func testDeleteQuery(t *testing.T, ds kolide.Datastore) { + user := newUser(t, ds, "Zach", "zwass", "zwass@kolide.co", true) + query := &kolide.Query{ Name: "foo", Query: "bar", Interval: 123, + AuthorID: user.ID, } query, err := ds.NewQuery(query) - assert.Nil(t, err) + require.Nil(t, err) assert.NotEqual(t, query.ID, 0) err = ds.DeleteQuery(query) - assert.Nil(t, err) + require.Nil(t, err) assert.NotEqual(t, query.ID, 0) _, err = ds.Query(query.ID) @@ -29,39 +32,47 @@ func testDeleteQuery(t *testing.T, ds kolide.Datastore) { } func testSaveQuery(t *testing.T, ds kolide.Datastore) { + user := newUser(t, ds, "Zach", "zwass", "zwass@kolide.co", true) + query := &kolide.Query{ - Name: "foo", - Query: "bar", + Name: "foo", + Query: "bar", + AuthorID: user.ID, } query, err := ds.NewQuery(query) - assert.Nil(t, err) + require.Nil(t, err) assert.NotEqual(t, 0, query.ID) query.Query = "baz" err = ds.SaveQuery(query) - assert.Nil(t, err) + require.Nil(t, err) queryVerify, err := ds.Query(query.ID) - assert.Nil(t, err) + require.Nil(t, err) assert.Equal(t, "baz", queryVerify.Query) + assert.Equal(t, "Zach", queryVerify.AuthorName) } func testListQuery(t *testing.T, ds kolide.Datastore) { + user := newUser(t, ds, "Zach", "zwass", "zwass@kolide.co", true) + for i := 0; i < 10; i++ { _, err := ds.NewQuery(&kolide.Query{ - Name: fmt.Sprintf("name%02d", i), - Query: fmt.Sprintf("query%02d", i), - Saved: true, + Name: fmt.Sprintf("name%02d", i), + Query: fmt.Sprintf("query%02d", i), + Saved: true, + AuthorID: user.ID, }) require.Nil(t, err) } // One unsaved query should not be returned _, err := ds.NewQuery(&kolide.Query{ - Name: "unsaved", - Query: "select * from time", - Saved: false, + Name: "unsaved", + Query: "select * from time", + Saved: false, + AuthorID: user.ID, }) require.Nil(t, err) @@ -78,8 +89,10 @@ func checkPacks(t *testing.T, expected []kolide.Pack, actual []kolide.Pack) { } func testLoadPacksForQueries(t *testing.T, ds kolide.Datastore) { - q1 := newQuery(t, ds, "q1", "select * from time") - q2 := newQuery(t, ds, "q2", "select * from osquery_info") + user := newUser(t, ds, "Zach", "zwass", "zwass@kolide.co", true) + + q1 := newQuery(t, ds, "q1", "select * from time", user.ID) + q2 := newQuery(t, ds, "q2", "select * from osquery_info", user.ID) p1 := newPack(t, ds, "p1") p2 := newPack(t, ds, "p2") diff --git a/server/datastore/inmem/queries.go b/server/datastore/inmem/queries.go index b9caf0b8bd..bf946529d2 100644 --- a/server/datastore/inmem/queries.go +++ b/server/datastore/inmem/queries.go @@ -49,6 +49,13 @@ func (orm *Datastore) DeleteQuery(query *kolide.Query) error { return nil } +func (orm *Datastore) getUserNameByID(id uint) string { + if u, ok := orm.users[id]; ok { + return u.Name + } + return "" +} + func (orm *Datastore) Query(id uint) (*kolide.Query, error) { orm.mtx.Lock() defer orm.mtx.Unlock() @@ -58,6 +65,8 @@ func (orm *Datastore) Query(id uint) (*kolide.Query, error) { return nil, errors.ErrNotFound } + query.AuthorName = orm.getUserNameByID(query.AuthorID) + if err := orm.loadPacksForQueries([]*kolide.Query{query}); err != nil { return nil, errors.DatabaseError(err) } @@ -78,8 +87,10 @@ func (orm *Datastore) ListQueries(opt kolide.ListOptions) ([]*kolide.Query, erro queries := []*kolide.Query{} for _, k := range keys { - if orm.queries[uint(k)].Saved { - queries = append(queries, orm.queries[uint(k)]) + q := orm.queries[uint(k)] + if q.Saved { + q.AuthorName = orm.getUserNameByID(q.AuthorID) + queries = append(queries, q) } } diff --git a/server/datastore/mysql/migrations/20161118212758_CreateTableUsers.go b/server/datastore/mysql/migrations/20161118212649_CreateTableUsers.go similarity index 89% rename from server/datastore/mysql/migrations/20161118212758_CreateTableUsers.go rename to server/datastore/mysql/migrations/20161118212649_CreateTableUsers.go index c0b46bb6ff..f381947b9d 100644 --- a/server/datastore/mysql/migrations/20161118212758_CreateTableUsers.go +++ b/server/datastore/mysql/migrations/20161118212649_CreateTableUsers.go @@ -7,10 +7,10 @@ import ( ) func init() { - goose.AddMigration(Up_20161118212758, Down_20161118212758) + goose.AddMigration(Up_20161118212649, Down_20161118212649) } -func Up_20161118212758(tx *sql.Tx) error { +func Up_20161118212649(tx *sql.Tx) error { _, err := tx.Exec( "CREATE TABLE `users` (" + "`id` int(10) unsigned NOT NULL AUTO_INCREMENT," + @@ -36,7 +36,7 @@ func Up_20161118212758(tx *sql.Tx) error { return err } -func Down_20161118212758(tx *sql.Tx) error { +func Down_20161118212649(tx *sql.Tx) error { _, err := tx.Exec("DROP TABLE IF EXISTS `users`;") return err } diff --git a/server/datastore/mysql/migrations/20161118212649_CreateTableQueries.go b/server/datastore/mysql/migrations/20161118212758_CreateTableQueries.go similarity index 76% rename from server/datastore/mysql/migrations/20161118212649_CreateTableQueries.go rename to server/datastore/mysql/migrations/20161118212758_CreateTableQueries.go index 9f857f8ad9..486ccbefcf 100644 --- a/server/datastore/mysql/migrations/20161118212649_CreateTableQueries.go +++ b/server/datastore/mysql/migrations/20161118212758_CreateTableQueries.go @@ -7,10 +7,10 @@ import ( ) func init() { - goose.AddMigration(Up_20161118212649, Down_20161118212649) + goose.AddMigration(Up_20161118212758, Down_20161118212758) } -func Up_20161118212649(tx *sql.Tx) error { +func Up_20161118212758(tx *sql.Tx) error { _, err := tx.Exec( "CREATE TABLE `queries` (" + "`id` int(10) unsigned NOT NULL AUTO_INCREMENT," + @@ -27,13 +27,15 @@ func Up_20161118212649(tx *sql.Tx) error { "`differential` tinyint(1) NOT NULL DEFAULT FALSE," + "`platform` varchar(255) DEFAULT NULL," + "`version` varchar(255) DEFAULT NULL," + - "PRIMARY KEY (`id`)" + + "`author_id` int(10) unsigned DEFAULT NULL," + + "PRIMARY KEY (`id`)," + + "FOREIGN KEY (`author_id`) REFERENCES `users`(`id`) ON DELETE SET NULL" + ") ENGINE=InnoDB DEFAULT CHARSET=utf8;", ) return err } -func Down_20161118212649(tx *sql.Tx) error { +func Down_20161118212758(tx *sql.Tx) error { _, err := tx.Exec("DROP TABLE IF EXISTS `queries`;") return err } diff --git a/server/datastore/mysql/queries.go b/server/datastore/mysql/queries.go index 8b8124e9f5..c4a324bbb2 100644 --- a/server/datastore/mysql/queries.go +++ b/server/datastore/mysql/queries.go @@ -11,13 +11,14 @@ func (d *Datastore) NewQuery(query *kolide.Query) (*kolide.Query, error) { sql := ` INSERT INTO queries (name, description, query, - saved, snapshot, differential, platform, version, ` + "`interval`" + `) - VALUES ( ?, ?, ?, ?, ?, ?, ?, ?, ? ) + saved, snapshot, differential, platform, version, + ` + "`interval`" + `, author_id) + VALUES ( ?, ?, ?, ?, ?, ?, ?, ?, ?, ? ) ` result, err := d.db.Exec(sql, query.Name, query.Description, query.Query, - query.Saved, query.Snapshot, query.Differential, query.Platform, - query.Version, query.Interval) + query.Saved, query.Snapshot, query.Differential, query.Platform, query.Version, + query.Interval, query.AuthorID) if err != nil { return nil, errors.DatabaseError(err) } @@ -32,11 +33,14 @@ func (d *Datastore) SaveQuery(q *kolide.Query) error { sql := ` UPDATE queries SET name = ?, description = ?, query = ?, ` + "`interval`" + ` = ?, - saved = ?, snapshot = ?, differential = ?, platform = ?, version = ? + saved = ?, snapshot = ?, differential = ?, platform = ?, version = ?, + author_id = ? WHERE id = ? AND NOT deleted ` _, err := d.db.Exec(sql, q.Name, q.Description, q.Query, q.Interval, - q.Saved, q.Snapshot, q.Differential, q.Platform, q.Version, q.ID) + q.Saved, q.Snapshot, q.Differential, q.Platform, q.Version, + q.AuthorID, + q.ID) if err != nil { return errors.DatabaseError(err) } @@ -64,7 +68,12 @@ func (d *Datastore) DeleteQuery(query *kolide.Query) error { // exists func (d *Datastore) Query(id uint) (*kolide.Query, error) { sql := ` - SELECT * FROM queries WHERE id = ? AND NOT deleted + SELECT q.*, IFNULL(u.name, "") AS author_name + FROM queries q + LEFT JOIN users u + ON q.author_id = u.id + WHERE q.id = ? + AND NOT q.deleted ` query := &kolide.Query{} if err := d.db.Get(query, sql, id); err != nil { @@ -82,9 +91,12 @@ func (d *Datastore) Query(id uint) (*kolide.Query, error) { // determined by passed in kolide.ListOptions func (d *Datastore) ListQueries(opt kolide.ListOptions) ([]*kolide.Query, error) { sql := ` - SELECT * FROM queries + SELECT q.*, IFNULL(u.name, "") AS author_name + FROM queries q + LEFT JOIN users u + ON q.author_id = u.id WHERE saved = true - AND NOT deleted + AND NOT q.deleted ` sql = appendListOptionsToSQL(sql, opt) results := []*kolide.Query{} diff --git a/server/datastore/test_util.go b/server/datastore/test_util.go index 3b044e76a7..0ad54feb68 100644 --- a/server/datastore/test_util.go +++ b/server/datastore/test_util.go @@ -8,10 +8,11 @@ import ( "github.com/stretchr/testify/require" ) -func newQuery(t *testing.T, ds kolide.Datastore, name, q string) *kolide.Query { +func newQuery(t *testing.T, ds kolide.Datastore, name, q string, authorID uint) *kolide.Query { query, err := ds.NewQuery(&kolide.Query{ - Name: name, - Query: q, + Name: name, + Query: q, + AuthorID: authorID, }) require.Nil(t, err) @@ -114,3 +115,19 @@ func newLabel(t *testing.T, ds kolide.Datastore, name, query string) *kolide.Lab return l } + +func newUser(t *testing.T, ds kolide.Datastore, name, username, email string, admin bool) *kolide.User { + u, err := ds.NewUser(&kolide.User{ + Password: []byte("garbage"), + Salt: "garbage", + Name: name, + Username: username, + Email: email, + Admin: admin, + }) + + require.Nil(t, err) + require.NotZero(t, u.ID) + + return u +} diff --git a/server/kolide/queries.go b/server/kolide/queries.go index bf563f0d32..60ac22becb 100644 --- a/server/kolide/queries.go +++ b/server/kolide/queries.go @@ -57,6 +57,10 @@ type Query struct { Differential bool `json:"differential"` Platform string `json:"platform"` Version string `json:"version"` + AuthorID uint `json:"author_id" db:"author_id"` + // AuthorName is retrieved with a join to the users table in the MySQL + // backend (using AuthorID) + AuthorName string `json:"author_name" db:"author_name"` // Packs is loaded when retrieving queries, but is stored in a join // table in the MySQL backend. Packs []Pack `json:"packs" db:"-"` diff --git a/server/service/service_campaigns.go b/server/service/service_campaigns.go index f6f8a9eb3a..9ecf559e94 100644 --- a/server/service/service_campaigns.go +++ b/server/service/service_campaigns.go @@ -18,9 +18,10 @@ func (svc service) NewDistributedQueryCampaign(ctx context.Context, queryString } query, err := svc.ds.NewQuery(&kolide.Query{ - Name: fmt.Sprintf("distributed_%s_%d", vc.Username(), time.Now().Unix()), - Query: queryString, - Saved: false, + Name: fmt.Sprintf("distributed_%s_%d", vc.Username(), time.Now().Unix()), + Query: queryString, + Saved: false, + AuthorID: vc.UserID(), }) if err != nil { return nil, errors.Wrap(err, "new query") diff --git a/server/service/service_queries.go b/server/service/service_queries.go index 496bcde3ab..da266c224c 100644 --- a/server/service/service_queries.go +++ b/server/service/service_queries.go @@ -1,6 +1,7 @@ package service import ( + "github.com/kolide/kolide-ose/server/contexts/viewer" "github.com/kolide/kolide-ose/server/kolide" "golang.org/x/net/context" ) @@ -48,6 +49,11 @@ func (svc service) NewQuery(ctx context.Context, p kolide.QueryPayload) (*kolide query.Version = *p.Version } + vc, ok := viewer.FromContext(ctx) + if ok { + query.AuthorID = vc.UserID() + } + query, err := svc.ds.NewQuery(query) if err != nil { return nil, err diff --git a/server/service/service_queries_test.go b/server/service/service_queries_test.go index cda5168136..0b2c74c06d 100644 --- a/server/service/service_queries_test.go +++ b/server/service/service_queries_test.go @@ -4,9 +4,11 @@ import ( "testing" "github.com/kolide/kolide-ose/server/config" + "github.com/kolide/kolide-ose/server/contexts/viewer" "github.com/kolide/kolide-ose/server/datastore/inmem" "github.com/kolide/kolide-ose/server/kolide" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "golang.org/x/net/context" ) @@ -63,10 +65,15 @@ func TestNewQuery(t *testing.T) { ds, err := inmem.New(config.TestConfig()) assert.Nil(t, err) + createTestUsers(t, ds) svc, err := newTestService(ds, nil) assert.Nil(t, err) + user, err := ds.User("admin1") + require.Nil(t, err) + ctx := context.Background() + ctx = viewer.NewContext(ctx, viewer.Viewer{User: user}) name := "foo" query := "select * from time;" @@ -79,7 +86,9 @@ func TestNewQuery(t *testing.T) { queries, err := ds.ListQueries(kolide.ListOptions{}) assert.Nil(t, err) - assert.Len(t, queries, 1) + if assert.Len(t, queries, 1) { + assert.Equal(t, "Test Name admin1", queries[0].AuthorName) + } } func TestModifyQuery(t *testing.T) { diff --git a/server/service/util_test.go b/server/service/util_test.go index 916d788b14..e0ed047f7d 100644 --- a/server/service/util_test.go +++ b/server/service/util_test.go @@ -22,6 +22,7 @@ func createTestUsers(t *testing.T, ds kolide.Datastore) map[string]kolide.User { users := make(map[string]kolide.User) for _, u := range testUsers { user := &kolide.User{ + Name: "Test Name " + u.Username, Username: u.Username, Email: u.Email, Admin: u.IsAdmin,