Add where support to kolide_wmi table (#658)

Add where support to wmi. This is implemented via a `whereclause` table
kolide · Sep 17, 2020 · b5be20d · b5be20d
1 parent a3ba681
commit b5be20d
Show file tree

Hide file tree

Showing 4 changed files with 107 additions and 40 deletions.
diff --git a/pkg/osquery/tables/wmitable/wmitable.go b/pkg/osquery/tables/wmitable/wmitable.go
@@ -7,6 +7,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/kolide/launcher/pkg/contexts/ctxlog"
 	"github.com/kolide/launcher/pkg/dataflatten"
 	"github.com/kolide/launcher/pkg/osquery/tables/tablehelpers"
 	"github.com/kolide/launcher/pkg/wmi"
@@ -39,11 +40,12 @@ func TablePlugin(client *osquery.ExtensionManagerClient, logger log.Logger) *tab
 		table.TextColumn("namespace"),
 		table.TextColumn("class"),
 		table.TextColumn("properties"),
+		table.TextColumn("whereclause"),
 	}
 
 	t := &Table{
 		client: client,
-		logger: level.NewFilter(logger, level.AllowInfo()),
+		logger: level.NewFilter(logger),
 	}
 
 	return table.NewPlugin("kolide_wmi", columns, t.generate)
@@ -70,6 +72,12 @@ func (t *Table) generate(ctx context.Context, queryContext table.QueryContext) (
 		tablehelpers.WithAllowedCharacters(allowedCharacters+`\`),
 	)
 
+	// Any whereclauses? These are not required
+	whereClauses := tablehelpers.GetConstraints(queryContext, "whereclause",
+		tablehelpers.WithDefaults(""),
+		tablehelpers.WithAllowedCharacters(allowedCharacters+`:\= '".`),
+	)
+
 	flattenQueries := tablehelpers.GetConstraints(queryContext, "query", tablehelpers.WithDefaults(""))
 
 	for _, class := range classes {
@@ -79,24 +87,30 @@ func (t *Table) generate(ctx context.Context, queryContext table.QueryContext) (
 				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
-				}
-
-				for _, dataQuery := range flattenQueries {
-					results = append(results, t.flattenRowsFromWmi(dataQuery, wmiResults, class, rawProperties, ns)...)
+				for _, whereClause := range whereClauses {
+					// Set a timeout in case wmi hangs
+					ctx, cancel := context.WithTimeout(ctx, 120*time.Second)
+					defer cancel()
+
+					// Add a logger in
+					ctx = ctxlog.NewContext(ctx, t.logger)
+
+					wmiResults, err := wmi.Query(ctx, class, properties, wmi.ConnectUseMaxWait(), wmi.ConnectNamespace(ns), wmi.WithWhere(whereClause))
+					if err != nil {
+						level.Info(t.logger).Log(
+							"msg", "wmi query failure",
+							"err", err,
+							"class", class,
+							"properties", rawProperties,
+							"namespace", ns,
+							"whereClause", whereClause,
+						)
+						continue
+					}
+
+					for _, dataQuery := range flattenQueries {
+						results = append(results, t.flattenRowsFromWmi(dataQuery, wmiResults, class, rawProperties, ns, whereClause)...)
+					}
 				}
 			}
 		}
@@ -105,15 +119,15 @@ func (t *Table) generate(ctx context.Context, queryContext table.QueryContext) (
 	return results, nil
 }
 
-func (t *Table) flattenRowsFromWmi(dataQuery string, wmiResults []map[string]interface{}, wmiClass, wmiProperties, wmiNamespace string) []map[string]string {
+func (t *Table) flattenRowsFromWmi(dataQuery string, wmiResults []map[string]interface{}, wmiClass, wmiProperties, wmiNamespace, whereClause string) []map[string]string {
 	flattenOpts := []dataflatten.FlattenOpts{}
 
 	if dataQuery != "" {
 		flattenOpts = append(flattenOpts, dataflatten.WithQuery(strings.Split(dataQuery, "/")))
 	}
 
 	if t.logger != nil {
-		flattenOpts = append(flattenOpts, dataflatten.WithLogger(t.logger))
+		flattenOpts = append(flattenOpts, dataflatten.WithLogger(level.NewFilter(t.logger, level.AllowInfo())))
 	}
 
 	// wmi.Query returns []map[string]interface{}, but dataflatten
@@ -135,14 +149,15 @@ func (t *Table) flattenRowsFromWmi(dataQuery string, wmiResults []map[string]int
 		p, k := row.ParentKey("/")
 
 		res := map[string]string{
-			"fullkey":    row.StringPath("/"),
-			"parent":     p,
-			"key":        k,
-			"value":      row.Value,
-			"query":      dataQuery,
-			"class":      wmiClass,
-			"properties": wmiProperties,
-			"namespace":  wmiNamespace,
+			"fullkey":     row.StringPath("/"),
+			"parent":      p,
+			"key":         k,
+			"value":       row.Value,
+			"query":       dataQuery,
+			"class":       wmiClass,
+			"properties":  wmiProperties,
+			"namespace":   wmiNamespace,
+			"whereclause": whereClause,
 		}
 		results = append(results, res)
 	}

diff --git a/pkg/osquery/tables/wmitable/wmitable_test.go b/pkg/osquery/tables/wmitable/wmitable_test.go
@@ -18,13 +18,14 @@ func TestQueries(t *testing.T) {
 	wmiTable := Table{logger: log.NewNopLogger()}
 
 	var tests = []struct {
-		name       string
-		class      string
-		properties []string
-		namespace  string
-		minRows    int
-		noData     bool
-		err        bool
+		name        string
+		class       string
+		properties  []string
+		namespace   string
+		whereClause string
+		minRows     int
+		noData      bool
+		err         bool
 	}{
 		{
 			name:       "simple operating system query",
@@ -79,14 +80,29 @@ func TestQueries(t *testing.T) {
 			namespace:  `root\wmi`,
 			minRows:    3,
 		},
+		{
+			name:        "where clause non-existent file",
+			class:       "CIM_DataFile",
+			properties:  []string{"name", "hidden"},
+			whereClause: `name = 'c:\\does\\not\\exist'`,
+			noData:      true,
+		},
+		{
+			name:        "where clause",
+			class:       "CIM_DataFile",
+			properties:  []string{"name", "hidden"},
+			whereClause: `name = 'c:\\windows\\system32\\notepad.exe'`,
+			minRows:     1,
+		},
 	}
 
 	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},
+				"class":       []string{tt.class},
+				"properties":  tt.properties,
+				"namespace":   []string{tt.namespace},
+				"whereclause": []string{tt.whereClause},
 			})
 
 			rows, err := wmiTable.generate(context.TODO(), mockQC)

diff --git a/pkg/wmi/wmi.go b/pkg/wmi/wmi.go
@@ -58,6 +58,7 @@ type querySettings struct {
 	connectLocale        string
 	connectAuthority     string
 	connectSecurityFlags uint
+	whereClause          string
 }
 
 // ConnectServerArgs returns an array suitable for being passed to ole
@@ -103,6 +104,13 @@ func ConnectUseMaxWait() Option {
 	}
 }
 
+// WithWhere will be used for the optional WHERE clause in wmi.
+func WithWhere(whereClause string) Option {
+	return func(qs *querySettings) {
+		qs.whereClause = whereClause
+	}
+}
+
 func Query(ctx context.Context, className string, properties []string, opts ...Option) ([]map[string]interface{}, error) {
 	logger := log.With(ctxlog.FromContext(ctx), "caller", "wmi.Query")
 	handler := NewOleHandler(ctx, properties)
@@ -113,11 +121,16 @@ func Query(ctx context.Context, className string, properties []string, opts ...O
 		opt(qs)
 	}
 
+	var whereClause string
+	if qs.whereClause != "" {
+		whereClause = fmt.Sprintf(" WHERE %s", qs.whereClause)
+	}
+
 	// If we query for the exact fields, _and_ one of the property
 	// names is wrong, we get no results. (clearly an error. but I
 	// can't find it) So query for `*`, and then fetch the
 	// property. More testing might show this needs to change
-	queryString := fmt.Sprintf("SELECT * FROM %s", className)
+	queryString := fmt.Sprintf("SELECT * FROM %s%s", className, whereClause)
 
 	// Initialize the COM system.
 	if err := ole.CoInitializeEx(0, ole.COINIT_MULTITHREADED); err != nil {
@@ -151,6 +164,8 @@ func Query(ctx context.Context, className string, properties []string, opts ...O
 	service := serviceRaw.ToIDispatch()
 	defer service.Release()
 
+	level.Debug(logger).Log("msg", "Running WMI query", "query", queryString)
+
 	// result is a SWBemObjectSet
 	resultRaw, err := oleutil.CallMethod(service, "ExecQuery", queryString)
 	if err != nil {

diff --git a/pkg/wmi/wmi_test.go b/pkg/wmi/wmi_test.go
@@ -96,6 +96,27 @@ func TestQuery(t *testing.T) {
 			options:    []Option{ConnectNamespace(`root\wmi`)},
 			minRows:    1,
 		},
+		{
+			name:       "where clause",
+			class:      "CIM_DataFile",
+			properties: []string{"name", "hidden"},
+			options:    []Option{WithWhere(`name = 'c:\\windows\\system32\\notepad.exe'`)},
+			minRows:    1,
+		},
+		{
+			name:       "where clause non-existent file",
+			class:      "CIM_DataFile",
+			properties: []string{"name", "hidden"},
+			options:    []Option{WithWhere(`name = 'c:\\does\\not\\exist'`)},
+			noData:     true,
+		},
+		{
+			name:       "where clause bad argument",
+			class:      "CIM_DataFile",
+			properties: []string{"name", "hidden"},
+			options:    []Option{WithWhere(`nameNope = 'c:\\does\\not\\exist'`)},
+			noData:     true,
+		},
 	}
 
 	for _, tt := range tests {