feat(cli): add lacework query preview-source

ALLY-879

cli/cmd/lql_preview.go

@@ -0,0 +1,127 @@
+//
+// Author:: Salim Afiune Maya (<[email protected]>)
+// Copyright:: Copyright 2020, Lacework Inc.
+// License:: Apache License, Version 2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+package cmd
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/pkg/errors"
+	"github.com/spf13/cobra"
+
+	"github.com/lacework/go-sdk/api"
+	"github.com/lacework/go-sdk/lwtime"
+)
+
+var (
+	queryPreviewSourceCmd = &cobra.Command{
+		Use:   "preview-source <datasource_id>",
+		Short: "Preview Lacework query data source",
+		Long:  `Preview Lacework query data source.`,
+		Args:  cobra.ExactArgs(1),
+		RunE:  previewQuerySource,
+	}
+	queryPreviewSourceTemplate = `CLIAdhocPreview { source { %s } return distinct { %s } }`
+)
+
+func init() {
+	queryCmd.AddCommand(queryPreviewSourceCmd)
+}
+
+func previewQuerySource(_ *cobra.Command, args []string) error {
+	cli.Log.Debugw("retrieving datasource", "id", args[0])
+
+	cli.StartProgress(" Retrieving datasource...")
+	datasourceResponse, err := cli.LwApi.V2.Datasources.Get(args[0])
+	cli.StopProgress()
+
+	if err != nil {
+		return errors.Wrap(err, "unable to retrieve data source")
+	}
+
+	// build returns list from datasource fields
+	var returns []string
+	for _, ret := range datasourceResponse.Data.ResultSchema {
+		returns = append(returns, ret.Name)
+	}
+	if len(returns) == 0 {
+		return errors.New("unable to parse data source schema")
+	}
+
+	// initialize query
+	executeQuery := api.ExecuteQueryRequest{
+		Query: api.ExecuteQuery{
+			QueryText: fmt.Sprintf(
+				queryPreviewSourceTemplate, args[0], strings.Join(returns, ",")),
+		},
+	}
+	if args[0] == "CloudTrailRawEvents" {
+		executeQuery.Query.EvaluatorID = "Cloudtrail"
+	}
+
+	// initialize time attempts
+	timeAttempts := []map[string]string{
+		{"start": "-24h", "end": "now"},
+		{"start": "-7d", "end": "-24h"},
+		{"start": "-30d", "end": "-7d"},
+	}
+
+	for _, timeAttempt := range timeAttempts {
+		start, _ := lwtime.ParseRelative(timeAttempt["start"])
+		end, _ := lwtime.ParseRelative(timeAttempt["end"])
+
+		executeQuery.Arguments = []api.ExecuteQueryArgument{
+			api.ExecuteQueryArgument{
+				Name:  "StartTimeRange",
+				Value: start.UTC().Format(lwtime.RFC3339Milli),
+			},
+			api.ExecuteQueryArgument{
+				Name:  "EndTimeRange",
+				Value: end.UTC().Format(lwtime.RFC3339Milli),
+			},
+		}
+
+		// execute query
+		cli.Log.Debugw("running query", "query", executeQuery.Query.QueryText)
+		cli.StartProgress(" Executing preview query...")
+		response, err := cli.LwApi.V2.Query.Execute(executeQuery)
+		cli.StopProgress()
+		if err != nil {
+			return errors.Wrap(err, "unable to preview data source")
+		}
+
+		// check and output
+		data, ok := response["data"]
+		if !ok {
+			err = errors.New("preview results missing data")
+			return errors.Wrap(err, "unable to preview data source")
+		}
+		if slice, ok := data.([]interface{}); ok {
+			if len(slice) == 0 {
+				continue
+			}
+			return cli.OutputJSON(slice[0])
+		} else {
+			err = errors.New("preview results data is not a slice")
+			return errors.Wrap(err, "unable to preview data source")
+		}
+	}
+	cli.OutputHuman("No results found for data source")
+	return nil
+}

cli/cmd/lql_sources.go

@@ -28,7 +28,6 @@ import (
 
 var (
 	queryListSourcesCmd = &cobra.Command{
-		Hidden:  false,
 		Aliases: []string{"sources"},
 		Use:     "list-sources",
 		Short:   "List Lacework query data sources",
@@ -38,7 +37,6 @@ var (
 	}
 
 	queryShowSourceCmd = &cobra.Command{
-		Hidden:  false,
 		Aliases: []string{"describe"},
 		Use:     "show-source <datasource_id>",
 		Short:   "Show Lacework query data source",

integration/lql_preview_test.go

@@ -0,0 +1,53 @@
+//go:build query
+
+// Author:: Salim Afiune Maya (<[email protected]>)
+// Copyright:: Copyright 2020, Lacework Inc.
+// License:: Apache License, Version 2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+package integration
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestQueryPreviewSourceHelp(t *testing.T) {
+	out, err, exitcode := LaceworkCLI("help", "query", "preview-source")
+	assert.Contains(t, out.String(), "lacework query preview-source <datasource_id> [flags]")
+	assert.Empty(t, err.String(), "STDERR should be empty")
+	assert.Equal(t, 0, exitcode, "EXITCODE is not the expected one")
+}
+
+func TestQueryPreviewSourceNoInput(t *testing.T) {
+	out, err, exitcode := LaceworkCLIWithTOMLConfig("query", "preview-source")
+	assert.Empty(t, out.String(), "STDOUT should be empty")
+	assert.Contains(t, err.String(), "ERROR accepts 1 arg(s), received 0")
+	assert.Equal(t, 1, exitcode, "EXITCODE is not the expected one")
+}
+
+func TestQueryPreviewSource(t *testing.T) {
+	out, err, exitcode := LaceworkCLIWithTOMLConfig("query", "preview-source", "CloudTrailRawEvents")
+	assert.Contains(t, out.String(), `"INSERT_ID"`)
+	assert.Empty(t, err.String(), "STDERR should be empty")
+	assert.Equal(t, 0, exitcode, "EXITCODE is not the expected one")
+}
+
+func TestQueryPreviewNoSuchSource(t *testing.T) {
+	_, err, exitcode := LaceworkCLIWithTOMLConfig("query", "preview-source", "NoSuchSource")
+	assert.Contains(t, err.String(), "unable to retrieve data source")
+	assert.Equal(t, 1, exitcode, "EXITCODE is not the expected one")
+}