Add namespace support to kolide_wmi (#623)

Adds support for specifying the connection parameters to the underlying wmi package, and plumbs namespace through to the table.

Additionally add `tablehelper` methods to improve code readability when parsing column constraints

Makefile

@@ -38,6 +38,9 @@ launcher: .pre-build
 
 table.ext: .pre-build
 	go run cmd/make/make.go -targets=table-extension -linkstamp
+table.ext-windows: .pre-build deps
+	go run cmd/make/make.go -targets=table-extension -linkstamp --os windows
+
 
 osqueryi-tables: table.ext
 	osqueryd -S --allow-unsafe --verbose --extension ./build/darwin/tables.ext

pkg/osquery/tables/tablehelpers/getconstraints.go

@@ -0,0 +1,30 @@
+package tablehelpers
+
+import (
+	"github.com/kolide/osquery-go/plugin/table"
+)
+
+// GetConstraints returns a []string of the constraint expressions on
+// a column. It's meant for the common, simple, usecase of iterating over them.
+func GetConstraints(queryContext table.QueryContext, columnName string, defaults ...string) []string {
+	q, ok := queryContext.Constraints[columnName]
+	if !ok || len(q.Constraints) == 0 {
+		return defaults
+	}
+
+	constraintSet := make(map[string]struct{})
+
+	for _, c := range q.Constraints {
+		constraintSet[c.Expression] = struct{}{}
+	}
+
+	constraints := make([]string, len(constraintSet))
+
+	i := 0
+	for key := range constraintSet {
+		constraints[i] = key
+		i++
+	}
+
+	return constraints
+}

pkg/osquery/tables/tablehelpers/getconstraints_test.go

@@ -0,0 +1,83 @@
+package tablehelpers
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestGetConstraints(t *testing.T) {
+	t.Parallel()
+
+	mockQC := MockQueryContext(map[string][]string{
+		"empty_array":      []string{},
+		"blank":            []string{""},
+		"single":           []string{"a"},
+		"double":           []string{"a", "b"},
+		"duplicates":       []string{"a", "a", "b", "b"},
+		"duplicate_blanks": []string{"a", "a", "", ""},
+	})
+
+	var tests = []struct {
+		name     string
+		expected []string
+		defaults []string
+	}{
+		{
+			name:     "does_not_exist",
+			expected: []string(nil),
+		},
+		{
+			name:     "does_not_exist_with_defaults",
+			expected: []string{"a", "b"},
+			defaults: []string{"a", "b"},
+		},
+		{
+			name:     "empty_array",
+			expected: []string{"a", "b"},
+			defaults: []string{"a", "b"},
+		},
+		{
+			name:     "empty_array",
+			expected: []string(nil),
+		},
+		{
+			name:     "blank",
+			expected: []string{""},
+		},
+		{
+			name:     "blank",
+			expected: []string{""},
+			defaults: []string{"a", "b"},
+		},
+		{
+			name:     "single",
+			expected: []string{"a"},
+		},
+		{
+			name:     "single",
+			expected: []string{"a"},
+			defaults: []string{"a", "b"},
+		},
+		{
+			name:     "double",
+			expected: []string{"a", "b"},
+		},
+		{
+			name:     "duplicates",
+			expected: []string{"a", "b"},
+		},
+		{
+			name:     "duplicate_blanks",
+			expected: []string{"a", ""},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			actual := GetConstraints(mockQC, tt.name, tt.defaults...)
+			require.Equal(t, tt.expected, actual)
+		})
+	}
+
+}

pkg/osquery/tables/wmitable/helpers.go

@@ -4,9 +4,11 @@ import "strings"
 
 const allowedCharacters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-"
 
-func onlyAllowedCharacters(input string) bool {
+func onlyAllowedCharacters(input string, extras ...string) bool {
+
+	allowed := strings.Join(append(extras, allowedCharacters), "")
 	for _, char := range input {
-		if !strings.ContainsRune(allowedCharacters, char) {
+		if !strings.ContainsRune(allowed, char) {
 			return false
 		}
 	}

pkg/osquery/tables/wmitable/wmitable.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/kolide/launcher/pkg/dataflatten"
+	"github.com/kolide/launcher/pkg/osquery/tables/tablehelpers"
 	"github.com/kolide/launcher/pkg/wmi"
 
 	"github.com/go-kit/kit/log"
@@ -33,6 +34,7 @@ func TablePlugin(client *osquery.ExtensionManagerClient, logger log.Logger) *tab
 		table.TextColumn("query"),
 
 		// wmi columns
+		table.TextColumn("namespace"),
 		table.TextColumn("class"),
 		table.TextColumn("properties"),
 	}
@@ -49,65 +51,72 @@ func TablePlugin(client *osquery.ExtensionManagerClient, logger log.Logger) *tab
 func (t *Table) generate(ctx context.Context, queryContext table.QueryContext) ([]map[string]string, error) {
 	var results []map[string]string
 
-	classQ, ok := queryContext.Constraints["class"]
-	if !ok || len(classQ.Constraints) == 0 {
+	classes := tablehelpers.GetConstraints(queryContext, "class")
+	for _, c := range classes {
+		if !onlyAllowedCharacters(c) {
+			return nil, errors.New("Disallowed character in class expression")
+		}
+	}
+	if len(classes) == 0 {
 		return nil, errors.New("The kolide_wmi table requires a wmi class")
 	}
 
-	propertiesQ, ok := queryContext.Constraints["properties"]
-	if !ok || len(propertiesQ.Constraints) == 0 {
+	properties := tablehelpers.GetConstraints(queryContext, "properties")
+	for _, p := range properties {
+		// Allow comma, since we're going to split on it later.
+		if !onlyAllowedCharacters(p, `,`) {
+			return nil, errors.New("Disallowed character in properties expression")
+		}
+	}
+	if len(properties) == 0 {
 		return nil, errors.New("The kolide_wmi table requires wmi properties")
 	}
 
-	for _, classConstraint := range classQ.Constraints {
-		if !onlyAllowedCharacters(classConstraint.Expression) {
-			level.Info(t.logger).Log("msg", "Disallowed character in class expression")
-			return nil, errors.New("Disallowed character in class expression")
+	// Get the list of namespaces to query. If not specified, that's
+	// okay too, default to ""
+	namespaces := tablehelpers.GetConstraints(queryContext, "namespace", "")
+	for _, ns := range namespaces {
+		if !onlyAllowedCharacters(ns, `\`) {
+			return nil, errors.New("Disallowed character in namespace expression")
 		}
+	}
 
-		for _, propertiesConstraint := range propertiesQ.Constraints {
-
-			if !onlyAllowedCharacters(strings.Replace(propertiesConstraint.Expression, ",", "", -1)) {
-				level.Info(t.logger).Log("msg", "Disallowed character in properties expression")
-				return nil, errors.New("Disallowed character in properties expression")
-			}
+	flattenQueries := tablehelpers.GetConstraints(queryContext, "query", "")
 
-			properties := strings.Split(propertiesConstraint.Expression, ",")
+	for _, class := range classes {
+		for _, rawProperties := range properties {
+			properties := strings.Split(rawProperties, ",")
 			if len(properties) == 0 {
 				continue
 			}
+			for _, ns := range namespaces {
+				// Set a timeout in case wmi hangs
+				ctx, cancel := context.WithTimeout(ctx, 120*time.Second)
+				defer cancel()
+
+				wmiResults, err := wmi.Query(ctx, class, properties, wmi.ConnectUseMaxWait(), wmi.ConnectNamespace(ns))
+				if err != nil {
+					level.Info(t.logger).Log(
+						"msg", "wmi query failure",
+						"err", err,
+						"class", class,
+						"properties", rawProperties,
+						"namespace", ns,
+					)
+					continue
+				}
 
-			// Set a timeout in case wmi hangs
-			ctx, cancel := context.WithTimeout(ctx, 120*time.Second)
-			defer cancel()
-
-			wmiResults, err := wmi.Query(ctx, classConstraint.Expression, properties)
-			if err != nil {
-				level.Info(t.logger).Log(
-					"msg", "wmi query failure",
-					"err", err,
-					"class", classConstraint.Expression,
-					"properties", propertiesConstraint.Expression,
-				)
-				continue
-			}
-
-			if q, ok := queryContext.Constraints["query"]; ok && len(q.Constraints) != 0 {
-				for _, constraint := range q.Constraints {
-					dataQuery := constraint.Expression
-					results = append(results, t.flattenRowsFromWmi(dataQuery, wmiResults, classConstraint.Expression, propertiesConstraint.Expression)...)
+				for _, dataQuery := range flattenQueries {
+					results = append(results, t.flattenRowsFromWmi(dataQuery, wmiResults, class, rawProperties, ns)...)
 				}
-			} else {
-				results = append(results, t.flattenRowsFromWmi("", wmiResults, classConstraint.Expression, propertiesConstraint.Expression)...)
 			}
-
 		}
 	}
 
 	return results, nil
 }
 
-func (t *Table) flattenRowsFromWmi(dataQuery string, wmiResults []map[string]interface{}, wmiClass, wmiProperties string) []map[string]string {
+func (t *Table) flattenRowsFromWmi(dataQuery string, wmiResults []map[string]interface{}, wmiClass, wmiProperties, wmiNamespace string) []map[string]string {
 	flattenOpts := []dataflatten.FlattenOpts{}
 
 	if dataQuery != "" {
@@ -144,6 +153,7 @@ func (t *Table) flattenRowsFromWmi(dataQuery string, wmiResults []map[string]int
 			"query":      dataQuery,
 			"class":      wmiClass,
 			"properties": wmiProperties,
+			"namespace":  wmiNamespace,
 		}
 		results = append(results, res)
 	}

pkg/osquery/tables/wmitable/wmitable_test.go

@@ -21,7 +21,9 @@ func TestQueries(t *testing.T) {
 		name       string
 		class      string
 		properties []string
+		namespace  string
 		minRows    int
+		noData     bool
 		err        bool
 	}{
 		{
@@ -61,13 +63,30 @@ func TestQueries(t *testing.T) {
 			properties: []string{"name,ver;sion"},
 			err:        true,
 		},
+
+		{
+			name:       "bad namespace",
+			class:      "Win32_OperatingSystem",
+			properties: []string{"name,version"},
+			namespace:  `root\unknown\namespace`,
+			noData:     true,
+		},
+
+		{
+			name:       "different namespace",
+			class:      "MSKeyboard_PortInformation",
+			properties: []string{"ConnectorType,FunctionKeys,Indicators"},
+			namespace:  `root\wmi`,
+			minRows:    3,
+		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			mockQC := tablehelpers.MockQueryContext(map[string][]string{
 				"class":      []string{tt.class},
 				"properties": tt.properties,
+				"namespace":  []string{tt.namespace},
 			})
 
 			rows, err := wmiTable.generate(context.TODO(), mockQC)
@@ -79,6 +98,11 @@ func TestQueries(t *testing.T) {
 
 			require.NoError(t, err)
 
+			if tt.noData {
+				assert.Empty(t, rows, "Expected no results")
+				return
+			}
+
 			// It's hard to model what we expect to get
 			// from a random windows machine. So let's
 			// just check for non-empty data.