empayre/fleet
14
0
Fork 0
mirror of https://github.com/empayre/fleet.git synced 2024-11-06 08:55:24 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add MVP support for software inventory to osquery-perf for load testing (#2751)

Browse Source 
* Add MVP support for software inventory to osquery-perf for load testing

* Fix test compile
This commit is contained in:
Lucas Manuel Rodriguez 2021-11-01 15:23:31 -03:00 committed by GitHub
parent a8735d55bb
commit 8642bb785e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 396 additions and 331 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
cmd/osquery-perf/README.md
View File

@ -11,26 +11,12 @@ The only requirement for running this tool is a working installation of
## Usage
Typically `go run` is used:
Typically `go run` is used.
You can use `--help` to view the available configuration:
```
go run agent.go --help
Usage of agent.go:
-config_interval duration
Interval for config requests (default 1m0s)
-enroll_secret string
Enroll secret to authenticate enrollment
-host_count int
Number of hosts to start (default 10) (default 10)
-query_interval duration
Interval for live query requests (default 10s)
-seed int
Seed for random generator (default current time) (default 1586310930917739000)
-server_url string
URL (with protocol and port of osquery server) (default "https://localhost:8080")
-start_period duration
Duration to spread start of hosts over (default 10s)
exit status 2
```
The tool should be invoked with the appropriate enroll secret. A typical

126
cmd/osquery-perf/agent.go
View File

@ -4,7 +4,6 @@ import (
"bytes"
"crypto/tls"
"embed"
_ "embed"
"encoding/json"
"flag"
"fmt"
@ -17,6 +16,8 @@ import (
"text/template"
"time"
"github.com/fleetdm/fleet/v4/server/fleet"
"github.com/fleetdm/fleet/v4/server/service"
"github.com/google/uuid"
"github.com/valyala/fasthttp"
)
@ -61,14 +62,14 @@ func (s *Stats) runLoop() {
}
}
type NodeKeyManager struct {
type nodeKeyManager struct {
filepath string
l sync.Mutex
nodekeys []string
}
func (n *NodeKeyManager) LoadKeys() {
func (n *nodeKeyManager) LoadKeys() {
if n.filepath == "" {
return
}
@ -85,7 +86,7 @@ func (n *NodeKeyManager) LoadKeys() {
fmt.Printf("loaded %d node keys\n", len(n.nodekeys))
}
func (n *NodeKeyManager) Get(i int) string {
func (n *nodeKeyManager) Get(i int) string {
n.l.Lock()
defer n.l.Unlock()
@ -95,7 +96,7 @@ func (n *NodeKeyManager) Get(i int) string {
return ""
}
func (n *NodeKeyManager) Add(nodekey string) {
func (n *nodeKeyManager) Add(nodekey string) {
if n.filepath == "" {
return
}
@ -115,26 +116,27 @@ func (n *NodeKeyManager) Add(nodekey string) {
}
}
type Agent struct {
type agent struct {
ServerAddress string
EnrollSecret string
NodeKey string
UUID string
FastClient fasthttp.Client
Client http.Client
ConfigInterval time.Duration
QueryInterval time.Duration
Templates *template.Template
strings map[string]string
Stats *Stats
NodeKeyManager *NodeKeyManager
SoftwareCount int
NodeKeyManager *nodeKeyManager
strings map[string]string
}
func NewAgent(serverAddress, enrollSecret string, templates *template.Template, configInterval, queryInterval time.Duration) *Agent {
func newAgent(serverAddress, enrollSecret string, templates *template.Template, configInterval, queryInterval time.Duration, softwareCount int) *agent {
transport := http.DefaultTransport.(*http.Transport).Clone()
transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
transport.DisableCompression = true
return &Agent{
return &agent{
ServerAddress: serverAddress,
EnrollSecret: enrollSecret,
Templates: templates,
@ -144,7 +146,7 @@ func NewAgent(serverAddress, enrollSecret string, templates *template.Template,
FastClient: fasthttp.Client{
TLSConfig: &tls.Config{InsecureSkipVerify: true},
},
Client: http.Client{Transport: transport},
SoftwareCount: softwareCount,
strings: make(map[string]string),
}
}
@ -157,12 +159,12 @@ type distributedReadResponse struct {
Queries map[string]string `json:"queries"`
}
func (a *Agent) runLoop(i int, onlyAlreadyEnrolled bool) {
if err := a.Enroll(i, onlyAlreadyEnrolled); err != nil {
func (a *agent) runLoop(i int, onlyAlreadyEnrolled bool) {
if err := a.enroll(i, onlyAlreadyEnrolled); err != nil {
return
}
a.Config()
a.config()
resp, err := a.DistributedRead()
if err != nil {
log.Println(err)
@ -177,7 +179,7 @@ func (a *Agent) runLoop(i int, onlyAlreadyEnrolled bool) {
for {
select {
case <-configTicker:
a.Config()
a.config()
case <-liveQueryTicker:
resp, err := a.DistributedRead()
if err != nil {
@ -191,17 +193,17 @@ func (a *Agent) runLoop(i int, onlyAlreadyEnrolled bool) {
}
}
func (a *Agent) waitingDo(req *fasthttp.Request, res *fasthttp.Response) {
func (a *agent) waitingDo(req *fasthttp.Request, res *fasthttp.Response) {
err := a.FastClient.Do(req, res)
for err != nil || res.StatusCode() != http.StatusOK {
fmt.Println(err, res.StatusCode())
a.Stats.RecordStats(1, 0, 0)
<-time.Tick(time.Duration(rand.Intn(120)+1) * time.Second)
err = fasthttp.Do(req, res)
err = a.FastClient.Do(req, res)
}
}
func (a *Agent) Enroll(i int, onlyAlreadyEnrolled bool) error {
func (a *agent) enroll(i int, onlyAlreadyEnrolled bool) error {
a.NodeKey = a.NodeKeyManager.Get(i)
if a.NodeKey != "" {
a.Stats.RecordStats(0, 1, 0)
@ -250,7 +252,7 @@ func (a *Agent) Enroll(i int, onlyAlreadyEnrolled bool) error {
return nil
}
func (a *Agent) Config() {
func (a *agent) config() {
body := bytes.NewBufferString(`{"node_key": "` + a.NodeKey + `"}`)
req := fasthttp.AcquireRequest()
@ -276,7 +278,7 @@ func (a *Agent) Config() {
const stringVals = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_."
func (a *Agent) randomString(n int) string {
func (a *agent) randomString(n int) string {
sb := strings.Builder{}
sb.Grow(n)
for i := 0; i < n; i++ {
@ -285,7 +287,7 @@ func (a *Agent) randomString(n int) string {
return sb.String()
}
func (a *Agent) CachedString(key string) string {
func (a *agent) CachedString(key string) string {
if val, ok := a.strings[key]; ok {
return val
}
@ -294,7 +296,20 @@ func (a *Agent) CachedString(key string) string {
return val
}
func (a *Agent) DistributedRead() (*distributedReadResponse, error) {
func (a *agent) SoftwareMacOS() []fleet.Software {
software := make([]fleet.Software, a.SoftwareCount)
for i := 0; i < len(software); i++ {
software[i] = fleet.Software{
Name: "Placeholder_Software",
Version: "0.0.1",
BundleIdentifier: "com.fleetdm.osquery-perf",
Source: "osquery-perf",
}
}
return software
}
func (a *agent) DistributedRead() (*distributedReadResponse, error) {
req := fasthttp.AcquireRequest()
req.SetBody([]byte(`{"node_key": "` + a.NodeKey + `"}`))
req.Header.SetMethod("POST")
@ -317,39 +332,41 @@ func (a *Agent) DistributedRead() (*distributedReadResponse, error) {
return &parsedResp, nil
}
type distributedWriteRequest struct {
Queries map[string]json.RawMessage `json:"queries"`
Statuses map[string]string `json:"statuses"`
NodeKey string `json:"node_key"`
var defaultQueryResult = []map[string]string{
{"foo": "bar"},
}
var defaultQueryResult = json.RawMessage(`[{"foo": "bar"}]`)
const statusSuccess = "0"
func (a *Agent) DistributedWrite(queries map[string]string) {
var body bytes.Buffer
if _, ok := queries["fleet_detail_query_network_interface"]; ok {
// Respond to label/detail queries
a.Templates.ExecuteTemplate(&body, "distributed_write", a)
} else {
// Return a generic response for any other queries
req := distributedWriteRequest{
Queries: make(map[string]json.RawMessage),
Statuses: make(map[string]string),
NodeKey: a.NodeKey,
func (a *agent) DistributedWrite(queries map[string]string) {
r := service.SubmitDistributedQueryResultsRequest{
Results: make(fleet.OsqueryDistributedQueryResults),
Statuses: make(map[string]fleet.OsqueryStatus),
}
r.NodeKey = a.NodeKey
for name := range queries {
req.Queries[name] = defaultQueryResult
req.Statuses[name] = statusSuccess
r.Results[name] = defaultQueryResult
r.Statuses[name] = fleet.StatusOK
if t := a.Templates.Lookup(name); t == nil {
continue
}
json.NewEncoder(&body).Encode(req)
var ni bytes.Buffer
err := a.Templates.ExecuteTemplate(&ni, name, a)
if err != nil {
panic(err)
}
var m []map[string]string
err = json.Unmarshal(ni.Bytes(), &m)
if err != nil {
panic(err)
}
r.Results[name] = m
}
body, err := json.Marshal(r)
if err != nil {
panic(err)
}
req := fasthttp.AcquireRequest()
req.SetBody(body.Bytes())
req.SetBody(body)
req.Header.SetMethod("POST")
req.Header.SetContentType("application/json")
req.Header.Add("User-Agent", "osquery/4.6.0")
@ -362,7 +379,6 @@ func (a *Agent) DistributedWrite(queries map[string]string) {
defer fasthttp.ReleaseResponse(res)
a.Stats.RecordStats(0, 0, 1)
// No need to read the distributed write body
}
@ -376,34 +392,34 @@ func main() {
queryInterval := flag.Duration("query_interval", 10*time.Second, "Interval for live query requests")
onlyAlreadyEnrolled := flag.Bool("only_already_enrolled", false, "Only start agents that are already enrolled")
nodeKeyFile := flag.String("node_key_file", "", "File with node keys to use")
softwareCount := flag.Int("software_count", 10, "Number of installed applications reported to fleet")
flag.Parse()
rand.Seed(*randSeed)
tmpl, err := template.ParseFS(templatesFS, "*.tmpl")
// Currently all hosts will be macOS.
tmpl, err := template.ParseFS(templatesFS, "mac10.14.6.tmpl")
if err != nil {
log.Fatal("parse templates: ", err)
}
// Spread starts over the interval to prevent thunering herd
// Spread starts over the interval to prevent thundering herd
sleepTime := *startPeriod / time.Duration(*hostCount)
stats := &Stats{}
go stats.runLoop()
nodeKeyManager := &NodeKeyManager{}
nodeKeyManager := &nodeKeyManager{}
if nodeKeyFile != nil {
nodeKeyManager.filepath = *nodeKeyFile
nodeKeyManager.LoadKeys()
}
var agents []*Agent
for i := 0; i < *hostCount; i++ {
a := NewAgent(*serverURL, *enrollSecret, tmpl, *configInterval, *queryInterval)
a := newAgent(*serverURL, *enrollSecret, tmpl, *configInterval, *queryInterval, *softwareCount)
a.Stats = stats
a.NodeKeyManager = nodeKeyManager
agents = append(agents, a)
go a.runLoop(i, onlyAlreadyEnrolled != nil && *onlyAlreadyEnrolled)
time.Sleep(sleepTime)
}

149
cmd/osquery-perf/mac10.14.6.tmpl
View File

@ -54,14 +54,12 @@
}
},
"host_identifier": "{{ .CachedString "hostname" }}",
"platform_type": "21"
"platform_type": "16"
}
{{- end }}
{{ define "distributed_write" -}}
{
"queries":{
"fleet_detail_query_network_interface":[
{{ define "fleet_detail_query_network_interface" -}}
[
{
"point_to_point":"",
"address":"fe80::8cb:112d:ff51:1e5d%en0",
@ -81,7 +79,6 @@
"idrops":"0",
"odrops":"0",
"last_change":"1582848084"
},
{
"point_to_point":"",
@ -102,7 +99,6 @@
"idrops":"0",
"odrops":"0",
"last_change":"1582848084"
},
{
"point_to_point":"127.0.0.1",
@ -123,7 +119,6 @@
"idrops":"0",
"odrops":"0",
"last_change":"1582840871"
},
{
"point_to_point":"::1",
@ -144,7 +139,6 @@
"idrops":"0",
"odrops":"0",
"last_change":"1582840871"
},
{
"point_to_point":"",
@ -165,7 +159,6 @@
"idrops":"0",
"odrops":"0",
"last_change":"1582840871"
},
{
"point_to_point":"",
@ -186,7 +179,6 @@
"idrops":"0",
"odrops":"0",
"last_change":"1582842892"
},
{
"point_to_point":"",
@ -207,11 +199,11 @@
"idrops":"0",
"odrops":"0",
"last_change":"1582840897"
}
],
"fleet_detail_query_os_version":[
]
{{- end }}
{{ define "fleet_detail_query_os_version" -}}
[
{
"name":"Mac OS X",
"version":"10.14.6",
@ -222,29 +214,27 @@
"platform":"darwin",
"platform_like":"darwin",
"codename":""
}
],
"fleet_detail_query_osquery_flags":[
]
{{- end }}
{{ define "fleet_detail_query_osquery_flags" -}}
[
{
"name":"config_refresh",
"value":"{{ printf "%.0f" .ConfigInterval.Seconds }}"
},
{
"name":"distributed_interval",
"value":"{{ printf "%.0f" .QueryInterval.Seconds }}"
},
{
"name":"logger_tls_period",
"value":"99999"
}
],
"fleet_detail_query_osquery_info":[
]
{{- end }}
{{ define "fleet_detail_query_osquery_info" -}}
[
{
"pid":"11287",
"uuid":"{{ .UUID }}",
@ -258,11 +248,11 @@
"start_time":"1582857555",
"watcher":"11286",
"platform_mask":"21"
}
],
"fleet_detail_query_system_info":[
]
{{- end }}
{{ define "fleet_detail_query_system_info" -}}
[
{
"hostname":"{{ .CachedString "hostname" }}",
"uuid":"4740D59F-699E-5B29-960B-979AAF9BBEEB",
@ -279,31 +269,104 @@
"hardware_serial":"C02R262BM8LN",
"computer_name":"{{ .CachedString "hostname" }}",
"local_hostname":"{{ .CachedString "hostname" }}"
}
],
"fleet_detail_query_uptime":[
]
{{- end }}
{{ define "fleet_detail_query_uptime" -}}
[
{
"days":"0",
"hours":"4",
"minutes":"38",
"seconds":"11",
"total_seconds":"16691"
}
]
{{- end }}
{{ define "fleet_detail_query_users" -}}
[
{
"uid":"0",
"username":"gandalf",
"type":"wizard",
"groupname":"bree"
},
"statuses":{
"fleet_detail_query_network_interface":0,
"fleet_detail_query_os_version":0,
"fleet_detail_query_osquery_flags":0,
"fleet_detail_query_osquery_info":0,
"fleet_detail_query_system_info":0,
"fleet_detail_query_uptime":0
{
"uid":"1",
"username":"frodo",
"type":"hobbit",
"groupname":"shire"
},
"node_key":"{{ .NodeKey }}"
{
"uid":"1",
"username":"sam",
"type":"hobbit",
"groupname":"shire"
},
{
"uid":"1",
"username":"legolas",
"type":"elve",
"groupname":"mirkwood"
}
]
{{- end }}
{{/* all hosts */}}
{{ define "fleet_label_query_6" -}}
[
{
"1": "1"
}
]
{{- end }}
{{/* All macOS hosts */}}
{{ define "fleet_label_query_7" -}}
[
{
"1": "1"
}
]
{{- end }}
{{/* All Ubuntu hosts */}}
{{ define "fleet_label_query_8" -}}
[]
{{- end }}
{{/* All CentOS hosts */}}
{{ define "fleet_label_query_9" -}}
[]
{{- end }}
{{/* All Windows hosts */}}
{{ define "fleet_label_query_10" -}}
[]
{{- end }}
{{/* All Red Hat hosts */}}
{{ define "fleet_label_query_11" -}}
[]
{{- end }}
{{/* All Linux distributions */}}
{{ define "fleet_label_query_12" -}}
[]
{{- end }}
{{ define "fleet_detail_query_software_macos" -}}
[
{{ range $index, $item := .SoftwareMacOS }}
{{if $index}},{{end}}
{
"name": "{{ .Name }}_{{ $index }}",
"version": "{{ .Version }}",
"type": "Application (macOS)",
"bundle_identifier": "{{ .BundleIdentifier }}",
"source": "apps"
}
{{- end }}
]
{{- end }}

4
server/service/endpoint_osquery.go
View File

@ -94,7 +94,7 @@ func makeGetDistributedQueriesEndpoint(svc fleet.Service) endpoint.Endpoint {
// Write Distributed Query Results
////////////////////////////////////////////////////////////////////////////////
type submitDistributedQueryResultsRequest struct {
type SubmitDistributedQueryResultsRequest struct {
NodeKey string `json:"node_key"`
Results fleet.OsqueryDistributedQueryResults `json:"queries"`
Statuses map[string]fleet.OsqueryStatus `json:"statuses"`
@ -109,7 +109,7 @@ func (r submitDistributedQueryResultsResponse) error() error { return r.Err }
func makeSubmitDistributedQueryResultsEndpoint(svc fleet.Service) endpoint.Endpoint {
return func(ctx context.Context, request interface{}) (interface{}, error) {
req := request.(submitDistributedQueryResultsRequest)
req := request.(SubmitDistributedQueryResultsRequest)
err := svc.SubmitDistributedQueryResults(ctx, req.Results, req.Statuses, req.Messages)
if err != nil {
return submitDistributedQueryResultsResponse{Err: err}, nil

10
server/service/integration_live_queries_test.go
View File

@ -95,7 +95,7 @@ func (s *liveQueriesTestSuite) TestLiveQueriesRestOneHostOneQuery() {
cid := getCIDForQ(s, q1)
distributedReq := submitDistributedQueryResultsRequest{
distributedReq := SubmitDistributedQueryResultsRequest{
NodeKey: host.NodeKey,
Results: map[string][]map[string]string{
hostDistributedQueryPrefix + cid: {{"col1": "a", "col2": "b"}},
@ -159,7 +159,7 @@ func (s *liveQueriesTestSuite) TestLiveQueriesRestOneHostMultipleQuery() {
cid1 := getCIDForQ(s, q1)
cid2 := getCIDForQ(s, q2)
distributedReq := submitDistributedQueryResultsRequest{
distributedReq := SubmitDistributedQueryResultsRequest{
NodeKey: host.NodeKey,
Results: map[string][]map[string]string{
hostDistributedQueryPrefix + cid1: {{"col1": "a", "col2": "b"}},
@ -257,7 +257,7 @@ func (s *liveQueriesTestSuite) TestLiveQueriesRestMultipleHostMultipleQuery() {
cid1 := getCIDForQ(s, q1)
cid2 := getCIDForQ(s, q2)
for i, h := range []*fleet.Host{h1, h2} {
distributedReq := submitDistributedQueryResultsRequest{
distributedReq := SubmitDistributedQueryResultsRequest{
NodeKey: h.NodeKey,
Results: map[string][]map[string]string{
hostDistributedQueryPrefix + cid1: {{"col1": fmt.Sprintf("a%d", i), "col2": fmt.Sprintf("b%d", i)}},
@ -352,7 +352,7 @@ func (s *liveQueriesTestSuite) TestLiveQueriesRestFailsOnSomeHost() {
// Give the above call a couple of seconds to create the campaign
time.Sleep(2 * time.Second)
cid1 := getCIDForQ(s, q1)
distributedReq := submitDistributedQueryResultsRequest{
distributedReq := SubmitDistributedQueryResultsRequest{
NodeKey: h1.NodeKey,
Results: map[string][]map[string]string{
hostDistributedQueryPrefix + cid1: {{"col1": "a", "col2": "b"}},
@ -367,7 +367,7 @@ func (s *liveQueriesTestSuite) TestLiveQueriesRestFailsOnSomeHost() {
distributedResp := submitDistributedQueryResultsResponse{}
s.DoJSON("POST", "/api/v1/osquery/distributed/write", distributedReq, http.StatusOK, &distributedResp)
distributedReq = submitDistributedQueryResultsRequest{
distributedReq = SubmitDistributedQueryResultsRequest{
NodeKey: h2.NodeKey,
Results: map[string][]map[string]string{
hostDistributedQueryPrefix + cid1: {},

2
server/service/transport_osquery.go
View File

@ -95,7 +95,7 @@ func decodeSubmitDistributedQueryResultsRequest(ctx context.Context, r *http.Req
}
}
req := submitDistributedQueryResultsRequest{
req := SubmitDistributedQueryResultsRequest{
NodeKey: shim.NodeKey,
Results: results,
Statuses: statuses,

2
server/service/transport_osquery_test.go
View File

@ -85,7 +85,7 @@ func TestDecodeSubmitDistributedQueryResultsRequest(t *testing.T) {
r, err := decodeSubmitDistributedQueryResultsRequest(context.Background(), request)
require.Nil(t, err)
params := r.(submitDistributedQueryResultsRequest)
params := r.(SubmitDistributedQueryResultsRequest)
assert.Equal(t, "key", params.NodeKey)
assert.Equal(t, fleet.OsqueryDistributedQueryResults{
"id1": {