Record author information with queries (#578)

This commit is contained in:
Zachary Wasserman 2016-12-07 12:22:31 -08:00 committed by GitHub
parent 6f75bf5393
commit d9190020fe
14 changed files with 228 additions and 136 deletions

View file

@ -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())

View file

@ -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)

View file

@ -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)

View file

@ -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")

View file

@ -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)
}
}

View file

@ -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
}

View file

@ -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
}

View file

@ -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{}

View file

@ -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
}

View file

@ -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:"-"`

View file

@ -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")

View file

@ -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

View file

@ -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) {

View file

@ -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,