fleet/server/service/queries.go

package service

import (
	"context"
	"fmt"
	"slices"
	"strings"

	"github.com/fleetdm/fleet/v4/server/authz"
	"github.com/fleetdm/fleet/v4/server/contexts/ctxerr"
	"github.com/fleetdm/fleet/v4/server/contexts/license"
	"github.com/fleetdm/fleet/v4/server/contexts/logging"
	"github.com/fleetdm/fleet/v4/server/contexts/viewer"
	"github.com/fleetdm/fleet/v4/server/fleet"
	"github.com/fleetdm/fleet/v4/server/ptr"
)

////////////////////////////////////////////////////////////////////////////////
// Get Query
////////////////////////////////////////////////////////////////////////////////

func getQueryEndpoint(ctx context.Context, request interface{}, svc fleet.Service) (fleet.Errorer, error) {
	req := request.(*fleet.GetQueryRequest)
	query, err := svc.GetQuery(ctx, req.ID)
	if err != nil {
		return fleet.GetQueryResponse{Err: err}, nil
	}
	return fleet.GetQueryResponse{Query: query, Report: query}, nil
}

func (svc *Service) GetQuery(ctx context.Context, id uint) (*fleet.Query, error) {
	// Load query first to get its teamID.
	query, err := svc.ds.Query(ctx, id)
	if err != nil {
		setAuthCheckedOnPreAuthErr(ctx)
		return nil, ctxerr.Wrap(ctx, err, "get query from datastore")
	}
	if err := svc.authz.Authorize(ctx, query, fleet.ActionRead); err != nil {
		return nil, err
	}
	return query, nil
}

////////////////////////////////////////////////////////////////////////////////
// List Queries
////////////////////////////////////////////////////////////////////////////////

func listQueriesEndpoint(ctx context.Context, request interface{}, svc fleet.Service) (fleet.Errorer, error) {
	req := request.(*fleet.ListQueriesRequest)

	var teamID *uint
	if req.TeamID != 0 {
		teamID = &req.TeamID
	}

	var urlPlatform *string
	if req.Platform != "" {
		urlPlatform = &req.Platform
	}

	queries, count, inheritedCount, meta, err := svc.ListQueries(ctx, req.ListOptions, teamID, nil, req.MergeInherited, urlPlatform)
	if err != nil {
		return fleet.ListQueriesResponse{Err: err}, nil
	}

	respQueries := make([]fleet.Query, 0, len(queries))
	for _, query := range queries {
		respQueries = append(respQueries, *query)
	}

	return fleet.ListQueriesResponse{
		Queries:             respQueries,
		Count:               count,
		InheritedQueryCount: inheritedCount,
		Meta:                meta,
	}, nil
}

func (svc *Service) ListQueries(ctx context.Context, opt fleet.ListOptions, teamID *uint, scheduled *bool, mergeInherited bool, urlPlatform *string) ([]*fleet.Query, int, int, *fleet.PaginationMetadata, error) {
	// Check the user is allowed to list queries on the given team.
	if err := svc.authz.Authorize(ctx, &fleet.Query{
		TeamID: teamID,
	}, fleet.ActionRead); err != nil {
		return nil, 0, 0, nil, err
	}

	// always include metadata for queries
	opt.IncludeMetadata = true

	var dbPlatform *string
	if urlPlatform != nil {
		// validate platform filter
		if *urlPlatform == "macos" {
			// More user-friendly API param "macos" is called "darwin" in the datastore
			dbPlatform = ptr.String("darwin")
		} else {
			dbPlatform = urlPlatform
		}
		if strings.Contains(*urlPlatform, ",") {
			return nil, 0, 0, nil, &fleet.BadRequestError{Message: "queries can only be filtered by one platform at a time"}
		}
		targetableDBPlatforms := []string{"darwin", "windows", "linux"}
		if !slices.Contains(targetableDBPlatforms, *dbPlatform) {
			return nil, 0, 0, nil, &fleet.BadRequestError{Message: fmt.Sprintf("platform %q cannot be a scheduled query target, supported platforms are: %s", *dbPlatform, strings.Join(targetableDBPlatforms, ","))}
		}
	}

	queries, count, inheritedCount, meta, err := svc.ds.ListQueries(ctx, fleet.ListQueryOptions{
		ListOptions:    opt,
		TeamID:         teamID,
		IsScheduled:    scheduled,
		MergeInherited: mergeInherited,
		Platform:       dbPlatform,
	})
	if err != nil {
		return nil, 0, 0, nil, err
	}

	return queries, count, inheritedCount, meta, nil
}

////////////////////////////////////////////////////////////////////////////////
// Query Reports
////////////////////////////////////////////////////////////////////////////////

func getQueryReportEndpoint(ctx context.Context, request interface{}, svc fleet.Service) (fleet.Errorer, error) {
	req := request.(*fleet.GetQueryReportRequest)
	queryReportResults, reportClipped, err := svc.GetQueryReportResults(ctx, req.ID, req.TeamID)
	if err != nil {
		return fleet.GetQueryReportResponse{Err: err}, nil
	}
	// Return an empty array if there are no results stored.
	results := []fleet.HostQueryResultRow{}
	if len(queryReportResults) > 0 {
		results = queryReportResults
	}
	return fleet.GetQueryReportResponse{
		QueryID:       req.ID,
		Results:       results,
		ReportClipped: reportClipped,
	}, nil
}

func (svc *Service) GetQueryReportResults(ctx context.Context, id uint, teamID *uint) ([]fleet.HostQueryResultRow, bool, error) {
	// Load query first to get its teamID.
	query, err := svc.ds.Query(ctx, id)
	if err != nil {
		setAuthCheckedOnPreAuthErr(ctx)
		return nil, false, ctxerr.Wrap(ctx, err, "get query from datastore")
	}
	if err := svc.authz.Authorize(ctx, query, fleet.ActionRead); err != nil {
		return nil, false, err
	}

	if query.DiscardData {
		return nil, false, nil
	}

	vc, ok := viewer.FromContext(ctx)
	if !ok {
		return nil, false, fleet.ErrNoContext
	}
	filter := fleet.TeamFilter{User: vc.User, IncludeObserver: true, TeamID: teamID}

	queryReportResultRows, err := svc.ds.QueryResultRows(ctx, id, filter)
	if err != nil {
		return nil, false, ctxerr.Wrap(ctx, err, "get query report results")
	}
	queryReportResults, err := fleet.MapQueryReportResultsToRows(queryReportResultRows)
	if err != nil {
		return nil, false, ctxerr.Wrap(ctx, err, "map db rows to results")
	}
	appConfig, err := svc.ds.AppConfig(ctx)
	if err != nil {
		return nil, false, ctxerr.Wrap(ctx, err, "get app config")
	}
	reportClipped, err := svc.QueryReportIsClipped(ctx, id, appConfig.ServerSettings.GetQueryReportCap())
	if err != nil {
		return nil, false, ctxerr.Wrap(ctx, err, "check query report is clipped")
	}
	return queryReportResults, reportClipped, nil
}

func (svc *Service) QueryReportIsClipped(ctx context.Context, queryID uint, maxQueryReportRows int) (bool, error) {
	query, err := svc.ds.Query(ctx, queryID)
	if err != nil {
		setAuthCheckedOnPreAuthErr(ctx)
		return false, ctxerr.Wrap(ctx, err, "get query from datastore")
	}
	if err := svc.authz.Authorize(ctx, query, fleet.ActionRead); err != nil {
		return false, err
	}

	count, err := svc.ds.ResultCountForQuery(ctx, queryID)
	if err != nil {
		return false, err
	}
	return count >= maxQueryReportRows, nil
}

////////////////////////////////////////////////////////////////////////////////
// Create Query
////////////////////////////////////////////////////////////////////////////////

func createQueryEndpoint(ctx context.Context, request interface{}, svc fleet.Service) (fleet.Errorer, error) {
	req := request.(*fleet.CreateQueryRequest)
	query, err := svc.NewQuery(ctx, req.QueryPayload)
	if err != nil {
		return fleet.CreateQueryResponse{Err: err}, nil
	}
	return fleet.CreateQueryResponse{Query: query, Report: query}, nil
}

func (svc *Service) NewQuery(ctx context.Context, p fleet.QueryPayload) (*fleet.Query, error) {
	// Check the user is allowed to create a new query on the team.
	if err := svc.authz.Authorize(ctx, fleet.Query{
		TeamID: p.TeamID,
	}, fleet.ActionWrite); err != nil {
		return nil, err
	}

	if p.Logging == nil || (p.Logging != nil && *p.Logging == "") {
		p.Logging = ptr.String(fleet.LoggingSnapshot)
	}

	// Targeting queries by label is a premium feature only.
	if (len(p.LabelsIncludeAny) > 0 || len(p.LabelsIncludeAll) > 0) && !license.IsPremium(ctx) {
		return nil, fleet.ErrMissingLicense
	}

	if err := p.Verify(); err != nil {
		return nil, ctxerr.Wrap(ctx, &fleet.BadRequestError{
			Message: fmt.Sprintf("query payload verification: %s", err),
		})
	}

	query := &fleet.Query{Saved: true, TeamID: p.TeamID}

	vc, ok := viewer.FromContext(ctx)
	if ok {
		query.AuthorID = ptr.Uint(vc.UserID())
		query.AuthorName = vc.FullName()
		query.AuthorEmail = vc.Email()
	}

	if err := verifyLabelsToAssociate(ctx, svc.ds, p.TeamID, slices.Concat(p.LabelsIncludeAny, p.LabelsIncludeAll), vc.User); err != nil {
		return nil, ctxerr.Wrap(ctx, err, "verify labels to associate")
	}

	if p.Name != nil {
		query.Name = *p.Name
	}
	if p.Description != nil {
		query.Description = *p.Description
	}
	if p.Query != nil {
		query.Query = *p.Query
	}
	if p.Interval != nil {
		query.Interval = *p.Interval
	}
	if p.Platform != nil {
		query.Platform = *p.Platform
	}
	if p.MinOsqueryVersion != nil {
		query.MinOsqueryVersion = *p.MinOsqueryVersion
	}
	if p.AutomationsEnabled != nil {
		query.AutomationsEnabled = *p.AutomationsEnabled
	}
	if p.Logging != nil {
		query.Logging = *p.Logging
	}
	if p.ObserverCanRun != nil {
		query.ObserverCanRun = *p.ObserverCanRun
	}
	if p.DiscardData != nil {
		query.DiscardData = *p.DiscardData
	}
	if len(p.LabelsIncludeAny) > 0 {
		query.LabelsIncludeAny = fleet.LabelNamesToIdents(p.LabelsIncludeAny)
	}
	if len(p.LabelsIncludeAll) > 0 {
		query.LabelsIncludeAll = fleet.LabelNamesToIdents(p.LabelsIncludeAll)
	}

	logging.WithExtras(ctx, "name", query.Name, "sql", query.Query)

	query, err := svc.ds.NewQuery(ctx, query)
	if err != nil {
		return nil, err
	}

	var teamID int64
	var teamName *string
	if query.TeamID != nil {
		teamID = int64(*query.TeamID) //nolint:gosec // dismiss G115
		if svc.EnterpriseOverrides != nil && svc.EnterpriseOverrides.TeamByIDOrName != nil {
			team, err := svc.EnterpriseOverrides.TeamByIDOrName(ctx, query.TeamID, nil)
			if err != nil {
				return nil, err
			}
			teamName = &team.Name
		}
	} else {
		teamID = -1 // Use -1 for global queries
		teamName = nil
	}

	if err := svc.NewActivity(
		ctx,
		authz.UserFromContext(ctx),
		fleet.ActivityTypeCreatedSavedQuery{
			ID:       query.ID,
			Name:     query.Name,
			TeamID:   teamID, // log teamID if available, else -1 for global
			TeamName: teamName,
		},
	); err != nil {
		return nil, ctxerr.Wrap(ctx, err, "create activity for query creation")
	}

	return query, nil
}

////////////////////////////////////////////////////////////////////////////////
// Modify Query
////////////////////////////////////////////////////////////////////////////////

func modifyQueryEndpoint(ctx context.Context, request interface{}, svc fleet.Service) (fleet.Errorer, error) {
	req := request.(*fleet.ModifyQueryRequest)
	query, err := svc.ModifyQuery(ctx, req.ID, req.QueryPayload)
	if err != nil {
		return fleet.ModifyQueryResponse{Err: err}, nil
	}
	return fleet.ModifyQueryResponse{Query: query, Report: query}, nil
}

func (svc *Service) ModifyQuery(ctx context.Context, id uint, p fleet.QueryPayload) (*fleet.Query, error) {
	// Load query first to determine if the user can modify it.
	query, err := svc.ds.Query(ctx, id)
	if err != nil {
		setAuthCheckedOnPreAuthErr(ctx)
		return nil, err
	}
	if err := svc.authz.Authorize(ctx, query, fleet.ActionWrite); err != nil {
		return nil, err
	}

	if p.Logging != nil && *p.Logging == "" {
		p.Logging = ptr.String(fleet.LoggingSnapshot)
	}

	// Targeting queries by label is a premium feature only.
	if (len(p.LabelsIncludeAny) > 0 || len(p.LabelsIncludeAll) > 0) && !license.IsPremium(ctx) {
		return nil, fleet.ErrMissingLicense
	}

	if err := p.Verify(); err != nil {
		return nil, ctxerr.Wrap(ctx, &fleet.BadRequestError{
			Message: fmt.Sprintf("query payload verification: %s", err),
		})
	}

	// We use query.TeamID because we do not allow changing the team.
	if err := verifyLabelsToAssociate(ctx, svc.ds, query.TeamID, slices.Concat(p.LabelsIncludeAny, p.LabelsIncludeAll), authz.UserFromContext(ctx)); err != nil {
		return nil, ctxerr.Wrap(ctx, err, "verify labels to associate")
	}

	shouldDiscardQueryResults, shouldDeleteStats := false, false

	if p.Name != nil {
		query.Name = *p.Name
	}
	if p.Description != nil {
		query.Description = *p.Description
	}
	if p.Query != nil {
		if query.Query != *p.Query {
			shouldDiscardQueryResults = true
			shouldDeleteStats = true
		}
		query.Query = *p.Query
	}
	if p.Interval != nil {
		query.Interval = *p.Interval
	}
	if p.Platform != nil {
		if !comparePlatforms(query.Platform, *p.Platform) {
			shouldDiscardQueryResults = true
		}
		query.Platform = *p.Platform
	}
	if p.MinOsqueryVersion != nil {
		if query.MinOsqueryVersion != *p.MinOsqueryVersion {
			shouldDiscardQueryResults = true
		}
		query.MinOsqueryVersion = *p.MinOsqueryVersion
	}
	if p.AutomationsEnabled != nil {
		query.AutomationsEnabled = *p.AutomationsEnabled
	}
	if p.Logging != nil {
		if query.Logging != *p.Logging && *p.Logging != fleet.LoggingSnapshot {
			shouldDiscardQueryResults = true
		}
		query.Logging = *p.Logging
	}
	if p.ObserverCanRun != nil {
		query.ObserverCanRun = *p.ObserverCanRun
	}
	if p.DiscardData != nil {
		if *p.DiscardData && *p.DiscardData != query.DiscardData {
			shouldDiscardQueryResults = true
		}
		query.DiscardData = *p.DiscardData
	}
	// If either label scope field is non-nil, treat both as authoritative.
	// Mutual exclusion is enforced upstream so at most one slice is non-nil;
	// the other is reset to empty so a scope switch implicitly clears the
	// previous scope.
	if p.LabelsIncludeAny != nil || p.LabelsIncludeAll != nil {
		query.LabelsIncludeAny = fleet.LabelNamesToIdents(p.LabelsIncludeAny)
		query.LabelsIncludeAll = fleet.LabelNamesToIdents(p.LabelsIncludeAll)
	}

	logging.WithExtras(ctx, "name", query.Name, "sql", query.Query)

	if err := svc.ds.SaveQuery(ctx, query, shouldDiscardQueryResults, shouldDeleteStats); err != nil {
		return nil, err
	}

	// If the query was modified in a way that requires discarding results,
	// reset the Redis count as well.
	if shouldDiscardQueryResults && svc.liveQueryStore != nil {
		err = svc.liveQueryStore.SetQueryResultsCount(query.ID, 0)
		if err != nil {
			// Log the error but don't fail the request; this will get cleaned up
			// in the "query_results_cleanup" job.
			svc.logger.ErrorContext(ctx, "failed to set query results count", "err", err, "query_id", query.ID)
		}
	}

	var teamID int64
	var teamName *string
	if query.TeamID != nil {
		teamID = int64(*query.TeamID) //nolint:gosec // dismiss G115
		if svc.EnterpriseOverrides != nil && svc.EnterpriseOverrides.TeamByIDOrName != nil {
			team, err := svc.EnterpriseOverrides.TeamByIDOrName(ctx, query.TeamID, nil)
			if err != nil {
				return nil, err
			}
			teamName = &team.Name
		}
	} else {
		teamID = -1
		teamName = nil
	}

	if err := svc.NewActivity(
		ctx,
		authz.UserFromContext(ctx),
		fleet.ActivityTypeEditedSavedQuery{
			ID:       query.ID,
			Name:     query.Name,
			TeamID:   teamID,
			TeamName: teamName,
		},
	); err != nil {
		return nil, ctxerr.Wrap(ctx, err, "create activity for query modification")
	}

	return query, nil
}

func comparePlatforms(platform1, platform2 string) bool {
	if platform1 == platform2 {
		return true
	}
	p1s := strings.Split(platform1, ",")
	slices.Sort(p1s)
	p2s := strings.Split(platform2, ",")
	slices.Sort(p2s)
	return slices.Compare(p1s, p2s) == 0
}

////////////////////////////////////////////////////////////////////////////////
// Delete Query
////////////////////////////////////////////////////////////////////////////////

func deleteQueryEndpoint(ctx context.Context, request interface{}, svc fleet.Service) (fleet.Errorer, error) {
	req := request.(*fleet.DeleteQueryRequest)
	var teamID *uint
	if req.TeamID != 0 {
		teamID = &req.TeamID
	}
	err := svc.DeleteQuery(ctx, teamID, req.Name)
	if err != nil {
		return fleet.DeleteQueryResponse{Err: err}, nil
	}
	return fleet.DeleteQueryResponse{}, nil
}

func (svc *Service) DeleteQuery(ctx context.Context, teamID *uint, name string) error {
	// Load query first to determine if the user can delete it.
	query, err := svc.ds.QueryByName(ctx, teamID, name)
	if err != nil {
		setAuthCheckedOnPreAuthErr(ctx)
		return err
	}
	if err := svc.authz.Authorize(ctx, query, fleet.ActionWrite); err != nil {
		return err
	}

	if err := svc.ds.DeleteQuery(ctx, teamID, name); err != nil {
		return err
	}

	// Delete the Redis counter for query results
	if svc.liveQueryStore != nil {
		if err = svc.liveQueryStore.DeleteQueryResultsCount(query.ID); err != nil {
			// Log the error but don't fail the request; this will get cleaned up
			// in the "query_results_cleanup" job.
			svc.logger.ErrorContext(ctx, "failed to delete query results count", "err", err, "query_id", query.ID)
		}
	}

	var logTeamID int64
	var teamName *string
	if query.TeamID != nil {
		logTeamID = int64(*query.TeamID) //nolint:gosec // dismiss G115
		if svc.EnterpriseOverrides != nil && svc.EnterpriseOverrides.TeamByIDOrName != nil {
			team, err := svc.EnterpriseOverrides.TeamByIDOrName(ctx, query.TeamID, nil)
			if err != nil {
				return err
			}
			teamName = &team.Name
		}
	} else {
		logTeamID = -1
		teamName = nil
	}

	if err := svc.NewActivity(
		ctx,
		authz.UserFromContext(ctx),
		fleet.ActivityTypeDeletedSavedQuery{
			Name:     name,
			TeamID:   logTeamID,
			TeamName: teamName,
		},
	); err != nil {
		return ctxerr.Wrap(ctx, err, "create activity for query deletion")
	}
	return nil
}

////////////////////////////////////////////////////////////////////////////////
// Delete Query By ID
////////////////////////////////////////////////////////////////////////////////

func deleteQueryByIDEndpoint(ctx context.Context, request interface{}, svc fleet.Service) (fleet.Errorer, error) {
	req := request.(*fleet.DeleteQueryByIDRequest)
	err := svc.DeleteQueryByID(ctx, req.ID)
	if err != nil {
		return fleet.DeleteQueryByIDResponse{Err: err}, nil
	}
	return fleet.DeleteQueryByIDResponse{}, nil
}

func (svc *Service) DeleteQueryByID(ctx context.Context, id uint) error {
	// Load query first to determine if the user can delete it.

	query, err := svc.ds.Query(ctx, id)
	if err != nil {
		setAuthCheckedOnPreAuthErr(ctx)
		return ctxerr.Wrap(ctx, err, "lookup query by ID")
	}
	if err := svc.authz.Authorize(ctx, query, fleet.ActionWrite); err != nil {
		return err
	}

	if err := svc.ds.DeleteQuery(ctx, query.TeamID, query.Name); err != nil {
		return ctxerr.Wrap(ctx, err, "delete query")
	}

	// Delete the Redis counter for query results
	if svc.liveQueryStore != nil {
		if err = svc.liveQueryStore.DeleteQueryResultsCount(query.ID); err != nil {
			// Log the error but don't fail the request; this will get cleaned up
			// in the "query_results_cleanup" job.
			svc.logger.ErrorContext(ctx, "failed to delete query results count", "err", err, "query_id", query.ID)
		}
	}

	var logTeamID int64
	var teamName *string
	if query.TeamID != nil {
		logTeamID = int64(*query.TeamID) //nolint:gosec // dismiss G115
		if svc.EnterpriseOverrides != nil && svc.EnterpriseOverrides.TeamByIDOrName != nil {
			team, err := svc.EnterpriseOverrides.TeamByIDOrName(ctx, query.TeamID, nil)
			if err != nil {
				return err
			}
			teamName = &team.Name
		}
	} else {
		logTeamID = -1
		teamName = nil
	}

	if err := svc.NewActivity(
		ctx,
		authz.UserFromContext(ctx),
		fleet.ActivityTypeDeletedSavedQuery{
			Name:     query.Name,
			TeamID:   logTeamID,
			TeamName: teamName,
		},
	); err != nil {
		return ctxerr.Wrap(ctx, err, "create activity for query deletion by id")
	}
	return nil
}

////////////////////////////////////////////////////////////////////////////////
// Delete Queries
////////////////////////////////////////////////////////////////////////////////

func deleteQueriesEndpoint(ctx context.Context, request interface{}, svc fleet.Service) (fleet.Errorer, error) {
	req := request.(*fleet.DeleteQueriesRequest)
	deleted, err := svc.DeleteQueries(ctx, req.IDs)
	if err != nil {
		return fleet.DeleteQueriesResponse{Err: err}, nil
	}
	return fleet.DeleteQueriesResponse{Deleted: deleted}, nil
}

func (svc *Service) DeleteQueries(ctx context.Context, ids []uint) (uint, error) {
	// Verify that the user is allowed to delete all the requested queries.
	var logTeamID int64 = -1
	var teamName *string
	for _, id := range ids {
		query, err := svc.ds.Query(ctx, id)
		if err != nil {
			setAuthCheckedOnPreAuthErr(ctx)
			return 0, ctxerr.Wrap(ctx, err, "lookup query by ID")
		}
		if err := svc.authz.Authorize(ctx, query, fleet.ActionWrite); err != nil {
			return 0, err
		}

		// Capture team information for activity logging.
		if query.TeamID != nil {
			logTeamID = int64(*query.TeamID) //nolint:gosec // dismiss G115
			if svc.EnterpriseOverrides != nil && svc.EnterpriseOverrides.TeamByIDOrName != nil {
				team, err := svc.EnterpriseOverrides.TeamByIDOrName(ctx, query.TeamID, nil)
				if err != nil {
					return 0, err
				}
				teamName = &team.Name
			}
		}
	}

	n, err := svc.ds.DeleteQueries(ctx, ids)
	if err != nil {
		return n, err
	}

	// Delete the Redis counters for query results
	if svc.liveQueryStore != nil {
		for _, id := range ids {
			if err = svc.liveQueryStore.DeleteQueryResultsCount(id); err != nil {
				// Log the error but don't fail the request; this will get cleaned up
				// in the "query_results_cleanup" job.
				svc.logger.ErrorContext(ctx, "failed to delete query results count", "err", err, "query_id", id)
			}
		}
	}

	if err := svc.NewActivity(
		ctx,
		authz.UserFromContext(ctx),
		fleet.ActivityTypeDeletedMultipleSavedQuery{
			IDs:      ids,
			Teamid:   logTeamID,
			TeamName: teamName,
		},
	); err != nil {
		return 0, ctxerr.Wrap(ctx, err, "create activity for query deletions")
	}

	return n, nil
}

////////////////////////////////////////////////////////////////////////////////
// Apply Query Specs
////////////////////////////////////////////////////////////////////////////////

func applyQuerySpecsEndpoint(ctx context.Context, request interface{}, svc fleet.Service) (fleet.Errorer, error) {
	req := request.(*fleet.ApplyQuerySpecsRequest)
	err := svc.ApplyQuerySpecs(ctx, req.Specs)
	if err != nil {
		return fleet.ApplyQuerySpecsResponse{Err: err}, nil
	}
	return fleet.ApplyQuerySpecsResponse{}, nil
}

func (svc *Service) ApplyQuerySpecs(ctx context.Context, specs []*fleet.QuerySpec) error {
	// 1. Turn specs into queries.
	queries := []*fleet.Query{}
	for _, spec := range specs {
		if len(spec.LabelsIncludeAny) > 0 && !license.IsPremium(ctx) {
			setAuthCheckedOnPreAuthErr(ctx)
			return fleet.ErrMissingLicense
		}
		query, err := svc.queryFromSpec(ctx, spec)
		if err != nil {
			setAuthCheckedOnPreAuthErr(ctx)
			return ctxerr.Wrap(ctx, err, "creating query from spec")
		}
		queries = append(queries, query)
	}
	// 2. Run authorization checks and verify their fields.
	for _, query := range queries {
		if err := svc.authz.Authorize(ctx, query, fleet.ActionWrite); err != nil {
			return err
		}
		if err := query.Verify(); err != nil {
			return ctxerr.Wrap(ctx, &fleet.BadRequestError{
				Message: fmt.Sprintf("query payload verification: %s", err),
			})
		}
	}
	// 3. Apply the queries.

	// first, find out if we should delete query results
	queriesToDiscardResults := make(map[uint]struct{})
	for _, query := range queries {
		dbQuery, err := svc.ds.QueryByName(ctx, query.TeamID, query.Name)
		if err != nil && !fleet.IsNotFound(err) {
			return ctxerr.Wrap(ctx, err, "fetching saved query")
		}

		if dbQuery == nil {
			// then we're creating a new query, so move on.
			continue
		}

		if (query.DiscardData && query.DiscardData != dbQuery.DiscardData) ||
			(query.Logging != dbQuery.Logging && query.Logging != fleet.LoggingSnapshot) ||
			query.Query != dbQuery.Query ||
			query.MinOsqueryVersion != dbQuery.MinOsqueryVersion ||
			!comparePlatforms(query.Platform, dbQuery.Platform) {
			queriesToDiscardResults[dbQuery.ID] = struct{}{}
		}
	}

	vc, ok := viewer.FromContext(ctx)
	if !ok {
		return ctxerr.New(ctx, "user must be authenticated to apply queries")
	}
	err := svc.ds.ApplyQueries(ctx, vc.UserID(), queries, queriesToDiscardResults)
	if err != nil {
		return ctxerr.Wrap(ctx, err, "applying queries")
	}

	// Reset the Redis counters for queries whose results were discarded
	if svc.liveQueryStore != nil {
		for queryID := range queriesToDiscardResults {
			if err = svc.liveQueryStore.SetQueryResultsCount(queryID, 0); err != nil {
				// Log the error but don't fail the request; this will get cleaned up
				// in the "query_results_cleanup" job.
				svc.logger.ErrorContext(ctx, "failed to set query results count", "err", err, "query_id", queryID)
			}
		}
	}

	if err := svc.NewActivity(
		ctx,
		authz.UserFromContext(ctx),
		fleet.ActivityTypeAppliedSpecSavedQuery{
			Specs: specs,
		},
	); err != nil {
		return ctxerr.Wrap(ctx, err, "create activity for query spec")
	}
	return nil
}

func (svc *Service) queryFromSpec(ctx context.Context, spec *fleet.QuerySpec) (*fleet.Query, error) {
	var teamID *uint
	if spec.TeamName != "" {
		team, err := svc.ds.TeamByName(ctx, spec.TeamName)
		if err != nil {
			return nil, ctxerr.Wrap(ctx, err, "get team by name")
		}
		teamID = &team.ID
	}
	logging := spec.Logging
	if logging == "" {
		logging = fleet.LoggingSnapshot
	}
	// Find labels by name.
	var queryLabels []fleet.LabelIdent
	if len(spec.LabelsIncludeAny) > 0 {
		vc, ok := viewer.FromContext(ctx)
		if !ok {
			return nil, fleet.ErrNoContext
		}

		labelsMap, err := svc.ds.LabelsByName(ctx, spec.LabelsIncludeAny, fleet.TeamFilter{User: vc.User, TeamID: teamID})
		if err != nil {
			return nil, ctxerr.Wrap(ctx, err, "get labels by name")
		}
		for labelName := range labelsMap {
			queryLabels = append(queryLabels, fleet.LabelIdent{LabelName: labelName, LabelID: labelsMap[labelName].ID})
		}
		// Make sure that all labels were found.
		for _, label := range spec.LabelsIncludeAny {
			if _, ok := labelsMap[label]; !ok {
				return nil, ctxerr.New(ctx, "label not found")
			}
		}
	}
	return &fleet.Query{
		Name:        spec.Name,
		Description: spec.Description,
		Query:       spec.Query,

		TeamID:             teamID,
		Interval:           spec.Interval,
		ObserverCanRun:     spec.ObserverCanRun,
		Platform:           spec.Platform,
		MinOsqueryVersion:  spec.MinOsqueryVersion,
		AutomationsEnabled: spec.AutomationsEnabled,
		Logging:            logging,
		DiscardData:        spec.DiscardData,
		LabelsIncludeAny:   queryLabels,
	}, nil
}

////////////////////////////////////////////////////////////////////////////////
// Get Query Specs
////////////////////////////////////////////////////////////////////////////////

func getQuerySpecsEndpoint(ctx context.Context, request interface{}, svc fleet.Service) (fleet.Errorer, error) {
	req := request.(*fleet.GetQuerySpecsRequest)
	var teamID *uint
	if req.TeamID != 0 {
		teamID = &req.TeamID
	}
	specs, err := svc.GetQuerySpecs(ctx, teamID)
	if err != nil {
		return fleet.GetQuerySpecsResponse{Err: err}, nil
	}
	return fleet.GetQuerySpecsResponse{Specs: specs}, nil
}

func (svc *Service) GetQuerySpecs(ctx context.Context, teamID *uint) ([]*fleet.QuerySpec, error) {
	queries, _, _, _, err := svc.ListQueries(ctx, fleet.ListOptions{}, teamID, nil, false, nil)
	if err != nil {
		return nil, ctxerr.Wrap(ctx, err, "getting queries")
	}

	// Turn queries into specs.
	var specs []*fleet.QuerySpec
	for _, query := range queries {
		spec, err := svc.specFromQuery(ctx, query)
		if err != nil {
			return nil, ctxerr.Wrap(ctx, err, "create spec from query")
		}
		specs = append(specs, spec)
	}
	return specs, nil
}

func (svc *Service) specFromQuery(ctx context.Context, query *fleet.Query) (*fleet.QuerySpec, error) {
	var teamName string
	if query.TeamID != nil {
		team, err := svc.ds.TeamLite(ctx, *query.TeamID)
		if err != nil {
			return nil, ctxerr.Wrap(ctx, err, "get team from id")
		}
		teamName = team.Name
	}
	labelsAny := []string{}
	for _, label := range query.LabelsIncludeAny {
		labelsAny = append(labelsAny, label.LabelName)
	}
	return &fleet.QuerySpec{
		Name:        query.Name,
		Description: query.Description,
		Query:       query.Query,

		TeamName:           teamName,
		Interval:           query.Interval,
		ObserverCanRun:     query.ObserverCanRun,
		Platform:           query.Platform,
		MinOsqueryVersion:  query.MinOsqueryVersion,
		AutomationsEnabled: query.AutomationsEnabled,
		Logging:            query.Logging,
		DiscardData:        query.DiscardData,
		LabelsIncludeAny:   labelsAny,
	}, nil
}

////////////////////////////////////////////////////////////////////////////////
// Get Query Spec
////////////////////////////////////////////////////////////////////////////////

func getQuerySpecEndpoint(ctx context.Context, request interface{}, svc fleet.Service) (fleet.Errorer, error) {
	req := request.(*fleet.GetQuerySpecRequest)
	var teamID *uint
	if req.TeamID != 0 {
		teamID = &req.TeamID
	}
	spec, err := svc.GetQuerySpec(ctx, teamID, req.Name)
	if err != nil {
		return fleet.GetQuerySpecResponse{Err: err}, nil
	}
	return fleet.GetQuerySpecResponse{Spec: spec}, nil
}

func (svc *Service) GetQuerySpec(ctx context.Context, teamID *uint, name string) (*fleet.QuerySpec, error) {
	// Check the user is allowed to get the query on the requested team.
	if err := svc.authz.Authorize(ctx, &fleet.Query{
		TeamID: teamID,
	}, fleet.ActionRead); err != nil {
		return nil, err
	}

	query, err := svc.ds.QueryByName(ctx, teamID, name)
	if err != nil {
		return nil, ctxerr.Wrap(ctx, err, "get query by name")
	}
	spec, err := svc.specFromQuery(ctx, query)
	if err != nil {
		return nil, ctxerr.Wrap(ctx, err, "create spec from query")
	}
	return spec, nil
}