action_handler.go

// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
// or more contributor license agreements. Licensed under the Elastic License;
// you may not use this file except in compliance with the Elastic License.

package beater

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/elastic/beats/v7/x-pack/osquerybeat/internal/action"
	"github.com/elastic/beats/v7/x-pack/osquerybeat/internal/config"
	"github.com/elastic/beats/v7/x-pack/osquerybeat/internal/ecs"
	"github.com/elastic/elastic-agent-libs/logp"
)

var (
	ErrNoPublisher     = errors.New("no publisher configured")
	ErrNoQueryExecutor = errors.New("no query executor configures")
)

type actionResultPublisher interface {
	PublishActionResult(req map[string]interface{}, res map[string]interface{})
}

type queryResultPublisher interface {
	Publish(index, idValue, idFieldKey, responseID, spaceID, packID, packName, queryName string, meta map[string]interface{}, hits []map[string]interface{}, ecsm ecs.Mapping, reqData interface{})
}

type scheduledResponsePublisher interface {
	PublishScheduledResponse(scheduleID, packID, packName, queryName, spaceID, responseID string, startedAt, completedAt, plannedScheduleTime time.Time, resultCount int, scheduleExecutionCount int64)
}

type queryProfilePublisher interface {
	PublishQueryProfile(index, queryName, actionID, responseID string, profile map[string]interface{}, reqData interface{})
}

type liveProfileRecorder interface {
	RecordLiveProfile(query string, profile map[string]interface{})
}

type scheduledQueryPublisher interface {
	queryResultPublisher
	scheduledResponsePublisher
	queryProfilePublisher
}

type actionQueryPublisher interface {
	queryResultPublisher
	queryProfilePublisher
}

type queryExecutor interface {
	Query(ctx context.Context, sql string, timeout time.Duration) ([]map[string]interface{}, error)
}

type namespaceProvider interface {
	GetNamespace() string
}

type actionHandler struct {
	log       *logp.Logger
	inputType string
	publisher actionQueryPublisher
	queryExec queryExecutor
	np        namespaceProvider
	profiles  liveProfileRecorder
}

func (a *actionHandler) Name() string {
	return a.inputType
}

// Execute handles the action request.
func (a *actionHandler) Execute(ctx context.Context, req map[string]interface{}) (map[string]interface{}, error) {

	start := time.Now().UTC()
	count, err := a.execute(ctx, req)
	end := time.Now().UTC()

	res := map[string]interface{}{
		"started_at":   start.Format(time.RFC3339Nano),
		"completed_at": end.Format(time.RFC3339Nano),
	}

	if err != nil {
		res["error"] = err.Error()
	} else {
		res["count"] = count
	}
	return res, nil
}

func (a *actionHandler) execute(ctx context.Context, req map[string]interface{}) (int, error) {
	ac, err := action.FromMap(req)
	if err != nil {
		return 0, fmt.Errorf("%w: %w", err, ErrQueryExecution)
	}
	if !ac.MatchesPlatform() {
		a.log.Debugf("Skipping query for platforms: %v", ac.Platforms)
		return 0, nil
	}

	return a.executeQuery(ctx, config.Datastream(a.namespace()), ac, "", req)
}

func (a *actionHandler) namespace() string {
	if a.np != nil {
		if ns := a.np.GetNamespace(); ns != "" {
			return ns
		}
	}
	return config.DefaultNamespace
}

func (a *actionHandler) executeQuery(ctx context.Context, index string, ac action.Action, responseID string, req map[string]interface{}) (int, error) {

	if a.queryExec == nil {
		return 0, ErrNoQueryExecutor
	}
	if a.publisher == nil {
		return 0, ErrNoPublisher
	}

	a.log.Debugf("Execute query: %s", ac.Query)

	var before runtimeSnapshot
	beforeReady := false
	shouldCollect := ac.Profile || a.profiles != nil
	if shouldCollect {
		snapshot, err := collectRuntimeSnapshot(ctx, a.queryExec)
		if err != nil {
			a.log.Debugf("failed to collect pre-query profile snapshot: %v", err)
		} else {
			before = snapshot
			beforeReady = true
		}
	}

	start := time.Now()

	hits, err := a.queryExec.Query(ctx, ac.Query, ac.Timeout)
	duration := time.Since(start)

	if shouldCollect && beforeReady {
		after, snapErr := collectRuntimeSnapshot(ctx, a.queryExec)
		if snapErr != nil {
			a.log.Debugf("failed to collect post-query profile snapshot: %v", snapErr)
		} else {
			profile := buildLiveQueryProfile(ac.Query, before, after, duration, err)
			if a.profiles != nil {
				a.profiles.RecordLiveProfile(ac.Query, profile)
			}
			if ac.Profile {
				a.publisher.PublishQueryProfile(config.QueryProfileDatastream(a.namespace()), "", ac.ID, responseID, profile, req["data"])
			}
		}
	} else if shouldCollect && !beforeReady {
		if ac.Profile {
			a.log.Debug("profile requested but skipped: pre-query snapshot was not collected")
		} else {
			a.log.Debug("profile storage skipped: pre-query snapshot was not collected")
		}
	}

	if err != nil {
		a.log.Errorf("Failed to execute query, err: %v", err)
		return 0, err
	}

	a.log.Debugf("Completed query in: %v", duration)

	a.publisher.Publish(index, ac.ID, "action_id", responseID, "", "", "", "", nil, hits, ac.ECSMapping, req["data"])

	return len(hits), nil
}