github.com/Uptycs/basequery-go@v0.8.0/plugin/distributed/distributed_test.go (about) 1 package distributed 2 3 import ( 4 "bytes" 5 "context" 6 "encoding/json" 7 "errors" 8 "fmt" 9 "sort" 10 "testing" 11 12 "github.com/Uptycs/basequery-go/gen/osquery" 13 "github.com/stretchr/testify/assert" 14 "github.com/stretchr/testify/require" 15 ) 16 17 var StatusOK = osquery.ExtensionStatus{Code: 0, Message: "OK"} 18 19 func TestDistributedPlugin(t *testing.T) { 20 var getCalled, writeCalled bool 21 var results []Result 22 plugin := NewPlugin( 23 "mock", 24 func(context.Context) (*GetQueriesResult, error) { 25 getCalled = true 26 return &GetQueriesResult{ 27 Queries: map[string]string{ 28 "query1": "select iso_8601 from time", 29 "query2": "select version from osquery_info", 30 "query3": "select foo from bar", 31 }, 32 }, nil 33 }, 34 func(ctx context.Context, res []Result) error { 35 writeCalled = true 36 results = res 37 return nil 38 }, 39 ) 40 41 // Basic methods 42 assert.Equal(t, "distributed", plugin.RegistryName()) 43 assert.Equal(t, "mock", plugin.Name()) 44 assert.Equal(t, StatusOK, plugin.Ping()) 45 assert.Equal(t, osquery.ExtensionPluginResponse{}, plugin.Routes()) 46 47 // Call getQueries 48 resp := plugin.Call(context.Background(), osquery.ExtensionPluginRequest{"action": "getQueries"}) 49 assert.True(t, getCalled) 50 assert.False(t, writeCalled) 51 assert.Equal(t, &StatusOK, resp.Status) 52 if assert.Len(t, resp.Response, 1) { 53 assert.JSONEq(t, `{"queries": {"query1": "select iso_8601 from time", "query2": "select version from osquery_info", "query3": "select foo from bar"}}`, 54 resp.Response[0]["results"]) 55 } 56 57 // Call writeResults 58 getCalled = false 59 resp = plugin.Call(context.Background(), osquery.ExtensionPluginRequest{"action": "writeResults", "results": `{"queries":{"query1":[{"iso_8601":"2017-07-10T22:08:40Z"}],"query2":[{"version":"2.4.0"}]},"statuses":{"query1":"0","query2":"0","query3":"1"}}`}) 60 assert.False(t, getCalled) 61 assert.True(t, writeCalled) 62 assert.Equal(t, &StatusOK, resp.Status) 63 // Ensure correct ordering for comparison 64 sort.Slice(results, func(i, j int) bool { return results[i].QueryName < results[j].QueryName }) 65 assert.Equal(t, []Result{ 66 {"query1", 0, []map[string]string{{"iso_8601": "2017-07-10T22:08:40Z"}}}, 67 {"query2", 0, []map[string]string{{"version": "2.4.0"}}}, 68 {"query3", 1, []map[string]string{}}, 69 }, 70 results) 71 } 72 73 func TestDistributedPluginAccelerateDiscovery(t *testing.T) { 74 plugin := NewPlugin( 75 "mock", 76 func(context.Context) (*GetQueriesResult, error) { 77 return &GetQueriesResult{ 78 Queries: map[string]string{ 79 "query1": "select * from time", 80 }, 81 Discovery: map[string]string{ 82 "query1": `select version from osquery_info where version = "2.4.0"`, 83 }, 84 AccelerateSeconds: 30, 85 }, nil 86 }, 87 func(ctx context.Context, res []Result) error { 88 return nil 89 }, 90 ) 91 92 // Call getQueries 93 resp := plugin.Call(context.Background(), osquery.ExtensionPluginRequest{"action": "getQueries"}) 94 assert.Equal(t, &StatusOK, resp.Status) 95 if assert.Len(t, resp.Response, 1) { 96 assert.JSONEq(t, `{"queries": {"query1": "select * from time"}, "discovery": {"query1": "select version from osquery_info where version = \"2.4.0\""}, "accelerate": 30}`, 97 resp.Response[0]["results"]) 98 } 99 } 100 101 func TestDistributedPluginErrors(t *testing.T) { 102 var getCalled, writeCalled bool 103 plugin := NewPlugin( 104 "mock", 105 func(context.Context) (*GetQueriesResult, error) { 106 getCalled = true 107 return nil, errors.New("getQueries failed") 108 }, 109 func(ctx context.Context, res []Result) error { 110 writeCalled = true 111 return errors.New("writeResults failed") 112 }, 113 ) 114 115 // Call with bad actions 116 assert.Equal(t, int32(1), plugin.Call(context.Background(), osquery.ExtensionPluginRequest{}).Status.Code) 117 assert.False(t, getCalled) 118 assert.False(t, writeCalled) 119 assert.Equal(t, int32(1), plugin.Call(context.Background(), osquery.ExtensionPluginRequest{"action": "bad"}).Status.Code) 120 assert.False(t, getCalled) 121 assert.False(t, writeCalled) 122 123 // Call with good action but getQueries fails 124 resp := plugin.Call(context.Background(), osquery.ExtensionPluginRequest{"action": "getQueries"}) 125 assert.True(t, getCalled) 126 assert.False(t, writeCalled) 127 assert.Equal(t, int32(1), resp.Status.Code) 128 assert.Equal(t, "error getting queries: getQueries failed", resp.Status.Message) 129 130 getCalled = false 131 132 // Call with good action but writeResults fails 133 // Error unmarshalling results 134 resp = plugin.Call(context.Background(), osquery.ExtensionPluginRequest{"action": "writeResults", "results": "foobar"}) 135 assert.False(t, getCalled) 136 assert.False(t, writeCalled) 137 assert.Equal(t, int32(1), resp.Status.Code) 138 assert.Contains(t, resp.Status.Message, "error unmarshalling results") 139 140 // Error converting status 141 resp = plugin.Call(context.Background(), osquery.ExtensionPluginRequest{"action": "writeResults", "results": `{"statuses": {"query1": "foo"}}`}) 142 assert.False(t, getCalled) 143 assert.False(t, writeCalled) 144 assert.Equal(t, int32(1), resp.Status.Code) 145 assert.Contains(t, resp.Status.Message, `error unmarshalling results: json: cannot unmarshal`) 146 147 // Error unmarshalling results 148 resp = plugin.Call(context.Background(), osquery.ExtensionPluginRequest{"action": "writeResults", "results": "{}"}) 149 assert.False(t, getCalled) 150 assert.True(t, writeCalled) 151 assert.Equal(t, int32(1), resp.Status.Code) 152 assert.Contains(t, resp.Status.Message, "error writing results") 153 } 154 155 var rawJSONQuery = "{\"queries\":{\"kolide_detail_query_network_interface\":[{\"interface\":\"en0\",\"mac\":\"78:4f:43:9c:3c:8d\",\"type\":\"\",\"mtu\":\"1500\",\"metric\":\"0\",\"ipackets\":\"7071136\",\"opackets\":\"6408727\",\"ibytes\":\"1481456771\",\"obytes\":\"1633052673\",\"ierrors\":\"0\",\"oerrors\":\"0\",\"idrops\":\"0\",\"odrops\":\"0\",\"last_change\":\"1501077669\",\"description\":\"\",\"manufacturer\":\"\",\"connection_id\":\"\",\"connection_status\":\"\",\"enabled\":\"\",\"physical_adapter\":\"\",\"speed\":\"\",\"dhcp_enabled\":\"\",\"dhcp_lease_expires\":\"\",\"dhcp_lease_obtained\":\"\",\"dhcp_server\":\"\",\"dns_domain\":\"\",\"dns_domain_suffix_search_order\":\"\",\"dns_host_name\":\"\",\"dns_server_search_order\":\"\",\"interface\":\"en0\",\"address\":\"192.168.1.135\",\"mask\":\"255.255.255.0\",\"broadcast\":\"192.168.1.255\",\"point_to_point\":\"\",\"type\":\"\"}],\"kolide_detail_query_os_version\":[{\"name\":\"Mac OS X\",\"version\":\"10.12.6\",\"major\":\"10\",\"minor\":\"12\",\"patch\":\"6\",\"build\":\"16G29\",\"platform\":\"darwin\",\"platform_like\":\"darwin\",\"codename\":\"\"}],\"kolide_detail_query_osquery_flags\":[{\"name\":\"config_refresh\",\"value\":\"10\"},{\"name\":\"distributed_interval\",\"value\":\"10\"},{\"name\":\"logger_tls_period\",\"value\":\"10\"}],\"kolide_detail_query_osquery_info\":[{\"pid\":\"75680\",\"uuid\":\"DE56C776-2F5A-56DF-81C7-F64EE1BBEC8C\",\"instance_id\":\"89f267fa-9a17-4a73-87d6-05197491f2e8\",\"version\":\"2.5.0\",\"config_hash\":\"960121acb9bcbb136ce49fe77000752f237fd0dd\",\"config_valid\":\"1\",\"extensions\":\"active\",\"build_platform\":\"darwin\",\"build_distro\":\"10.12\",\"start_time\":\"1502371429\",\"watcher\":\"75678\"}],\"kolide_detail_query_system_info\":[{\"hostname\":\"Johns-MacBook-Pro.local\",\"uuid\":\"DE56C776-2F5A-56DF-81C7-F64EE1BBEC8C\",\"cpu_type\":\"x86_64h\",\"cpu_subtype\":\"Intel x86-64h Haswell\",\"cpu_brand\":\"Intel(R) Core(TM) i7-6820HQ CPU @ 2.70GHz\",\"cpu_physical_cores\":\"4\",\"cpu_logical_cores\":\"8\",\"physical_memory\":\"17179869184\",\"hardware_vendor\":\"Apple Inc.\",\"hardware_model\":\"MacBookPro13,3\",\"hardware_version\":\"1.0\",\"hardware_serial\":\"C02SP067H040\",\"computer_name\":\"\",\"local_hostname\":\"Johns-MacBook-Pro\"}],\"kolide_detail_query_uptime\":[{\"days\":\"21\",\"hours\":\"18\",\"minutes\":\"44\",\"seconds\":\"28\",\"total_seconds\":\"1881868\"}],\"kolide_label_query_6\":[{\"1\":\"1\"}],\"kolide_label_query_9\":\"\",\"kolide_detail_query_network_interface\":[{\"interface\":\"en0\",\"mac\":\"78:4f:43:9c:3c:8d\",\"type\":\"\",\"mtu\":\"1500\",\"metric\":\"0\",\"ipackets\":\"7071178\",\"opackets\":\"6408775\",\"ibytes\":\"1481473778\",\"obytes\":\"1633061382\",\"ierrors\":\"0\",\"oerrors\":\"0\",\"idrops\":\"0\",\"odrops\":\"0\",\"last_change\":\"1501077680\",\"description\":\"\",\"manufacturer\":\"\",\"connection_id\":\"\",\"connection_status\":\"\",\"enabled\":\"\",\"physical_adapter\":\"\",\"speed\":\"\",\"dhcp_enabled\":\"\",\"dhcp_lease_expires\":\"\",\"dhcp_lease_obtained\":\"\",\"dhcp_server\":\"\",\"dns_domain\":\"\",\"dns_domain_suffix_search_order\":\"\",\"dns_host_name\":\"\",\"dns_server_search_order\":\"\",\"interface\":\"en0\",\"address\":\"192.168.1.135\",\"mask\":\"255.255.255.0\",\"broadcast\":\"192.168.1.255\",\"point_to_point\":\"\",\"type\":\"\"}],\"kolide_detail_query_os_version\":[{\"name\":\"Mac OS X\",\"version\":\"10.12.6\",\"major\":\"10\",\"minor\":\"12\",\"patch\":\"6\",\"build\":\"16G29\",\"platform\":\"darwin\",\"platform_like\":\"darwin\",\"codename\":\"\"}],\"kolide_detail_query_osquery_flags\":[{\"name\":\"config_refresh\",\"value\":\"10\"},{\"name\":\"distributed_interval\",\"value\":\"10\"},{\"name\":\"logger_tls_period\",\"value\":\"10\"}],\"kolide_detail_query_osquery_info\":[{\"pid\":\"75680\",\"uuid\":\"DE56C776-2F5A-56DF-81C7-F64EE1BBEC8C\",\"instance_id\":\"89f267fa-9a17-4a73-87d6-05197491f2e8\",\"version\":\"2.5.0\",\"config_hash\":\"960121acb9bcbb136ce49fe77000752f237fd0dd\",\"config_valid\":\"1\",\"extensions\":\"active\",\"build_platform\":\"darwin\",\"build_distro\":\"10.12\",\"start_time\":\"1502371429\",\"watcher\":\"75678\"}],\"kolide_detail_query_system_info\":[{\"hostname\":\"Johns-MacBook-Pro.local\",\"uuid\":\"DE56C776-2F5A-56DF-81C7-F64EE1BBEC8C\",\"cpu_type\":\"x86_64h\",\"cpu_subtype\":\"Intel x86-64h Haswell\",\"cpu_brand\":\"Intel(R) Core(TM) i7-6820HQ CPU @ 2.70GHz\",\"cpu_physical_cores\":\"4\",\"cpu_logical_cores\":\"8\",\"physical_memory\":\"17179869184\",\"hardware_vendor\":\"Apple Inc.\",\"hardware_model\":\"MacBookPro13,3\",\"hardware_version\":\"1.0\",\"hardware_serial\":\"C02SP067H040\",\"computer_name\":\"\",\"local_hostname\":\"Johns-MacBook-Pro\"}],\"kolide_detail_query_uptime\":[{\"days\":\"21\",\"hours\":\"18\",\"minutes\":\"44\",\"seconds\":\"38\",\"total_seconds\":\"1881878\"}],\"kolide_label_query_6\":[{\"1\":\"1\"}],\"kolide_label_query_9\":\"\",\"kolide_detail_query_network_interface\":[{\"interface\":\"en0\",\"mac\":\"78:4f:43:9c:3c:8d\",\"type\":\"\",\"mtu\":\"1500\",\"metric\":\"0\",\"ipackets\":\"7071216\",\"opackets\":\"6408814\",\"ibytes\":\"1481486677\",\"obytes\":\"1633066361\",\"ierrors\":\"0\",\"oerrors\":\"0\",\"idrops\":\"0\",\"odrops\":\"0\",\"last_change\":\"1501077688\",\"description\":\"\",\"manufacturer\":\"\",\"connection_id\":\"\",\"connection_status\":\"\",\"enabled\":\"\",\"physical_adapter\":\"\",\"speed\":\"\",\"dhcp_enabled\":\"\",\"dhcp_lease_expires\":\"\",\"dhcp_lease_obtained\":\"\",\"dhcp_server\":\"\",\"dns_domain\":\"\",\"dns_domain_suffix_search_order\":\"\",\"dns_host_name\":\"\",\"dns_server_search_order\":\"\",\"interface\":\"en0\",\"address\":\"192.168.1.135\",\"mask\":\"255.255.255.0\",\"broadcast\":\"192.168.1.255\",\"point_to_point\":\"\",\"type\":\"\"}],\"kolide_detail_query_os_version\":[{\"name\":\"Mac OS X\",\"version\":\"10.12.6\",\"major\":\"10\",\"minor\":\"12\",\"patch\":\"6\",\"build\":\"16G29\",\"platform\":\"darwin\",\"platform_like\":\"darwin\",\"codename\":\"\"}],\"kolide_detail_query_osquery_flags\":[{\"name\":\"config_refresh\",\"value\":\"10\"},{\"name\":\"distributed_interval\",\"value\":\"10\"},{\"name\":\"logger_tls_period\",\"value\":\"10\"}],\"kolide_detail_query_osquery_info\":[{\"pid\":\"75680\",\"uuid\":\"DE56C776-2F5A-56DF-81C7-F64EE1BBEC8C\",\"instance_id\":\"89f267fa-9a17-4a73-87d6-05197491f2e8\",\"version\":\"2.5.0\",\"config_hash\":\"960121acb9bcbb136ce49fe77000752f237fd0dd\",\"config_valid\":\"1\",\"extensions\":\"active\",\"build_platform\":\"darwin\",\"build_distro\":\"10.12\",\"start_time\":\"1502371429\",\"watcher\":\"75678\"}],\"kolide_detail_query_system_info\":[{\"hostname\":\"Johns-MacBook-Pro.local\",\"uuid\":\"DE56C776-2F5A-56DF-81C7-F64EE1BBEC8C\",\"cpu_type\":\"x86_64h\",\"cpu_subtype\":\"Intel x86-64h Haswell\",\"cpu_brand\":\"Intel(R) Core(TM) i7-6820HQ CPU @ 2.70GHz\",\"cpu_physical_cores\":\"4\",\"cpu_logical_cores\":\"8\",\"physical_memory\":\"17179869184\",\"hardware_vendor\":\"Apple Inc.\",\"hardware_model\":\"MacBookPro13,3\",\"hardware_version\":\"1.0\",\"hardware_serial\":\"C02SP067H040\",\"computer_name\":\"\",\"local_hostname\":\"Johns-MacBook-Pro\"}],\"kolide_detail_query_uptime\":[{\"days\":\"21\",\"hours\":\"18\",\"minutes\":\"44\",\"seconds\":\"49\",\"total_seconds\":\"1881889\"}],\"kolide_label_query_6\":[{\"1\":\"1\"}],\"kolide_label_query_9\":\"\"},\"statuses\":{\"kolide_detail_query_network_interface\":\"0\",\"kolide_detail_query_os_version\":\"0\",\"kolide_detail_query_osquery_flags\":\"0\",\"kolide_detail_query_osquery_info\":\"0\",\"kolide_detail_query_system_info\":\"0\",\"kolide_detail_query_uptime\":\"0\",\"kolide_label_query_6\":\"0\",\"kolide_label_query_9\":\"0\"}}\n" 156 157 func TestUnmarshalResults(t *testing.T) { 158 var rs ResultsStruct 159 err := json.NewDecoder(bytes.NewBufferString(rawJSONQuery)).Decode(&rs) 160 require.Nil(t, err) 161 results, err := rs.toResults() 162 require.Nil(t, err) 163 assert.Len(t, results, 8) 164 } 165 166 func TestUnmarshalStatus(t *testing.T) { 167 testCases := []struct { 168 json []byte 169 success bool 170 expected OsqueryInt 171 }{ 172 {[]byte{}, true, 0}, 173 {[]byte(`""`), true, 0}, 174 {[]byte(`"23"`), true, 23}, 175 {[]byte(`"0000"`), true, 0}, 176 {[]byte(`"-12"`), true, -12}, 177 {[]byte(`"0"`), true, 0}, 178 {[]byte(`"foo"`), false, 0}, 179 {[]byte(`0`), true, 0}, 180 {[]byte(`1`), true, 1}, 181 } 182 for i, testCase := range testCases { 183 t.Run(fmt.Sprintf("#%.2d", i), func(t *testing.T) { 184 var i OsqueryInt 185 err := i.UnmarshalJSON(testCase.json) 186 require.Equal(t, testCase.success, (err == nil), fmt.Sprintf("Trying to convert %s", string(testCase.json))) 187 assert.Equal(t, testCase.expected, i) 188 }) 189 } 190 } 191 192 func TestHandleResults(t *testing.T) { 193 called := false 194 var results []Result 195 plugin := NewPlugin( 196 "mock", 197 nil, 198 func(ctx context.Context, res []Result) error { 199 called = true 200 results = res 201 return nil 202 }, 203 ) 204 resp := plugin.Call(context.Background(), osquery.ExtensionPluginRequest{"action": "writeResults", "results": rawJSONQuery}) 205 require.True(t, called) 206 assert.Len(t, results, 8) 207 assert.Equal(t, &StatusOK, resp.Status) 208 }