github.com/Uptycs/basequery-go@v0.8.0/plugin/table/table.go

// Package table creates an osquery table plugin.
package table

import (
	"context"
	"encoding/json"
	"strconv"

	"github.com/Uptycs/basequery-go/gen/osquery"
	"github.com/pkg/errors"
)

// GenerateFunc returns the rows generated by the table. The ctx argument
// should be checked for cancellation if the generation performs a
// substantial amount of work. The queryContext argument provides the
// deserialized JSON query context from osquery.
type GenerateFunc func(ctx context.Context, queryContext QueryContext) ([]map[string]string, error)

// InsertFunc is optional implementation that can be used to implement insert SQL semantics
type InsertFunc func(ctx context.Context, autoRowId bool, row []interface{}) ([]map[string]string, error)

// UpdateFunc is optional implementation that can be used to implement update SQL semantics
type UpdateFunc func(ctx context.Context, rowID int64, row []interface{}) error

// DeleteFunc is optional implementation that can be used to implement delete SQL semantics
type DeleteFunc func(ctx context.Context, rowID int64) error

// Plugin structure holds the plugin details.
type Plugin struct {
	name     string
	columns  []ColumnDefinition
	generate GenerateFunc
	insert   InsertFunc
	update   UpdateFunc
	delete   DeleteFunc
}

// NewPlugin is helper method to create plugin structure.
func NewPlugin(name string, columns []ColumnDefinition, gen GenerateFunc) *Plugin {
	return &Plugin{
		name:     name,
		columns:  columns,
		generate: gen,
	}
}

// NewMutablePlugin is helper method to create mutable plugin structure.
func NewMutablePlugin(name string, columns []ColumnDefinition, gen GenerateFunc, ins InsertFunc, upd UpdateFunc, del DeleteFunc) *Plugin {
	return &Plugin{
		name:     name,
		columns:  columns,
		generate: gen,
		insert:   ins,
		update:   upd,
		delete:   del,
	}
}

func createError(prefix string, err error) osquery.ExtensionResponse {
	msg := prefix
	if err != nil {
		msg += err.Error()
	}
	return osquery.ExtensionResponse{
		Status: &osquery.ExtensionStatus{
			Code:    1,
			Message: msg,
		},
	}
}

// Name returns the plugin name.
func (t *Plugin) Name() string {
	return t.name
}

// RegistryName returns the plugin type, which is always "table" for table plugin.
func (t *Plugin) RegistryName() string {
	return "table"
}

// Routes returns the table columns definitions.
func (t *Plugin) Routes() osquery.ExtensionPluginResponse {
	routes := []map[string]string{}
	for _, col := range t.columns {
		routes = append(routes, map[string]string{
			"id":   "column",
			"name": col.Name,
			"type": string(col.Type),
			"op":   strconv.Itoa(int(col.Op)),
		})
	}
	return routes
}

// Call is invoked to generate the table contents or to get the column details.
func (t *Plugin) Call(ctx context.Context, request osquery.ExtensionPluginRequest) osquery.ExtensionResponse {
	ok := osquery.ExtensionStatus{Code: 0, Message: "OK"}
	switch request["action"] {
	case "generate":
		queryContext, err := parseQueryContext(request["context"])
		if err != nil {
			return createError("error parsing context JSON: ", err)
		}

		rows, err := t.generate(ctx, *queryContext)
		if err != nil {
			return createError("error generating table: ", err)
		}

		return osquery.ExtensionResponse{Status: &ok, Response: rows}

	case "insert":
		if t.insert == nil {
			return createError("'insert' not implemented by table: "+t.name, nil)
		}

		row, err := parseRow(request["json_value_array"])
		if err != nil {
			return createError("invalid data to insert: ", err)
		}

		autoRowID, err := strconv.ParseBool(request["auto_rowid"])
		if err != nil {
			return createError("invalid value for auto_rowid: ", err)
		}

		rows, err := t.insert(ctx, autoRowID, row)
		if err != nil {
			return createError("error inserting into table: ", err)
		}

		return osquery.ExtensionResponse{Status: &ok, Response: rows}

	case "update":
		if t.update == nil {
			return createError("'update' not implemented by table: "+t.name, nil)
		}

		rowID, err := strconv.ParseInt(request["id"], 10, 64)
		if err != nil {
			return createError("invalid row id to update: ", err)
		}

		row, err := parseRow(request["json_value_array"])
		if err != nil {
			return createError("invalid data to update: ", err)
		}

		err = t.update(ctx, rowID, row)
		if err != nil {
			return createError("error updating table: ", err)
		}

		return osquery.ExtensionResponse{Status: &ok, Response: []map[string]string{{"status": "success"}}}

	case "delete":
		if t.delete == nil {
			return createError("'delete' not implemented by table: "+t.name, nil)
		}

		rowID, err := strconv.ParseInt(request["id"], 10, 64)
		if err != nil {
			return createError("invalid row id to delete: ", err)
		}

		err = t.delete(ctx, rowID)
		if err != nil {
			return createError("error deleting from table: ", err)
		}

		return osquery.ExtensionResponse{Status: &ok, Response: []map[string]string{{"status": "success"}}}

	case "columns":
		return osquery.ExtensionResponse{Status: &ok, Response: t.Routes()}

	default:
		return createError("unknown action: "+request["action"], nil)
	}

}

// Ping returns static OK response.
func (t *Plugin) Ping() osquery.ExtensionStatus {
	return osquery.ExtensionStatus{Code: 0, Message: "OK"}
}

// Shutdown is a no-op for table plugins.
func (t *Plugin) Shutdown() {}

// ColumnDefinition defines the relevant information for a column in a table
// plugin. Both values are mandatory. Prefer using the *Column helpers to
// create ColumnDefinition structs.
type ColumnDefinition struct {
	Name string
	Type ColumnType
	Op   ColumnOptions
}

// TextColumn is a helper for defining columns containing strings.
func TextColumn(name string, options ...ColumnOptions) ColumnDefinition {
	return ColumnDefinition{
		Name: name,
		Type: ColumnTypeText,
		Op:   getColumnOption(options...),
	}
}

// IntegerColumn is a helper for defining columns containing integers.
func IntegerColumn(name string, options ...ColumnOptions) ColumnDefinition {
	return ColumnDefinition{
		Name: name,
		Type: ColumnTypeInteger,
		Op:   getColumnOption(options...),
	}
}

// BigIntColumn is a helper for defining columns containing big integers.
func BigIntColumn(name string, options ...ColumnOptions) ColumnDefinition {
	return ColumnDefinition{
		Name: name,
		Type: ColumnTypeBigInt,
		Op:   getColumnOption(options...),
	}
}

// DoubleColumn is a helper for defining columns containing floating point
// values.
func DoubleColumn(name string, options ...ColumnOptions) ColumnDefinition {
	return ColumnDefinition{
		Name: name,
		Type: ColumnTypeDouble,
		Op:   getColumnOption(options...),
	}
}

func getColumnOption(options ...ColumnOptions) ColumnOptions {
	op := DEFAULT
	if len(options) > 0 {
		for _, option := range options {
			op |= option
		}
	}
	return op
}

// ColumnType is a strongly typed representation of the data type string for a
// column definition. The named constants should be used.
type ColumnType string

// The following column types are defined in osquery tables.h.
const (
	ColumnTypeText    ColumnType = "TEXT"
	ColumnTypeInteger            = "INTEGER"
	ColumnTypeBigInt             = "BIGINT"
	ColumnTypeDouble             = "DOUBLE"
)

// ColumnOptions for marking columns
type ColumnOptions int

const (
	// DEFAULT means no special column options.
	DEFAULT ColumnOptions = iota // 0

	// INDEX treats this column as a primary key.
	INDEX = iota // 1

	// REQUIRED column MUST be included in the query predicate.
	REQUIRED = iota // 2

	// ADDITIONAL column is used to generate additional information.
	ADDITIONAL = iota + 1 // 3 + 1

	// OPTIMIZED column can be used to optimize the query.
	OPTIMIZED = iota + 4 // 4 + 4

	// HIDDEN column should be hidden from '*'' selects.
	HIDDEN = iota + 11 // 5 + 11
)

// QueryContext contains the constraints from the WHERE clause of the query,
// that can optionally be used to optimize the table generation. Note that the
// osquery SQLite engine will perform the filtering with these constraints, so
// it is not mandatory that they be used in table generation.
type QueryContext struct {
	// Constraints is a map from column name to the details of the
	// constraints on that column.
	Constraints map[string]ConstraintList
}

// ConstraintList contains the details of the constraints for the given column.
type ConstraintList struct {
	Affinity    ColumnType
	Constraints []Constraint
}

// Constraint contains both an operator and an expression that are applied as
// constraints in the query.
type Constraint struct {
	Operator   Operator
	Expression string
}

// Operator is an enum of the osquery operators.
type Operator int

// The following operators are dfined in osquery tables.h.
const (
	OperatorEquals              Operator = 2
	OperatorGreaterThan                  = 4
	OperatorLessThanOrEquals             = 8
	OperatorLessThan                     = 16
	OperatorGreaterThanOrEquals          = 32
	OperatorMatch                        = 64
	OperatorLike                         = 65
	OperatorGlob                         = 66
	OperatorRegexp                       = 67
	OperatorUnique                       = 1
)

// The following types and functions exist for parsing of the queryContext
// JSON and are not made public.
type queryContextJSON struct {
	Constraints []constraintListJSON `json:"constraints"`
}

type constraintListJSON struct {
	Name     string          `json:"name"`
	Affinity string          `json:"affinity"`
	List     json.RawMessage `json:"list"`
}

func parseQueryContext(ctxJSON string) (*QueryContext, error) {
	ctx := QueryContext{map[string]ConstraintList{}}
	if ctxJSON == "" {
		return &ctx, nil
	}

	var parsed queryContextJSON
	err := json.Unmarshal([]byte(ctxJSON), &parsed)
	if err != nil {
		return nil, errors.Wrap(err, "unmarshaling context JSON")
	}

	for _, cList := range parsed.Constraints {
		constraints, err := parseConstraintList(cList.List)
		if err != nil {
			return nil, err
		}

		ctx.Constraints[cList.Name] = ConstraintList{
			Affinity:    ColumnType(cList.Affinity),
			Constraints: constraints,
		}
	}

	return &ctx, nil
}

func parseRow(row string) ([]interface{}, error) {
	if row == "" {
		return nil, errors.Errorf("invalid data to insert")
	}

	var parsed []interface{}
	err := json.Unmarshal([]byte(row), &parsed)
	if err != nil {
		return nil, errors.Wrap(err, "unmarshaling JSON")
	}

	return parsed, nil
}

func parseConstraintList(constraints json.RawMessage) ([]Constraint, error) {
	var str string
	err := json.Unmarshal(constraints, &str)
	if err == nil {
		// string indicates empty list
		return []Constraint{}, nil
	}

	var cList []map[string]interface{}
	err = json.Unmarshal(constraints, &cList)
	if err != nil {
		// cannot do anything with other types
		return nil, errors.Errorf("unexpected context list: %s", string(constraints))
	}

	cl := []Constraint{}
	for _, c := range cList {
		var op Operator
		switch opVal := c["op"].(type) {
		case string: // osquery < 3.0 with stringy types
			opInt, err := strconv.Atoi(opVal)
			if err != nil {
				return nil, errors.Errorf("parsing operator int: %s", c["op"])
			}
			op = Operator(opInt)
		case float64: // osquery > 3.0 with strong types
			op = Operator(opVal)
		default:
			return nil, errors.Errorf("cannot parse type %T", opVal)
		}

		expr, ok := c["expr"].(string)
		if !ok {
			return nil, errors.Errorf("expr should be string: %s", c["expr"])
		}

		cl = append(cl, Constraint{
			Operator:   op,
			Expression: expr,
		})
	}
	return cl, nil
}