Elgato_dark/fleet
1
0
Fork 0
mirror of https://github.com/fleetdm/fleet synced 2026-04-21 21:47:20 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 117
fleet/server/datastore/mysql/scripts.go

2556 lines
80 KiB
Go
Raw Blame History

package mysql
import (
"context"
"crypto/md5" //nolint:gosec
"database/sql"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"strings"
"time"
"unicode/utf8"
constants "github.com/fleetdm/fleet/v4/pkg/scripts"
"github.com/fleetdm/fleet/v4/server/contexts/ctxerr"
"github.com/fleetdm/fleet/v4/server/fleet"
"github.com/fleetdm/fleet/v4/server/ptr"
"github.com/go-kit/log/level"
"github.com/google/uuid"
"github.com/jmoiron/sqlx"
)
func (ds *Datastore) NewHostScriptExecutionRequest(ctx context.Context, request *fleet.HostScriptRequestPayload) (*fleet.HostScriptResult, error) {
var res *fleet.HostScriptResult
return res, ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
var err error
if request.ScriptContentID == 0 {
// then we are doing a sync execution, so create the contents first
scRes, err := insertScriptContents(ctx, tx, request.ScriptContents)
if err != nil {
return err
}
id, _ := scRes.LastInsertId()
request.ScriptContentID = uint(id) //nolint:gosec // dismiss G115
}
res, err = ds.newHostScriptExecutionRequest(ctx, tx, request, false)
return err
})
}
func (ds *Datastore) newHostScriptExecutionRequest(ctx context.Context, tx sqlx.ExtContext, request *fleet.HostScriptRequestPayload, isInternal bool) (*fleet.HostScriptResult, error) {
const (
getStmt = `
SELECT
ua.id, ua.host_id, ua.execution_id, ua.created_at, sua.script_id, sua.policy_id, ua.user_id,
payload->'$.sync_request' AS sync_request,
sc.contents as script_contents, sua.setup_experience_script_id
FROM
upcoming_activities ua
INNER JOIN script_upcoming_activities sua
ON ua.id = sua.upcoming_activity_id
INNER JOIN script_contents sc
ON sua.script_content_id = sc.id
WHERE
ua.id = ?
`
)
_, activityID, err := ds.insertNewHostScriptExecution(ctx, tx, request, isInternal)
if err != nil {
return nil, ctxerr.Wrap(ctx, err, "inserting new script execution request")
}
var script fleet.HostScriptResult
err = sqlx.GetContext(ctx, tx, &script, getStmt, activityID)
if err != nil {
return nil, ctxerr.Wrap(ctx, err, "getting the created host script activity to return")
}
return &script, nil
}
func (ds *Datastore) insertNewHostScriptExecution(ctx context.Context, tx sqlx.ExtContext, request *fleet.HostScriptRequestPayload, isInternal bool) (string, int64, error) {
const (
insUAStmt = `
INSERT INTO upcoming_activities
(host_id, priority, user_id, fleet_initiated, activity_type, execution_id, payload)
VALUES
(?, ?, ?, ?, 'script', ?,
JSON_OBJECT(
'sync_request', ?,
'is_internal', ?,
'user', (SELECT JSON_OBJECT('name', name, 'email', email, 'gravatar_url', gravatar_url) FROM users WHERE id = ?)
)
)`
insSUAStmt = `
INSERT INTO script_upcoming_activities
(upcoming_activity_id, script_id, script_content_id, policy_id, setup_experience_script_id)
VALUES
(?, ?, ?, ?, ?)
`
)
execID := uuid.New().String()
result, err := tx.ExecContext(ctx, insUAStmt,
request.HostID,
request.Priority(),
request.UserID,
request.PolicyID != nil, // fleet-initiated if request is via a policy failure
execID,
request.SyncRequest,
isInternal,
request.UserID,
)
if err != nil {
return "", 0, ctxerr.Wrap(ctx, err, "new script upcoming activity")
}
activityID, _ := result.LastInsertId()
_, err = tx.ExecContext(ctx, insSUAStmt,
activityID,
request.ScriptID,
request.ScriptContentID,
request.PolicyID,
request.SetupExperienceScriptID,
)
if err != nil {
return "", 0, ctxerr.Wrap(ctx, err, "new join script upcoming activity")
}
if _, err := ds.activateNextUpcomingActivity(ctx, tx, request.HostID, ""); err != nil {
return "", 0, ctxerr.Wrap(ctx, err, "activate next activity")
}
return execID, activityID, nil
}
func truncateScriptResult(output string) string {
const maxOutputRuneLen = 10000
if len(output) > utf8.UTFMax*maxOutputRuneLen {
// truncate the bytes as we know the output is too long, no point
// converting more bytes than needed to runes.
output = output[len(output)-(utf8.UTFMax*maxOutputRuneLen):]
}
if utf8.RuneCountInString(output) > maxOutputRuneLen {
outputRunes := []rune(output)
output = string(outputRunes[len(outputRunes)-maxOutputRuneLen:])
}
return output
}
func (ds *Datastore) SetHostScriptExecutionResult(ctx context.Context, result *fleet.HostScriptResultPayload) (*fleet.HostScriptResult,
string, error,
) {
const resultExistsStmt = `
SELECT
1
FROM
host_script_results
WHERE
host_id = ? AND
execution_id = ? AND
exit_code IS NOT NULL
`
const updStmt = `
UPDATE host_script_results SET
output = ?,
runtime = ?,
exit_code = ?,
timeout = ?
WHERE
host_id = ? AND
execution_id = ?`
const hostMDMActionsStmt = `
SELECT 'uninstall' AS action
FROM
host_software_installs
WHERE
execution_id = :execution_id AND host_id = :host_id
UNION -- host_mdm_actions query (and thus row in union) must be last to avoid #25144
SELECT
CASE
WHEN lock_ref = :execution_id THEN 'lock_ref'
WHEN unlock_ref = :execution_id THEN 'unlock_ref'
WHEN wipe_ref = :execution_id THEN 'wipe_ref'
ELSE ''
END AS action
FROM
host_mdm_actions
WHERE
host_id = :host_id
`
output := truncateScriptResult(result.Output)
var hsr *fleet.HostScriptResult
var action string
err := ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
var resultExists bool
err := sqlx.GetContext(ctx, tx, &resultExists, resultExistsStmt, result.HostID, result.ExecutionID)
if err != nil && !errors.Is(err, sql.ErrNoRows) {
return ctxerr.Wrap(ctx, err, "check if host script result exists")
}
if resultExists {
level.Debug(ds.logger).Log("msg", "duplicate script execution result sent, will be ignored (original result is preserved)",
"host_id", result.HostID,
"execution_id", result.ExecutionID,
)
// still do the activate next activity to ensure progress as there was
// an unexpected flow if we get here.
if _, err := ds.activateNextUpcomingActivity(ctx, tx, result.HostID, result.ExecutionID); err != nil {
return ctxerr.Wrap(ctx, err, "activate next activity")
}
// succeed but leave hsr nil
return nil
}
res, err := tx.ExecContext(ctx, updStmt,
output,
result.Runtime,
// Windows error codes are signed 32-bit integers, but are
// returned as unsigned integers by the windows API. The
// software that receives them is responsible for casting
// it to a 32-bit signed integer.
// See /orbit/pkg/scripts/exec_windows.go
int32(result.ExitCode), //nolint:gosec // dismiss G115
result.Timeout,
result.HostID,
result.ExecutionID,
)
if err != nil {
return ctxerr.Wrap(ctx, err, "update host script result")
}
if n, _ := res.RowsAffected(); n > 0 {
// it did update, so return the updated result
hsr, err = ds.getHostScriptExecutionResultDB(ctx, tx, result.ExecutionID, scriptExecutionSearchOpts{IncludeCanceled: true})
if err != nil {
return ctxerr.Wrap(ctx, err, "load updated host script result")
}
// look up if that script was a lock/unlock/wipe/uninstall script for that host,
// and if so update the host_mdm_actions table accordingly.
namedArgs := map[string]any{
"host_id": result.HostID,
"execution_id": result.ExecutionID,
}
stmt, args, err := sqlx.Named(hostMDMActionsStmt, namedArgs)
if err != nil {
return ctxerr.Wrap(ctx, err, "build named query for host mdm actions")
}
err = sqlx.GetContext(ctx, tx, &action, stmt, args...)
if err != nil && !errors.Is(err, sql.ErrNoRows) { // ignore ErrNoRows, refCol will be empty
return ctxerr.Wrap(ctx, err, "lookup host script corresponding mdm action")
}
switch action {
case "":
// do nothing
case "uninstall":
err = ds.updateUninstallStatusFromResult(ctx, tx, result.HostID, result.ExecutionID, result.ExitCode)
if err != nil {
return ctxerr.Wrap(ctx, err, "update host uninstall action based on script result")
}
default: // lock/unlock/wipe
err = updateHostLockWipeStatusFromResult(ctx, tx, result.HostID, action, result.ExitCode == 0)
if err != nil {
return ctxerr.Wrap(ctx, err, "update host mdm action based on script result")
}
}
}
if _, err := ds.activateNextUpcomingActivity(ctx, tx, result.HostID, result.ExecutionID); err != nil {
return ctxerr.Wrap(ctx, err, "activate next activity")
}
return nil
})
if err != nil {
return nil, "", err
}
return hsr, action, nil
}
func (ds *Datastore) ListPendingHostScriptExecutions(ctx context.Context, hostID uint, onlyShowInternal bool) ([]*fleet.HostScriptResult, error) {
return ds.listUpcomingHostScriptExecutions(ctx, hostID, onlyShowInternal, false)
}
func (ds *Datastore) ListReadyToExecuteScriptsForHost(ctx context.Context, hostID uint, onlyShowInternal bool) ([]*fleet.HostScriptResult, error) {
return ds.listUpcomingHostScriptExecutions(ctx, hostID, onlyShowInternal, true)
}
func (ds *Datastore) listUpcomingHostScriptExecutions(ctx context.Context, hostID uint, onlyShowInternal, onlyReadyToExecute bool) ([]*fleet.HostScriptResult, error) {
extraWhere := ""
if onlyShowInternal {
// software_uninstalls are implicitly internal
extraWhere = " AND COALESCE(ua.payload->'$.is_internal', 1) = 1"
}
if onlyReadyToExecute {
extraWhere += " AND ua.activated_at IS NOT NULL"
}
// this selects software uninstalls too as they run as scripts
listStmt := fmt.Sprintf(`
SELECT
id,
host_id,
execution_id,
script_id,
created_at
FROM (
SELECT
ua.id,
ua.host_id,
ua.execution_id,
sua.script_id,
ua.priority,
ua.created_at,
IF(ua.activated_at IS NULL, 0, 1) AS topmost
FROM
upcoming_activities ua
-- left join because software_uninstall has no script join
LEFT JOIN script_upcoming_activities sua
ON ua.id = sua.upcoming_activity_id
WHERE
ua.host_id = ? AND
ua.activity_type IN ('script', 'software_uninstall')
%s
ORDER BY topmost DESC, priority DESC, created_at ASC) t`, extraWhere)
var results []*fleet.HostScriptResult
if err := sqlx.SelectContext(ctx, ds.reader(ctx), &results, listStmt, hostID); err != nil {
return nil, ctxerr.Wrap(ctx, err, "list pending host script executions")
}
return results, nil
}
func (ds *Datastore) IsExecutionPendingForHost(ctx context.Context, hostID uint, scriptID uint) (bool, error) {
const getStmt = `
SELECT
1
FROM
upcoming_activities ua
INNER JOIN script_upcoming_activities sua
ON ua.id = sua.upcoming_activity_id
WHERE
ua.host_id = ? AND
ua.activity_type = 'script' AND
sua.script_id = ?
`
var results []*uint
if err := sqlx.SelectContext(ctx, ds.reader(ctx), &results, getStmt, hostID, scriptID); err != nil {
return false, ctxerr.Wrap(ctx, err, "is execution pending for host")
}
return len(results) > 0, nil
}
type scriptExecutionSearchOpts struct {
IncludeCanceled bool
UninstallHostID uint
}
func (ds *Datastore) GetHostScriptExecutionResult(ctx context.Context, execID string) (*fleet.HostScriptResult, error) {
return ds.getHostScriptExecutionResultDB(ctx, ds.reader(ctx), execID, scriptExecutionSearchOpts{})
}
func (ds *Datastore) GetSelfServiceUninstallScriptExecutionResult(ctx context.Context, execID string, hostID uint) (*fleet.HostScriptResult, error) {
return ds.getHostScriptExecutionResultDB(ctx, ds.reader(ctx), execID, scriptExecutionSearchOpts{UninstallHostID: hostID})
}
func (ds *Datastore) getHostScriptExecutionResultDB(ctx context.Context, q sqlx.QueryerContext, execID string, opts scriptExecutionSearchOpts) (*fleet.HostScriptResult, error) {
var activeParams []any
canceledCondition := ""
if !opts.IncludeCanceled {
canceledCondition = " AND hsr.canceled = 0"
}
uninstallCondition := ""
if opts.UninstallHostID > 0 {
uninstallCondition = `JOIN host_software_installs hsi ON hsi.execution_id = hsr.execution_id
AND hsi.uninstall = TRUE AND hsr.host_id = ?`
activeParams = append(activeParams, opts.UninstallHostID)
}
activeParams = append(activeParams, execID)
getActiveStmt := fmt.Sprintf(`
SELECT
hsr.id,
hsr.host_id,
hsr.execution_id,
sc.contents as script_contents,
hsr.script_id,
hsr.policy_id,
hsr.output,
hsr.runtime,
hsr.exit_code,
hsr.timeout,
hsr.created_at,
hsr.user_id,
hsr.sync_request,
hsr.host_deleted_at,
hsr.setup_experience_script_id,
hsr.canceled,
bahr.batch_execution_id
FROM
host_script_results hsr
LEFT JOIN
batch_activity_host_results bahr ON hsr.execution_id = bahr.host_execution_id
JOIN
script_contents sc
%s
WHERE
hsr.execution_id = ? AND
hsr.script_content_id = sc.id
%s
`, uninstallCondition, canceledCondition)
// We don't include upcoming uninstall script executions in results (different activity type, and they're blank anyway)
const getUpcomingStmt = `
SELECT
0 as id,
ua.host_id,
ua.execution_id,
sc.contents as script_contents,
sua.script_id,
sua.policy_id,
'' as output,
0 as runtime,
NULL as exit_code,
NULL as timeout,
ua.created_at,
ua.user_id,
COALESCE(ua.payload->'$.sync_request', 0) as sync_request,
NULL as host_deleted_at,
sua.setup_experience_script_id,
0 as canceled
FROM
upcoming_activities ua
INNER JOIN script_upcoming_activities sua
ON ua.id = sua.upcoming_activity_id
INNER JOIN
script_contents sc
ON sua.script_content_id = sc.id
WHERE
ua.execution_id = ? AND
ua.activity_type = 'script'
`
var result fleet.HostScriptResult
if err := sqlx.GetContext(ctx, q, &result, getActiveStmt, activeParams...); err != nil {
if errors.Is(err, sql.ErrNoRows) {
// try with upcoming activities
err = sqlx.GetContext(ctx, q, &result, getUpcomingStmt, execID)
if errors.Is(err, sql.ErrNoRows) {
return nil, ctxerr.Wrap(ctx, notFound("HostScriptResult").WithName(execID))
}
}
if err != nil {
return nil, ctxerr.Wrap(ctx, err, "get host script result")
}
}
return &result, nil
}
func (ds *Datastore) NewScript(ctx context.Context, script *fleet.Script) (*fleet.Script, error) {
var res sql.Result
err := ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
var err error
// first insert script contents
scRes, err := insertScriptContents(ctx, tx, script.ScriptContents)
if err != nil {
return err
}
id, _ := scRes.LastInsertId()
// then create the script entity
res, err = insertScript(ctx, tx, script, uint(id)) //nolint:gosec // dismiss G115
return err
})
if err != nil {
return nil, err
}
id, _ := res.LastInsertId()
return ds.getScriptDB(ctx, ds.writer(ctx), uint(id)) //nolint:gosec // dismiss G115
}
func (ds *Datastore) UpdateScriptContents(ctx context.Context, scriptID uint, scriptContents string) (*fleet.Script, error) {
err := ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
// Get the current script_content_id
var oldContentID int64
getCurrentStmt := `SELECT script_content_id FROM scripts WHERE id = ?`
err := sqlx.GetContext(ctx, tx, &oldContentID, getCurrentStmt, scriptID)
if err != nil {
return ctxerr.Wrap(ctx, err, "getting current script content id")
}
// Insert or get existing content (insertScriptContents handles deduplication)
scRes, err := insertScriptContents(ctx, tx, scriptContents)
if err != nil {
return ctxerr.Wrap(ctx, err, "inserting/getting script contents")
}
newContentID, _ := scRes.LastInsertId()
// Update the script to point to the new content
if newContentID != oldContentID {
updateStmt := `
UPDATE scripts
SET script_content_id = ?
WHERE id = ?
`
_, err = tx.ExecContext(ctx, updateStmt, newContentID, scriptID)
if err != nil {
return ctxerr.Wrap(ctx, err, "updating script content reference")
}
// Try to clean up the old content if no longer used
// Don't fail the transaction if cleanup fails; just log it
if err := ds.cleanupScriptContent(ctx, tx, uint(oldContentID)); err != nil { //nolint:gosec
level.Error(ds.logger).Log("msg", "failed to cleanup orphaned script content",
"script_id", scriptID, "old_content_id", oldContentID, "err", err)
ctxerr.Handle(ctx, err)
}
} else {
// Just update the timestamp
_, err = tx.ExecContext(ctx, "UPDATE scripts SET updated_at = NOW() WHERE id = ?", scriptID)
if err != nil {
return ctxerr.Wrap(ctx, err, "updating script updated_at time")
}
}
// Cancel pending executions
if err := ds.cancelUpcomingScriptActivities(ctx, tx, scriptID); err != nil {
return ctxerr.Wrap(ctx, err, "canceling upcoming script executions")
}
return nil
})
if err != nil {
return nil, ctxerr.Wrap(ctx, err, "updating script contents")
}
return ds.Script(ctx, scriptID)
}
func (ds *Datastore) cancelUpcomingScriptActivities(ctx context.Context, db sqlx.ExtContext, scriptID uint) error {
const stmt = `
SELECT
ua.execution_id,
ua.host_id
FROM
script_upcoming_activities sua
INNER JOIN
upcoming_activities ua ON ua.id = sua.upcoming_activity_id
WHERE
sua.script_id = ?
`
var upcomingExecutions []struct {
ExecutionID string `db:"execution_id"`
HostID uint `db:"host_id"`
}
if err := sqlx.SelectContext(ctx, db, &upcomingExecutions, stmt, scriptID); err != nil {
return ctxerr.Wrap(ctx, err, "selecting upcoming script executions")
}
for _, upcomingExecution := range upcomingExecutions {
if _, err := ds.cancelHostUpcomingActivity(ctx, db, upcomingExecution.HostID, upcomingExecution.ExecutionID); err != nil {
return ctxerr.Wrap(ctx, err, "canceling upcoming activity")
}
}
return nil
}
func insertScript(ctx context.Context, tx sqlx.ExtContext, script *fleet.Script, scriptContentsID uint) (sql.Result, error) {
const insertStmt = `
INSERT INTO
scripts (
team_id, global_or_team_id, name, script_content_id
)
VALUES
(?, ?, ?, ?)
`
var globalOrTeamID uint
if script.TeamID != nil {
globalOrTeamID = *script.TeamID
}
res, err := tx.ExecContext(ctx, insertStmt,
script.TeamID, globalOrTeamID, script.Name, scriptContentsID)
if err != nil {
if IsDuplicate(err) {
// name already exists for this team/global
err = alreadyExists("Script", script.Name)
} else if isChildForeignKeyError(err) {
// team does not exist
err = foreignKey("scripts", fmt.Sprintf("team_id=%v", script.TeamID))
}
return nil, ctxerr.Wrap(ctx, err, "insert script")
}
return res, nil
}
func insertScriptContents(ctx context.Context, tx sqlx.ExtContext, contents string) (sql.Result, error) {
const insertStmt = `
INSERT INTO
script_contents (
md5_checksum, contents
)
VALUES (UNHEX(?),?)
ON DUPLICATE KEY UPDATE
id=LAST_INSERT_ID(id)
`
md5Checksum := md5ChecksumScriptContent(contents)
res, err := tx.ExecContext(ctx, insertStmt, md5Checksum, contents)
if err != nil {
return nil, ctxerr.Wrap(ctx, err, "insert script contents")
}
return res, nil
}
func md5ChecksumScriptContent(s string) string {
return md5ChecksumBytes([]byte(s))
}
func md5ChecksumBytes(b []byte) string {
rawChecksum := md5.Sum(b) //nolint:gosec
return strings.ToUpper(hex.EncodeToString(rawChecksum[:]))
}
func (ds *Datastore) cleanupScriptContent(ctx context.Context, tx sqlx.ExtContext, contentID uint) error {
// Check if this content is still being used anywhere
var usageCount int
stmt := `
SELECT COUNT(*) FROM (
SELECT 1 FROM scripts WHERE script_content_id = ?
UNION ALL
SELECT 1 FROM setup_experience_scripts WHERE script_content_id = ?
UNION ALL
SELECT 1 FROM software_installers WHERE
install_script_content_id = ?
OR uninstall_script_content_id = ?
OR post_install_script_content_id = ?
UNION ALL
SELECT 1 FROM script_upcoming_activities WHERE script_content_id = ?
UNION ALL
SELECT 1 FROM host_script_results WHERE script_content_id = ?
) t
`
err := sqlx.GetContext(ctx, tx, &usageCount, stmt,
contentID, contentID, contentID, contentID, contentID, contentID, contentID)
if err != nil {
return ctxerr.Wrap(ctx, err, "checking script content usage for cleanup")
}
if usageCount == 0 {
// Not being used, safe to delete
deleteStmt := `DELETE FROM script_contents WHERE id = ?`
_, err = tx.ExecContext(ctx, deleteStmt, contentID)
if err != nil {
return ctxerr.Wrap(ctx, err, "deleting unused script content")
}
}
return nil
}
func (ds *Datastore) Script(ctx context.Context, id uint) (*fleet.Script, error) {
return ds.getScriptDB(ctx, ds.reader(ctx), id)
}
func (ds *Datastore) getScriptDB(ctx context.Context, q sqlx.QueryerContext, id uint) (*fleet.Script, error) {
const getStmt = `
SELECT
id,
team_id,
name,
created_at,
updated_at,
script_content_id
FROM
scripts
WHERE
id = ?
`
var script fleet.Script
if err := sqlx.GetContext(ctx, q, &script, getStmt, id); err != nil {
if err == sql.ErrNoRows {
return nil, notFound("Script").WithID(id)
}
return nil, ctxerr.Wrap(ctx, err, "get script")
}
return &script, nil
}
func (ds *Datastore) GetScriptContents(ctx context.Context, id uint) ([]byte, error) {
const getStmt = `
SELECT
sc.contents
FROM
script_contents sc
JOIN scripts s ON s.script_content_id = sc.id
WHERE
s.id = ?
`
var contents []byte
if err := sqlx.GetContext(ctx, ds.reader(ctx), &contents, getStmt, id); err != nil {
if err == sql.ErrNoRows {
return nil, notFound("Script").WithID(id)
}
return nil, ctxerr.Wrap(ctx, err, "get script contents")
}
return contents, nil
}
func (ds *Datastore) GetAnyScriptContents(ctx context.Context, id uint) ([]byte, error) {
const getStmt = `
SELECT
sc.contents
FROM
script_contents sc
WHERE
sc.id = ?
`
var contents []byte
if err := sqlx.GetContext(ctx, ds.reader(ctx), &contents, getStmt, id); err != nil {
if errors.Is(err, sql.ErrNoRows) {
return nil, notFound("Script").WithID(id)
}
return nil, ctxerr.Wrap(ctx, err, "get any script contents")
}
return contents, nil
}
var errDeleteScriptWithAssociatedPolicy = &fleet.ConflictError{Message: "Couldn't delete. Policy automation uses this script. Please remove this script from associated policy automations and try again."}
func (ds *Datastore) DeleteScript(ctx context.Context, id uint) error {
var activateAffectedHosts []uint
err := ds.withTx(ctx, func(tx sqlx.ExtContext) error {
_, err := tx.ExecContext(ctx, `DELETE FROM host_script_results WHERE script_id = ?
AND exit_code IS NULL AND (sync_request = 0 OR created_at >= NOW() - INTERVAL ? SECOND)`,
id, int(constants.MaxServerWaitTime.Seconds()),
)
if err != nil {
return ctxerr.Wrapf(ctx, err, "cancel pending script executions")
}
// load hosts that will have their upcoming_activities deleted, if that
// activity is "activated", as that means we will have to call
// activateNextUpcomingActivity for those hosts.
loadAffectedHostsStmt := `
SELECT
DISTINCT host_id
FROM
upcoming_activities ua
INNER JOIN script_upcoming_activities sua
ON ua.id = sua.upcoming_activity_id
WHERE sua.script_id = ? AND
ua.activity_type = 'script' AND
ua.activated_at IS NOT NULL AND
(ua.payload->'$.sync_request' = 0 OR
ua.created_at >= NOW() - INTERVAL ? SECOND)`
var affectedHosts []uint
if err := sqlx.SelectContext(ctx, tx, &affectedHosts, loadAffectedHostsStmt,
id, int(constants.MaxServerWaitTime.Seconds())); err != nil {
return ctxerr.Wrapf(ctx, err, "load affected hosts")
}
activateAffectedHosts = affectedHosts
_, err = tx.ExecContext(ctx, `DELETE FROM upcoming_activities
USING upcoming_activities
INNER JOIN script_upcoming_activities sua
ON upcoming_activities.id = sua.upcoming_activity_id
WHERE sua.script_id = ? AND
upcoming_activities.activity_type = 'script' AND
(upcoming_activities.payload->'$.sync_request' = 0 OR
upcoming_activities.created_at >= NOW() - INTERVAL ? SECOND)
`,
id, int(constants.MaxServerWaitTime.Seconds()),
)
if err != nil {
return ctxerr.Wrapf(ctx, err, "cancel upcoming pending script executions")
}
_, err = tx.ExecContext(ctx, `DELETE FROM scripts WHERE id = ?`, id)
if err != nil {
if isMySQLForeignKey(err) {
// Check if the script is referenced by a policy automation.
var count int
if err := sqlx.GetContext(ctx, tx, &count, `SELECT COUNT(*) FROM policies WHERE script_id = ?`, id); err != nil {
return ctxerr.Wrapf(ctx, err, "getting reference from policies")
}
if count > 0 {
return ctxerr.Wrap(ctx, errDeleteScriptWithAssociatedPolicy, "delete script")
}
}
return ctxerr.Wrap(ctx, err, "delete script")
}
return nil
})
if err != nil {
return err
}
// we call this outside of the transaction to avoid a
// long-running/deadlock-prone transaction, as many hosts could be affected.
return ds.activateNextUpcomingActivityForBatchOfHosts(ctx, activateAffectedHosts)
}
// deletePendingHostScriptExecutionsForPolicy should be called when a policy is deleted to remove any pending script executions
func (ds *Datastore) deletePendingHostScriptExecutionsForPolicy(ctx context.Context, teamID *uint, policyID uint) error {
var globalOrTeamID uint
if teamID != nil {
globalOrTeamID = *teamID
}
deletePendingFunc := func(stmt string, args ...any) error {
_, err := ds.writer(ctx).ExecContext(ctx, stmt, args...)
return ctxerr.Wrap(ctx, err, "delete pending host script executions for policy")
}
deleteHSRStmt := `
DELETE FROM
host_script_results
WHERE
policy_id = ? AND
script_id IN (
SELECT id FROM scripts WHERE scripts.global_or_team_id = ?
) AND
exit_code IS NULL
`
if err := deletePendingFunc(deleteHSRStmt, policyID, globalOrTeamID); err != nil {
return err
}
loadAffectedHostsStmt := `
SELECT
DISTINCT host_id
FROM
upcoming_activities ua
INNER JOIN script_upcoming_activities sua
ON ua.id = sua.upcoming_activity_id
WHERE
ua.activity_type = 'script' AND
ua.activated_at IS NOT NULL AND
sua.policy_id = ? AND
sua.script_id IN (
SELECT id FROM scripts WHERE scripts.global_or_team_id = ?
)`
var affectedHosts []uint
if err := sqlx.SelectContext(ctx, ds.reader(ctx), &affectedHosts,
loadAffectedHostsStmt, policyID, globalOrTeamID); err != nil {
return err
}
deleteUAStmt := `
DELETE FROM
upcoming_activities
USING
upcoming_activities
INNER JOIN script_upcoming_activities sua
ON upcoming_activities.id = sua.upcoming_activity_id
WHERE
upcoming_activities.activity_type = 'script' AND
sua.policy_id = ? AND
sua.script_id IN (
SELECT id FROM scripts WHERE scripts.global_or_team_id = ?
)
`
if err := deletePendingFunc(deleteUAStmt, policyID, globalOrTeamID); err != nil {
return err
}
return ds.activateNextUpcomingActivityForBatchOfHosts(ctx, affectedHosts)
}
func (ds *Datastore) ListScripts(ctx context.Context, teamID *uint, opt fleet.ListOptions) ([]*fleet.Script, *fleet.PaginationMetadata, error) {
var scripts []*fleet.Script
const selectStmt = `
SELECT
s.id,
s.team_id,
s.name,
s.created_at,
s.updated_at
FROM
scripts s
WHERE
s.global_or_team_id = ?
`
var globalOrTeamID uint
if teamID != nil {
globalOrTeamID = *teamID
}
args := []any{globalOrTeamID}
stmt, args := appendListOptionsWithCursorToSQL(selectStmt, args, &opt)
if err := sqlx.SelectContext(ctx, ds.reader(ctx), &scripts, stmt, args...); err != nil {
return nil, nil, ctxerr.Wrap(ctx, err, "select scripts")
}
var metaData *fleet.PaginationMetadata
if opt.IncludeMetadata {
metaData = &fleet.PaginationMetadata{HasPreviousResults: opt.Page > 0}
if len(scripts) > int(opt.PerPage) { //nolint:gosec // dismiss G115
metaData.HasNextResults = true
scripts = scripts[:len(scripts)-1]
}
}
return scripts, metaData, nil
}
func (ds *Datastore) GetScriptIDByName(ctx context.Context, name string, teamID *uint) (uint, error) {
const selectStmt = `
SELECT
id
FROM
scripts
WHERE
global_or_team_id = ?
AND name = ?
`
var globalOrTeamID uint
if teamID != nil {
globalOrTeamID = *teamID
}
var id uint
if err := sqlx.GetContext(ctx, ds.reader(ctx), &id, selectStmt, globalOrTeamID, name); err != nil {
if err == sql.ErrNoRows {
return 0, notFound("Script").WithName(name)
}
return 0, ctxerr.Wrap(ctx, err, "get script by name")
}
return id, nil
}
func (ds *Datastore) GetHostScriptDetails(ctx context.Context, hostID uint, teamID *uint, opt fleet.ListOptions, hostPlatform string) ([]*fleet.HostScriptDetail, *fleet.PaginationMetadata, error) {
var globalOrTeamID uint
if teamID != nil {
globalOrTeamID = *teamID
}
var extension string
switch {
case hostPlatform == "windows":
// filter by .ps1 extension
extension = `%.ps1`
case fleet.IsUnixLike(hostPlatform):
// filter by .sh extension
extension = `%.sh`
default:
// no extension filter
}
type row struct {
ScriptID uint `db:"script_id"`
Name string `db:"name"`
HSRID *uint `db:"hsr_id"`
ExecutionID *string `db:"execution_id"`
ExecutedAt *time.Time `db:"executed_at"`
ExitCode *int64 `db:"exit_code"`
}
sql := `
WITH all_latest_activities AS (
-- Use window function to efficiently find the latest execution per script
-- This is O(n) (a self-join approach would be O(n²))
SELECT * FROM (
SELECT
id,
host_id,
script_id,
execution_id,
created_at,
exit_code,
'completed' as source,
ROW_NUMBER() OVER (
PARTITION BY script_id
ORDER BY created_at DESC, id DESC
) AS row_num
FROM
host_script_results
WHERE
host_id = ? AND
canceled = 0
) completed_ranked
WHERE row_num = 1
UNION ALL
-- latest from upcoming_activities
SELECT * FROM (
SELECT
NULL as id,
ua.host_id,
sua.script_id,
ua.execution_id,
ua.created_at,
NULL as exit_code,
'upcoming' as source,
ROW_NUMBER() OVER (
PARTITION BY sua.script_id
ORDER BY ua.created_at DESC, ua.id DESC
) AS row_num
FROM
upcoming_activities ua
INNER JOIN script_upcoming_activities sua
ON ua.id = sua.upcoming_activity_id
WHERE
ua.host_id = ? AND
ua.activity_type = 'script'
) upcoming_ranked
WHERE row_num = 1
)
SELECT
s.id AS script_id,
s.name,
latest.id AS hsr_id,
latest.created_at AS executed_at,
latest.execution_id,
latest.exit_code
FROM
scripts s
LEFT JOIN (
-- Pick the most recent between completed and upcoming for each script
SELECT * FROM (
SELECT
*,
ROW_NUMBER() OVER (
PARTITION BY script_id
ORDER BY
CASE WHEN source = 'upcoming' THEN 1 ELSE 2 END, -- Prefer upcoming over completed
created_at DESC,
id DESC
) AS final_rn
FROM all_latest_activities
) final_ranked
WHERE final_rn = 1
) latest
ON s.id = latest.script_id
WHERE
(latest.host_id IS NULL OR latest.host_id = ?)
AND s.global_or_team_id = ?
`
args := []any{hostID, hostID, hostID, globalOrTeamID}
if len(extension) > 0 {
args = append(args, extension)
sql += `
AND s.name LIKE ?
`
}
stmt, args := appendListOptionsWithCursorToSQL(sql, args, &opt)
var rows []*row
if err := sqlx.SelectContext(ctx, ds.reader(ctx), &rows, stmt, args...); err != nil {
return nil, nil, ctxerr.Wrap(ctx, err, "get host script details")
}
var metaData *fleet.PaginationMetadata
if opt.IncludeMetadata {
metaData = &fleet.PaginationMetadata{HasPreviousResults: opt.Page > 0}
if len(rows) > int(opt.PerPage) { //nolint:gosec // dismiss G115
metaData.HasNextResults = true
rows = rows[:len(rows)-1]
}
}
results := make([]*fleet.HostScriptDetail, 0, len(rows))
for _, r := range rows {
results = append(results, fleet.NewHostScriptDetail(hostID, r.ScriptID, r.Name, r.ExecutionID, r.ExecutedAt, r.ExitCode, r.HSRID))
}
return results, metaData, nil
}
func (ds *Datastore) BatchSetScripts(ctx context.Context, tmID *uint, scripts []*fleet.Script) ([]fleet.ScriptResponse, error) {
const loadExistingScripts = `
SELECT
name
FROM
scripts
WHERE
global_or_team_id = ? AND
name IN (?)
`
const deleteAllScriptsInTeam = `
DELETE FROM
scripts
WHERE
global_or_team_id = ?
`
const unsetAllScriptsFromPolicies = `UPDATE policies SET script_id = NULL WHERE team_id = ?`
const clearAllPendingExecutionsHSR = `DELETE FROM host_script_results WHERE
exit_code IS NULL AND (sync_request = 0 OR created_at >= NOW() - INTERVAL ? SECOND)
AND script_id IN (SELECT id FROM scripts WHERE global_or_team_id = ?)`
const loadAffectedHostsAllPendingExecutionsUA = `
SELECT
DISTINCT host_id
FROM
upcoming_activities ua
INNER JOIN script_upcoming_activities sua
ON ua.id = sua.upcoming_activity_id
WHERE
ua.activity_type = 'script'
AND ua.activated_at IS NOT NULL
AND (ua.payload->'$.sync_request' = 0 OR ua.created_at >= NOW() - INTERVAL ? SECOND)
AND sua.script_id IN (SELECT id FROM scripts WHERE global_or_team_id = ?)`
const clearAllPendingExecutionsUA = `DELETE FROM upcoming_activities
USING
upcoming_activities
INNER JOIN script_upcoming_activities sua
ON upcoming_activities.id = sua.upcoming_activity_id
WHERE
upcoming_activities.activity_type = 'script'
AND (upcoming_activities.payload->'$.sync_request' = 0 OR upcoming_activities.created_at >= NOW() - INTERVAL ? SECOND)
AND sua.script_id IN (SELECT id FROM scripts WHERE global_or_team_id = ?)`
const unsetScriptsNotInListFromPolicies = `
UPDATE policies SET script_id = NULL
WHERE script_id IN (SELECT id FROM scripts WHERE global_or_team_id = ? AND name NOT IN (?))
`
const deleteScriptsNotInList = `
DELETE FROM
scripts
WHERE
global_or_team_id = ? AND
name NOT IN (?)
`
const clearPendingExecutionsNotInListHSR = `DELETE FROM host_script_results WHERE
exit_code IS NULL AND (sync_request = 0 OR created_at >= NOW() - INTERVAL ? SECOND)
AND script_id IN (SELECT id FROM scripts WHERE global_or_team_id = ? AND name NOT IN (?))`
const loadAffectedHostsPendingExecutionsNotInListUA = `
SELECT
DISTINCT host_id
FROM
upcoming_activities ua
INNER JOIN script_upcoming_activities sua
ON ua.id = sua.upcoming_activity_id
WHERE
ua.activity_type = 'script'
AND ua.activated_at IS NOT NULL
AND (ua.payload->'$.sync_request' = 0 OR ua.created_at >= NOW() - INTERVAL ? SECOND)
AND sua.script_id IN (SELECT id FROM scripts WHERE global_or_team_id = ? AND name NOT IN (?))`
const clearPendingExecutionsNotInListUA = `DELETE FROM upcoming_activities
USING
upcoming_activities
INNER JOIN script_upcoming_activities sua
ON upcoming_activities.id = sua.upcoming_activity_id
WHERE
upcoming_activities.activity_type = 'script'
AND (upcoming_activities.payload->'$.sync_request' = 0 OR upcoming_activities.created_at >= NOW() - INTERVAL ? SECOND)
AND sua.script_id IN (SELECT id FROM scripts WHERE global_or_team_id = ? AND name NOT IN (?))`
const insertNewOrEditedScript = `
INSERT INTO
scripts (
team_id, global_or_team_id, name, script_content_id
)
VALUES
(?, ?, ?, ?)
ON DUPLICATE KEY UPDATE
script_content_id = VALUES(script_content_id), id=LAST_INSERT_ID(id)
`
const clearPendingExecutionsWithObsoleteScriptHSR = `DELETE FROM host_script_results WHERE
exit_code IS NULL AND (sync_request = 0 OR created_at >= NOW() - INTERVAL ? SECOND)
AND script_id = ? AND script_content_id != ?`
const loadAffectedHostsPendingExecutionsWithObsoleteScriptUA = `
SELECT
DISTINCT host_id
FROM
upcoming_activities ua
INNER JOIN script_upcoming_activities sua
ON ua.id = sua.upcoming_activity_id
WHERE
ua.activity_type = 'script'
AND ua.activated_at IS NOT NULL
AND (ua.payload->'$.sync_request' = 0 OR ua.created_at >= NOW() - INTERVAL ? SECOND)
AND sua.script_id = ? AND sua.script_content_id != ?`
const clearPendingExecutionsWithObsoleteScriptUA = `DELETE FROM upcoming_activities
USING
upcoming_activities
INNER JOIN script_upcoming_activities sua
ON upcoming_activities.id = sua.upcoming_activity_id
WHERE
upcoming_activities.activity_type = 'script'
AND (upcoming_activities.payload->'$.sync_request' = 0 OR upcoming_activities.created_at >= NOW() - INTERVAL ? SECOND)
AND sua.script_id = ? AND sua.script_content_id != ?`
const loadInsertedScripts = `SELECT id, team_id, name FROM scripts WHERE global_or_team_id = ?`
// use a team id of 0 if no-team
var globalOrTeamID uint
if tmID != nil {
globalOrTeamID = *tmID
}
// build a list of names for the incoming scripts, will keep the
// existing ones if there's a match and no change
incomingNames := make([]string, len(scripts))
// at the same time, index the incoming scripts keyed by name for ease
// of processing
incomingScripts := make(map[string]*fleet.Script, len(scripts))
for i, p := range scripts {
incomingNames[i] = p.Name
incomingScripts[p.Name] = p
}
var insertedScripts []fleet.ScriptResponse
var activateAffectedHosts []uint
if err := ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
var existingScripts []*fleet.Script
if len(incomingNames) > 0 {
// load existing scripts that match the incoming scripts by names
stmt, args, err := sqlx.In(loadExistingScripts, globalOrTeamID, incomingNames)
if err != nil {
return ctxerr.Wrap(ctx, err, "build query to load existing scripts")
}
if err := sqlx.SelectContext(ctx, tx, &existingScripts, stmt, args...); err != nil {
return ctxerr.Wrap(ctx, err, "load existing scripts")
}
}
// figure out if we need to delete any scripts
keepNames := make([]string, 0, len(incomingNames))
for _, p := range existingScripts {
if newS := incomingScripts[p.Name]; newS != nil {
keepNames = append(keepNames, p.Name)
}
}
var (
scriptsStmt string
scriptsArgs []any
policiesStmt string
policiesArgs []any
executionsStmt string
executionsArgs []any
extraExecStmt string
extraExecArgs []any
err error
affectedHostIDs []uint
)
if len(keepNames) > 0 {
// delete the obsolete scripts
scriptsStmt, scriptsArgs, err = sqlx.In(deleteScriptsNotInList, globalOrTeamID, keepNames)
if err != nil {
return ctxerr.Wrap(ctx, err, "build statement to delete obsolete scripts")
}
policiesStmt, policiesArgs, err = sqlx.In(unsetScriptsNotInListFromPolicies, globalOrTeamID, keepNames)
if err != nil {
return ctxerr.Wrap(ctx, err, "build statement to unset obsolete scripts from policies")
}
executionsStmt, executionsArgs, err = sqlx.In(clearPendingExecutionsNotInListHSR, int(constants.MaxServerWaitTime.Seconds()), globalOrTeamID, keepNames)
if err != nil {
return ctxerr.Wrap(ctx, err, "build statement to clear pending script executions from obsolete scripts")
}
loadAffectedStmt, args, err := sqlx.In(loadAffectedHostsPendingExecutionsNotInListUA,
int(constants.MaxServerWaitTime.Seconds()), globalOrTeamID, keepNames)
if err != nil {
return ctxerr.Wrap(ctx, err, "build query to load affected hosts for upcoming script executions")
}
if err := sqlx.SelectContext(ctx, tx, &affectedHostIDs, loadAffectedStmt, args...); err != nil {
return ctxerr.Wrap(ctx, err, "load affected hosts for upcoming script executions")
}
extraExecStmt, extraExecArgs, err = sqlx.In(clearPendingExecutionsNotInListUA, int(constants.MaxServerWaitTime.Seconds()), globalOrTeamID, keepNames)
if err != nil {
return ctxerr.Wrap(ctx, err, "build statement to clear upcoming pending script executions from obsolete scripts")
}
} else {
scriptsStmt = deleteAllScriptsInTeam
scriptsArgs = []any{globalOrTeamID}
policiesStmt = unsetAllScriptsFromPolicies
policiesArgs = []any{globalOrTeamID}
executionsStmt = clearAllPendingExecutionsHSR
executionsArgs = []any{int(constants.MaxServerWaitTime.Seconds()), globalOrTeamID}
if err := sqlx.SelectContext(ctx, tx, &affectedHostIDs,
loadAffectedHostsAllPendingExecutionsUA, int(constants.MaxServerWaitTime.Seconds()), globalOrTeamID); err != nil {
return ctxerr.Wrap(ctx, err, "load affected hosts for upcoming script executions")
}
extraExecStmt = clearAllPendingExecutionsUA
extraExecArgs = []any{int(constants.MaxServerWaitTime.Seconds()), globalOrTeamID}
}
if _, err := tx.ExecContext(ctx, policiesStmt, policiesArgs...); err != nil {
return ctxerr.Wrap(ctx, err, "unset obsolete scripts from policies")
}
if _, err := tx.ExecContext(ctx, executionsStmt, executionsArgs...); err != nil {
return ctxerr.Wrap(ctx, err, "clear obsolete script pending executions")
}
if _, err := tx.ExecContext(ctx, extraExecStmt, extraExecArgs...); err != nil {
return ctxerr.Wrap(ctx, err, "clear obsolete upcoming script pending executions")
}
if _, err := tx.ExecContext(ctx, scriptsStmt, scriptsArgs...); err != nil {
return ctxerr.Wrap(ctx, err, "delete obsolete scripts")
}
activateAffectedHosts = affectedHostIDs
// insert the new scripts and the ones that have changed
for _, s := range incomingScripts {
scRes, err := insertScriptContents(ctx, tx, s.ScriptContents)
if err != nil {
return ctxerr.Wrapf(ctx, err, "inserting script contents for script with name %q", s.Name)
}
contentID, _ := scRes.LastInsertId()
insertRes, err := tx.ExecContext(ctx, insertNewOrEditedScript, tmID, globalOrTeamID, s.Name, uint(contentID)) //nolint:gosec // dismiss G115
if err != nil {
return ctxerr.Wrapf(ctx, err, "insert new/edited script with name %q", s.Name)
}
scriptID, _ := insertRes.LastInsertId()
if _, err := tx.ExecContext(ctx, clearPendingExecutionsWithObsoleteScriptHSR, int(constants.MaxServerWaitTime.Seconds()), scriptID, contentID); err != nil {
return ctxerr.Wrapf(ctx, err, "clear obsolete pending script executions with name %q", s.Name)
}
var affectedHosts []uint
if err := sqlx.SelectContext(ctx, tx, &affectedHosts, loadAffectedHostsPendingExecutionsWithObsoleteScriptUA,
int(constants.MaxServerWaitTime.Seconds()), scriptID, contentID); err != nil {
return ctxerr.Wrapf(ctx, err, "load affected hosts for upcoming script executions with name %q", s.Name)
}
activateAffectedHosts = append(activateAffectedHosts, affectedHosts...)
if _, err = tx.ExecContext(ctx, clearPendingExecutionsWithObsoleteScriptUA, int(constants.MaxServerWaitTime.Seconds()), scriptID, contentID); err != nil {
return ctxerr.Wrapf(ctx, err, "clear obsolete upcoming pending script executions with name %q", s.Name)
}
}
if err := sqlx.SelectContext(ctx, tx, &insertedScripts, loadInsertedScripts, globalOrTeamID); err != nil {
return ctxerr.Wrap(ctx, err, "load inserted scripts")
}
return nil
}); err != nil {
return nil, err
}
if err := ds.activateNextUpcomingActivityForBatchOfHosts(ctx, activateAffectedHosts); err != nil {
return nil, ctxerr.Wrap(ctx, err, "activate next upcoming activity for batch of hosts")
}
return insertedScripts, nil
}
func (ds *Datastore) GetHostLockWipeStatus(ctx context.Context, host *fleet.Host) (*fleet.HostLockWipeStatus, error) {
const stmt = `
SELECT
lock_ref,
wipe_ref,
unlock_ref,
unlock_pin,
fleet_platform
FROM
host_mdm_actions
WHERE
host_id = ?
`
var mdmActions struct {
LockRef *string `db:"lock_ref"`
WipeRef *string `db:"wipe_ref"`
UnlockRef *string `db:"unlock_ref"`
UnlockPIN *string `db:"unlock_pin"`
FleetPlatform string `db:"fleet_platform"`
}
fleetPlatform := host.FleetPlatform()
status := &fleet.HostLockWipeStatus{
HostFleetPlatform: fleetPlatform,
}
if err := sqlx.GetContext(ctx, ds.reader(ctx), &mdmActions, stmt, host.ID); err != nil {
if err == sql.ErrNoRows {
// do not return a Not Found error, return the zero-value status, which
// will report the correct states.
return status, nil
}
return nil, ctxerr.Wrap(ctx, err, "get host lock/wipe status")
}
// if we have a fleet platform stored in host_mdm_actions, use it instead of
// the host.FleetPlatform() because the platform can be overwritten with an
// unknown OS name when a Wipe gets executed.
if mdmActions.FleetPlatform != "" {
fleetPlatform = mdmActions.FleetPlatform
status.HostFleetPlatform = fleetPlatform
}
switch fleetPlatform {
case "darwin", "ios", "ipados":
if mdmActions.UnlockPIN != nil && fleetPlatform == "darwin" {
// Unlock PIN is only available for macOS hosts
status.UnlockPIN = *mdmActions.UnlockPIN
}
if mdmActions.UnlockRef != nil && fleetPlatform == "darwin" {
// the unlock reference is a timestamp
// (we only store the timestamp for macOS unlocks)
var err error
status.UnlockRequestedAt, err = time.Parse(time.DateTime, *mdmActions.UnlockRef)
if err != nil {
// if the format is unexpected but there's something in UnlockRef, just
// replace it with the current timestamp, it should still indicate that
// an unlock was requested (e.g. in case someone plays with the data
// directly in the DB and messes up the format).
status.UnlockRequestedAt = time.Now().UTC()
}
} else if mdmActions.UnlockRef != nil && fleetPlatform != "darwin" {
// the unlock reference is an MDM command uuid
cmd, cmdRes, err := ds.getHostMDMAppleCommand(ctx, *mdmActions.UnlockRef, host.UUID)
if err != nil {
return nil, ctxerr.Wrap(ctx, err, "get unlock reference")
}
status.UnlockMDMCommand = cmd
status.UnlockMDMCommandResult = cmdRes
}
if mdmActions.LockRef != nil {
// the lock reference is an MDM command
cmd, cmdRes, err := ds.getHostMDMAppleCommand(ctx, *mdmActions.LockRef, host.UUID)
if err != nil {
return nil, ctxerr.Wrap(ctx, err, "get lock reference")
}
status.LockMDMCommand = cmd
status.LockMDMCommandResult = cmdRes
}
if mdmActions.WipeRef != nil {
// the wipe reference is an MDM command
cmd, cmdRes, err := ds.getHostMDMAppleCommand(ctx, *mdmActions.WipeRef, host.UUID)
if err != nil {
return nil, ctxerr.Wrap(ctx, err, "get wipe reference")
}
status.WipeMDMCommand = cmd
status.WipeMDMCommandResult = cmdRes
}
case "windows", "linux":
// lock and unlock references are scripts
if mdmActions.LockRef != nil {
hsr, err := ds.getHostScriptExecutionResultDB(ctx, ds.reader(ctx), *mdmActions.LockRef, scriptExecutionSearchOpts{IncludeCanceled: true})
if err != nil {
return nil, ctxerr.Wrap(ctx, err, "get lock reference script result")
}
status.LockScript = hsr
}
if mdmActions.UnlockRef != nil {
hsr, err := ds.getHostScriptExecutionResultDB(ctx, ds.reader(ctx), *mdmActions.UnlockRef, scriptExecutionSearchOpts{IncludeCanceled: true})
if err != nil {
return nil, ctxerr.Wrap(ctx, err, "get unlock reference script result")
}
status.UnlockScript = hsr
}
// wipe is an MDM command on Windows, a script on Linux
if mdmActions.WipeRef != nil {
if fleetPlatform == "windows" {
cmd, cmdRes, err := ds.getHostMDMWindowsCommand(ctx, *mdmActions.WipeRef, host.UUID)
if err != nil {
return nil, ctxerr.Wrap(ctx, err, "get wipe reference")
}
status.WipeMDMCommand = cmd
status.WipeMDMCommandResult = cmdRes
} else {
hsr, err := ds.getHostScriptExecutionResultDB(ctx, ds.reader(ctx), *mdmActions.WipeRef, scriptExecutionSearchOpts{IncludeCanceled: true})
if err != nil {
return nil, ctxerr.Wrap(ctx, err, "get wipe reference script result")
}
status.WipeScript = hsr
}
}
}
return status, nil
}
func (ds *Datastore) getHostMDMWindowsCommand(ctx context.Context, cmdUUID, hostUUID string) (*fleet.MDMCommand, *fleet.MDMCommandResult, error) {
cmd, err := ds.getMDMCommand(ctx, ds.reader(ctx), cmdUUID)
if err != nil {
return nil, nil, ctxerr.Wrap(ctx, err, "get Windows MDM command")
}
// get the MDM command result, which may be not found (indicating the command doesn't exist).
// If it is pending, then it returns 101, and result will be empty.
cmdResults, err := ds.GetMDMWindowsCommandResults(ctx, cmdUUID, "")
if err != nil {
return nil, nil, ctxerr.Wrap(ctx, err, "get Windows MDM command result")
}
// each item in the slice returned by GetMDMWindowsCommandResults is
// potentially a result for a different host, we need to find the one for
// that specific host.
var cmdRes *fleet.MDMCommandResult
for _, r := range cmdResults {
if r.HostUUID != hostUUID {
continue
}
if r.Status == "101" || string(r.Result) == "" {
// command is still pending
continue
}
// all statuses for Windows indicate end of processing of the command
// (there is no equivalent of "NotNow" or "Idle" as for Apple).
cmdRes = r
break
}
return cmd, cmdRes, nil
}
func (ds *Datastore) getHostMDMAppleCommand(ctx context.Context, cmdUUID, hostUUID string) (*fleet.MDMCommand, *fleet.MDMCommandResult, error) {
cmd, err := ds.getMDMCommand(ctx, ds.reader(ctx), cmdUUID)
if err != nil {
return nil, nil, ctxerr.Wrap(ctx, err, "get Apple MDM command")
}
// get the MDM command result, which may be not found (indicating the command
// is pending). Note that it doesn't return ErrNoRows if not found, it
// returns success and an empty cmdRes slice.
cmdResults, err := ds.GetMDMAppleCommandResults(ctx, cmdUUID, hostUUID)
if err != nil {
return nil, nil, ctxerr.Wrap(ctx, err, "get Apple MDM command result")
}
// filter by result status to preserve old behavior of this method where it doesn't return pending results.
var cmdRes *fleet.MDMCommandResult
for _, r := range cmdResults {
if r.HostUUID != hostUUID {
// this should never happen because we already filter by hostUUID, but just in case
continue
}
if r.Status == fleet.MDMAppleStatusAcknowledged || r.Status == fleet.MDMAppleStatusError || r.Status == fleet.MDMAppleStatusCommandFormatError {
cmdRes = r
break
}
}
return cmd, cmdRes, nil
}
// LockHostViaScript will create the script execution request and update
// host_mdm_actions in a single transaction.
func (ds *Datastore) LockHostViaScript(ctx context.Context, request *fleet.HostScriptRequestPayload, hostFleetPlatform string) error {
var res *fleet.HostScriptResult
return ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
var err error
scRes, err := insertScriptContents(ctx, tx, request.ScriptContents)
if err != nil {
return err
}
id, _ := scRes.LastInsertId()
request.ScriptContentID = uint(id) //nolint:gosec // dismiss G115
res, err = ds.newHostScriptExecutionRequest(ctx, tx, request, true)
if err != nil {
return ctxerr.Wrap(ctx, err, "lock host via script create execution")
}
// on duplicate we don't clear any other existing state because at this
// point in time, this is just a request to lock the host that is recorded,
// it is pending execution. The host's state should be updated to "locked"
// only when the script execution is successfully completed, and then any
// unlock or wipe references should be cleared.
const stmt = `
INSERT INTO host_mdm_actions
(
host_id,
lock_ref,
fleet_platform
)
VALUES (?,?,?)
ON DUPLICATE KEY UPDATE
lock_ref = VALUES(lock_ref)
`
_, err = tx.ExecContext(ctx, stmt,
request.HostID,
res.ExecutionID,
hostFleetPlatform,
)
if err != nil {
return ctxerr.Wrap(ctx, err, "lock host via script update mdm actions")
}
return nil
})
}
// UnlockHostViaScript will create the script execution request and update
// host_mdm_actions in a single transaction.
func (ds *Datastore) UnlockHostViaScript(ctx context.Context, request *fleet.HostScriptRequestPayload, hostFleetPlatform string) error {
var res *fleet.HostScriptResult
return ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
var err error
scRes, err := insertScriptContents(ctx, tx, request.ScriptContents)
if err != nil {
return err
}
id, _ := scRes.LastInsertId()
request.ScriptContentID = uint(id) //nolint:gosec // dismiss G115
res, err = ds.newHostScriptExecutionRequest(ctx, tx, request, true)
if err != nil {
return ctxerr.Wrap(ctx, err, "unlock host via script create execution")
}
// on duplicate we don't clear any other existing state because at this
// point in time, this is just a request to unlock the host that is
// recorded, it is pending execution. The host's state should be updated to
// "unlocked" only when the script execution is successfully completed, and
// then any lock or wipe references should be cleared.
const stmt = `
INSERT INTO host_mdm_actions
(
host_id,
unlock_ref,
fleet_platform
)
VALUES (?,?,?)
ON DUPLICATE KEY UPDATE
unlock_ref = VALUES(unlock_ref),
unlock_pin = NULL
`
_, err = tx.ExecContext(ctx, stmt,
request.HostID,
res.ExecutionID,
hostFleetPlatform,
)
if err != nil {
return ctxerr.Wrap(ctx, err, "unlock host via script update mdm actions")
}
return err
})
}
// WipeHostViaScript creates the script execution request and updates the
// host_mdm_actions table in a single transaction.
func (ds *Datastore) WipeHostViaScript(ctx context.Context, request *fleet.HostScriptRequestPayload, hostFleetPlatform string) error {
var res *fleet.HostScriptResult
return ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
var err error
scRes, err := insertScriptContents(ctx, tx, request.ScriptContents)
if err != nil {
return err
}
id, _ := scRes.LastInsertId()
request.ScriptContentID = uint(id) //nolint:gosec // dismiss G115
res, err = ds.newHostScriptExecutionRequest(ctx, tx, request, true)
if err != nil {
return ctxerr.Wrap(ctx, err, "wipe host via script create execution")
}
// on duplicate we don't clear any other existing state because at this
// point in time, this is just a request to wipe the host that is recorded,
// it is pending execution, so if it was locked, it is still locked (so the
// lock_ref info must still be there).
const stmt = `
INSERT INTO host_mdm_actions
(
host_id,
wipe_ref,
fleet_platform
)
VALUES (?,?,?)
ON DUPLICATE KEY UPDATE
wipe_ref = VALUES(wipe_ref)
`
_, err = tx.ExecContext(ctx, stmt,
request.HostID,
res.ExecutionID,
hostFleetPlatform,
)
if err != nil {
return ctxerr.Wrap(ctx, err, "wipe host via script update mdm actions")
}
return err
})
}
func (ds *Datastore) UnlockHostManually(ctx context.Context, hostID uint, hostFleetPlatform string, ts time.Time) error {
const stmt = `
INSERT INTO host_mdm_actions
(
host_id,
unlock_ref,
fleet_platform
)
VALUES (?, ?, ?)
ON DUPLICATE KEY UPDATE
-- do not overwrite if a value is already set
unlock_ref = IF(unlock_ref IS NULL, VALUES(unlock_ref), unlock_ref)
`
// for macOS, the unlock_ref is just the timestamp at which the user first
// requested to unlock the host. This then indicates in the host's status
// that it's pending an unlock (which requires manual intervention by
// entering a PIN on the device). The /unlock endpoint can be called multiple
// times, so we record the timestamp of the first time it was requested and
// from then on, the host is marked as "pending unlock" until the device is
// actually unlocked with the PIN. The actual unlocking happens when the
// device sends an Idle MDM request.
unlockRef := ts.Format(time.DateTime)
_, err := ds.writer(ctx).ExecContext(ctx, stmt, hostID, unlockRef, hostFleetPlatform)
return ctxerr.Wrap(ctx, err, "record manual unlock host request")
}
func buildHostLockWipeStatusUpdateStmt(refCol string, succeeded bool, joinPart string, setUnlockRef bool) string {
var alias string
stmt := `UPDATE host_mdm_actions `
if joinPart != "" {
stmt += `hma ` + joinPart
alias = "hma."
}
stmt += ` SET `
if succeeded {
switch refCol {
case "lock_ref":
// Note that this must not clear the unlock_pin, because recording the
// lock request does generate the PIN and store it there to be used by an
// eventual unlock.
if !setUnlockRef {
stmt += fmt.Sprintf("%sunlock_ref = NULL, %[1]swipe_ref = NULL", alias)
} else {
// Currently only used for Apple MDM devices.
// We set the unlock_ref to current time since the device can be unlocked any time after the lock.
// Apple MDM does not have a concept of unlock pending.
stmt += fmt.Sprintf("%sunlock_ref = '%s', %[1]swipe_ref = NULL", alias, time.Now().Format(time.DateTime))
}
case "unlock_ref":
// a successful unlock clears itself as well as the lock ref, because
// unlock is the default state so we don't need to keep its unlock_ref
// around once it's confirmed.
stmt += fmt.Sprintf("%slock_ref = NULL, %[1]sunlock_ref = NULL, %[1]sunlock_pin = NULL, %[1]swipe_ref = NULL", alias)
case "wipe_ref":
stmt += fmt.Sprintf("%slock_ref = NULL, %[1]sunlock_ref = NULL, %[1]sunlock_pin = NULL", alias)
}
} else {
// if the action failed, then we clear the reference to that action itself so
// the host stays in the previous state (it doesn't transition to the new
// state).
stmt += fmt.Sprintf("%s"+refCol+" = NULL", alias)
}
return stmt
}
func (ds *Datastore) UpdateHostLockWipeStatusFromAppleMDMResult(ctx context.Context, hostUUID, cmdUUID, requestType string, succeeded bool) error {
// a bit of MDM protocol leaking in the mysql layer, but it's either that or
// the other way around (MDM protocol would translate to database column)
var refCol string
var setUnlockRef bool
switch requestType {
case "EraseDevice":
refCol = "wipe_ref"
case "DeviceLock":
refCol = "lock_ref"
setUnlockRef = true
case "EnableLostMode":
refCol = "lock_ref"
case "DisableLostMode":
refCol = "unlock_ref"
default:
return nil
}
return updateHostLockWipeStatusFromResultAndHostUUID(ctx, ds.writer(ctx), hostUUID, refCol, cmdUUID, succeeded, setUnlockRef)
}
func updateHostLockWipeStatusFromResultAndHostUUID(
ctx context.Context, tx sqlx.ExtContext, hostUUID, refCol, cmdUUID string, succeeded bool, setUnlockRef bool,
) error {
stmt := buildHostLockWipeStatusUpdateStmt(refCol, succeeded, `JOIN hosts h ON hma.host_id = h.id`, setUnlockRef)
stmt += ` WHERE h.uuid = ? AND hma.` + refCol + ` = ?`
_, err := tx.ExecContext(ctx, stmt, hostUUID, cmdUUID)
return ctxerr.Wrap(ctx, err, "update host lock/wipe status from result via host uuid")
}
func updateHostLockWipeStatusFromResult(ctx context.Context, tx sqlx.ExtContext, hostID uint, refCol string, succeeded bool) error {
stmt := buildHostLockWipeStatusUpdateStmt(refCol, succeeded, "", false)
stmt += ` WHERE host_id = ?`
_, err := tx.ExecContext(ctx, stmt, hostID)
return ctxerr.Wrap(ctx, err, "update host lock/wipe status from result")
}
func (ds *Datastore) updateUninstallStatusFromResult(ctx context.Context, tx sqlx.ExtContext, hostID uint, executionID string, exitCode int) error {
stmt := `
UPDATE host_software_installs SET uninstall_script_exit_code = ? WHERE execution_id = ? AND host_id = ?
`
if _, err := tx.ExecContext(ctx, stmt, exitCode, executionID, hostID); err != nil {
return ctxerr.Wrap(ctx, err, "update uninstall status from result")
}
// NOTE: no need to call activateNextUpcomingActivity here as this function
// is called from SetHostScriptExecutionResult which will call it before
// completing.
return nil
}
func (ds *Datastore) CleanupUnusedScriptContents(ctx context.Context) error {
deleteStmt := `
DELETE FROM
script_contents
WHERE
NOT EXISTS (
SELECT 1 FROM host_script_results WHERE script_content_id = script_contents.id)
AND NOT EXISTS (
SELECT 1 FROM scripts WHERE script_content_id = script_contents.id)
AND NOT EXISTS (
SELECT 1 FROM software_installers si
WHERE script_contents.id IN (si.install_script_content_id, si.post_install_script_content_id, si.uninstall_script_content_id)
)
AND NOT EXISTS (
SELECT 1 FROM setup_experience_scripts WHERE script_content_id = script_contents.id
)
AND NOT EXISTS (
SELECT 1 FROM script_upcoming_activities WHERE script_content_id = script_contents.id
)
`
_, err := ds.writer(ctx).ExecContext(ctx, deleteStmt)
if err != nil {
return ctxerr.Wrap(ctx, err, "cleaning up unused script contents")
}
return nil
}
func (ds *Datastore) getOrGenerateScriptContentsID(ctx context.Context, contents string) (uint, error) {
csum := md5ChecksumScriptContent(contents)
scriptContentsID, err := ds.optimisticGetOrInsert(ctx,
&parameterizedStmt{
Statement: `SELECT id FROM script_contents WHERE md5_checksum = UNHEX(?)`,
Args: []interface{}{csum},
},
&parameterizedStmt{
Statement: `INSERT INTO script_contents (md5_checksum, contents) VALUES (UNHEX(?), ?)`,
Args: []interface{}{csum, contents},
},
)
if err != nil {
return 0, err
}
return scriptContentsID, nil
}
func teamIDEq(teamID1, teamID2 *uint) bool {
sameTeamNoTeam := teamID1 == nil && teamID2 == nil
sameTeamNumber := teamID1 != nil && teamID2 != nil && *teamID1 == *teamID2
return sameTeamNoTeam || sameTeamNumber
}
func (ds *Datastore) batchExecuteScript(ctx context.Context, userID *uint, scriptID uint, hostIDs []uint, batchExecID string) error {
script, err := ds.Script(ctx, scriptID)
if err != nil {
return fleet.NewInvalidArgumentError("script_id", err.Error())
}
invalidHostIDPlatform := "batch-invalid-hostid"
// We need full host info to check if hosts are able to run scripts, see svc.RunHostScript
fullHosts := make([]*fleet.Host, 0, len(hostIDs))
// The execution results to be stored in the database
executions := make([]fleet.BatchExecutionHost, 0, len(fullHosts))
// Check that all hosts exist before attempting to process them
for _, hostID := range hostIDs {
host, err := ds.Host(ctx, hostID)
if err != nil {
fullHosts = append(fullHosts, &fleet.Host{
ID: hostID,
Platform: invalidHostIDPlatform,
})
continue
}
fullHosts = append(fullHosts, host)
}
if err := ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
for _, host := range fullHosts {
// Host doesn't exist anymore
if host.Platform == invalidHostIDPlatform {
executions = append(executions, fleet.BatchExecutionHost{
HostID: host.ID,
Error: &fleet.BatchExecuteInvalidHost,
})
continue
}
// Non-orbit-enrolled host (iOS, android)
noNodeKey := host.OrbitNodeKey == nil || *host.OrbitNodeKey == ""
// Scripts disabled on host
scriptsDisabled := host.ScriptsEnabled != nil && !*host.ScriptsEnabled
if noNodeKey || scriptsDisabled {
executions = append(executions, fleet.BatchExecutionHost{
HostID: host.ID,
Error: &fleet.BatchExecuteIncompatibleFleetd,
})
continue
}
if !fleet.ValidateScriptPlatform(script.Name, host.Platform) {
executions = append(executions, fleet.BatchExecutionHost{
HostID: host.ID,
Error: &fleet.BatchExecuteIncompatiblePlatform,
})
continue
}
executionID, _, err := ds.insertNewHostScriptExecution(ctx, tx, &fleet.HostScriptRequestPayload{
HostID: host.ID,
UserID: userID,
ScriptID: &script.ID,
ScriptContentID: script.ScriptContentID,
}, false)
if err != nil {
return ctxerr.Wrap(ctx, err, "queueing script for bulk execution")
}
executions = append(executions, fleet.BatchExecutionHost{
HostID: host.ID,
ExecutionID: &executionID,
})
}
_, err := tx.ExecContext(
ctx,
`INSERT INTO batch_activities (execution_id, script_id, status, activity_type, num_targeted, started_at) VALUES (?, ?, ?, ?, ?, NOW())
ON DUPLICATE KEY UPDATE status = VALUES(status), started_at = VALUES(started_at)`,
batchExecID,
script.ID,
fleet.ScheduledBatchExecutionStarted,
fleet.BatchExecutionActivityScript,
len(hostIDs),
)
if err != nil {
return ctxerr.Wrap(ctx, err, "failed to insert new batch execution")
}
args := make([]map[string]any, 0, len(executions))
for _, execHost := range executions {
args = append(args, map[string]any{
"batch_id": batchExecID,
"host_id": execHost.HostID,
"host_execution_id": execHost.ExecutionID,
"error": execHost.Error,
})
}
insertStmt := `
INSERT INTO batch_activity_host_results (
batch_execution_id,
host_id,
host_execution_id,
error
) VALUES (
:batch_id,
:host_id,
:host_execution_id,
:error
) ON DUPLICATE KEY UPDATE host_execution_id = VALUES(host_execution_id), error = VALUES(error)`
if _, err := sqlx.NamedExecContext(ctx, tx, insertStmt, args); err != nil {
return ctxerr.Wrap(ctx, err, "associating script executions with batch job")
}
return nil
}); err != nil {
return fmt.Errorf("creating bulk execution order: %w", err)
}
return nil
}
func (ds *Datastore) BatchExecuteScript(ctx context.Context, userID *uint, scriptID uint, hostIDs []uint) (string, error) {
batchExecID := uuid.New().String()
script, err := ds.Script(ctx, scriptID)
if err != nil {
return "", fleet.NewInvalidArgumentError("script_id", err.Error())
}
for _, hostID := range hostIDs {
host, err := ds.HostLite(ctx, hostID)
if err != nil {
return "", fmt.Errorf("unable to load host information for %d: %w", hostID, err)
}
if !teamIDEq(host.TeamID, script.TeamID) {
return "", ctxerr.Errorf(ctx, "all hosts must be on the same team as the script")
}
}
if err := ds.batchExecuteScript(ctx, userID, scriptID, hostIDs, batchExecID); err != nil {
return "", ctxerr.Wrap(ctx, err, "immediate batch execution")
}
return batchExecID, nil
}
func (ds *Datastore) BatchScheduleScript(ctx context.Context, userID *uint, scriptID uint, hostIDs []uint, notBefore time.Time) (string, error) {
batchExecID := uuid.New().String()
const batchActivitiesStmt = `INSERT INTO batch_activities (execution_id, job_id, script_id, user_id, status, activity_type, num_targeted) VALUES (?, ?, ?, ?, ?, ?, ?)`
const batchHostsStmt = `INSERT INTO batch_activity_host_results (batch_execution_id, host_id) VALUES (:exec_id, :host_id)`
argBytes, err := json.Marshal(fleet.BatchActivityScriptJobArgs{
ExecutionID: batchExecID,
})
if err != nil {
return "", ctxerr.Wrap(ctx, err, "encooding job args")
}
if err := ds.withTx(ctx, func(tx sqlx.ExtContext) error {
job, err := ds.NewJob(ctx, &fleet.Job{
Name: fleet.BatchActivityScriptsJobName,
Args: (*json.RawMessage)(&argBytes),
State: fleet.JobStateQueued,
NotBefore: notBefore.UTC(),
})
if err != nil {
return ctxerr.Wrap(ctx, err, "creating new job")
}
_, err = tx.ExecContext(
ctx,
batchActivitiesStmt,
batchExecID,
job.ID,
scriptID,
userID,
fleet.ScheduledBatchExecutionScheduled,
fleet.BatchExecutionActivityScript,
len(hostIDs),
)
if err != nil {
return ctxerr.Wrap(ctx, err, "inserting new batch activity")
}
args := make([]map[string]any, 0, len(hostIDs))
for _, hostID := range hostIDs {
args = append(args, map[string]any{
"exec_id": batchExecID,
"host_id": hostID,
})
}
if _, err := sqlx.NamedExecContext(ctx, tx, batchHostsStmt, args); err != nil {
return ctxerr.Wrap(ctx, err, "inserting batch host results")
}
return nil
}); err != nil {
return "", ctxerr.Wrap(ctx, err, "creating scheduled script execution")
}
return batchExecID, nil
}
func (ds *Datastore) CancelBatchScript(ctx context.Context, executionID string) error {
stmt := `
SELECT
bahr.host_execution_id,
bahr.host_id
FROM
batch_activity_host_results bahr
LEFT JOIN
host_script_results hsr ON bahr.host_execution_id = hsr.execution_id -- I think?
WHERE
bahr.batch_execution_id = ?
AND
hsr.canceled = 0
AND
hsr.exit_code IS NULL
AND
bahr.error IS NULL`
stmtSetCanceled := `
UPDATE
batch_activities ba
SET
finished_at = NOW(),
status = 'finished',
canceled = 1,
num_canceled = (SELECT COUNT(*) FROM batch_activity_host_results WHERE batch_execution_id = ba.execution_id)
WHERE
ba.execution_id = ?`
stmtCanceled := `
UPDATE
batch_activities
SET
canceled = 1
WHERE
execution_id = ?`
activity, err := ds.GetBatchActivity(ctx, executionID)
if err != nil {
return ctxerr.Wrap(ctx, err, "getting batch activity")
}
if activity.Status == fleet.ScheduledBatchExecutionFinished {
return nil
}
if err := ds.withRetryTxx(ctx, func(tx sqlx.ExtContext) error {
// If job worker exists, mark it as complete to stop it from running
if jobID := activity.JobID; jobID != nil {
job, err := ds.GetJob(ctx, *jobID)
if err != nil {
return ctxerr.Wrap(ctx, err, "failed to find job associated with batch activity")
}
job.State = fleet.JobStateSuccess
if _, err := ds.updateJob(ctx, tx, *jobID, job); err != nil {
return ctxerr.Wrap(ctx, err, "updating batch activity job")
}
}
if activity.Status == fleet.ScheduledBatchExecutionStarted {
// If the batch activity has started, we need to cancel anything in progress or queued
toCancel := []struct {
HostExecutionID string `db:"host_execution_id"`
HostID uint `db:"host_id"`
}{}
if err := sqlx.SelectContext(ctx, tx, &toCancel, stmt, executionID); err != nil {
return ctxerr.Wrap(ctx, err, "selecting hosts to cancel")
}
for _, host := range toCancel {
if _, err := ds.cancelHostUpcomingActivity(ctx, tx, host.HostID, host.HostExecutionID); err != nil {
return ctxerr.Wrap(ctx, err, "canceling upcoming activity")
}
}
if _, err := tx.ExecContext(ctx, stmtCanceled, executionID); err != nil {
return ctxerr.Wrap(ctx, err, "setting canceled column")
}
if err := ds.markActivitiesAsCompleted(ctx, tx); err != nil {
return ctxerr.Wrap(ctx, err, "marking job as complete and summarizing counts")
}
} else {
// The batch activity is scheduled, but not started
if _, err := tx.ExecContext(ctx, stmtSetCanceled, executionID); err != nil {
return ctxerr.Wrap(ctx, err, "setting canceled host count")
}
}
return nil
}); err != nil {
return ctxerr.Wrap(ctx, err, "cancel batch script db transaction")
}
return nil
}
func (ds *Datastore) GetBatchActivity(ctx context.Context, executionID string) (*fleet.BatchActivity, error) {
const stmt = `
SELECT
ba.id,
ba.script_id,
s.name as script_name,
ba.execution_id,
ba.user_id,
ba.job_id,
ba.status,
ba.activity_type,
ba.num_targeted,
ba.num_pending,
ba.num_ran,
ba.num_errored,
ba.num_incompatible,
ba.num_canceled,
ba.created_at,
ba.updated_at,
ba.started_at,
ba.finished_at,
ba.canceled
FROM
batch_activities ba
LEFT JOIN
scripts s ON s.id = ba.script_id
WHERE
execution_id = ?`
batchActivity := &fleet.BatchActivity{}
if err := sqlx.GetContext(ctx, ds.reader(ctx), batchActivity, stmt, executionID); err != nil {
return nil, ctxerr.Wrap(ctx, err, "selecting batch activity")
}
return batchActivity, nil
}
func (ds *Datastore) GetBatchActivityHostResults(ctx context.Context, executionID string) ([]*fleet.BatchActivityHostResult, error) {
const stmt = `
SELECT
id,
batch_execution_id,
host_id,
host_execution_id,
error
FROM
batch_activity_host_results
WHERE
batch_execution_id = ?`
results := []*fleet.BatchActivityHostResult{}
if err := sqlx.SelectContext(ctx, ds.reader(ctx), &results, stmt, executionID); err != nil {
return nil, ctxerr.Wrap(ctx, err, "selecting batch activity host results")
}
return results, nil
}
func (ds *Datastore) RunScheduledBatchActivity(ctx context.Context, executionID string) error {
batchActivity, err := ds.GetBatchActivity(ctx, executionID)
if err != nil {
return ctxerr.Wrap(ctx, err, "getting batch activity")
}
if batchActivity.Status != fleet.ScheduledBatchExecutionScheduled {
return ctxerr.New(ctx, "batch job has already been started")
}
if batchActivity.Canceled {
return ctxerr.New(ctx, "batch job was canceled")
}
if batchActivity.ScriptID == nil {
return ctxerr.New(ctx, "no script ID present in batch activity")
}
script, err := ds.Script(ctx, *batchActivity.ScriptID)
if err != nil {
return ctxerr.Wrap(ctx, err, "could not get script")
}
results, err := ds.GetBatchActivityHostResults(ctx, executionID)
if err != nil {
return ctxerr.Wrap(ctx, err, "getting batch activity host results")
}
hostIDs := []uint{}
for _, result := range results {
hostIDs = append(hostIDs, result.HostID)
}
if err := ds.batchExecuteScript(ctx, batchActivity.UserID, script.ID, hostIDs, batchActivity.BatchExecutionID); err != nil {
return ctxerr.Wrap(ctx, err, "scheduled batch script execution")
}
return nil
}
// Deprecated; will be removed in favor of ListBatchScriptExecutions when the batch script details page is ready.
func (ds *Datastore) BatchExecuteSummary(ctx context.Context, executionID string) (*fleet.BatchActivity, error) {
stmtExecutions := `
SELECT
COUNT(*) as num_targeted,
COUNT(bsehr.error) as num_did_not_run,
COUNT(CASE WHEN hsr.exit_code = 0 THEN 1 END) as num_succeeded,
COUNT(CASE WHEN hsr.exit_code <> 0 THEN 1 END) as num_failed,
COUNT(CASE WHEN hsr.canceled = 1 AND hsr.exit_code IS NULL THEN 1 END) as num_cancelled
FROM
batch_activity_host_results bsehr
LEFT JOIN
host_script_results hsr
ON bsehr.host_execution_id = hsr.execution_id
WHERE
bsehr.batch_execution_id = ?`
stmtScriptDetails := `
SELECT
script_id,
s.name as script_name,
s.team_id as team_id,
bse.created_at as created_at
FROM
batch_activities bse
JOIN
scripts s
ON bse.script_id = s.id
WHERE
bse.execution_id = ?`
var summary fleet.BatchActivity
var temp_summary struct {
NumTargeted uint `db:"num_targeted"`
NumDidNotRun uint `db:"num_did_not_run"`
NumSucceeded uint `db:"num_succeeded"`
NumFailed uint `db:"num_failed"`
NumCancelled uint `db:"num_cancelled"`
}
// Fill out the execution details
if err := sqlx.GetContext(ctx, ds.reader(ctx), &temp_summary, stmtExecutions, executionID); err != nil {
return nil, ctxerr.Wrap(ctx, err, "selecting execution information for bulk execution summary")
}
summary.NumTargeted = &temp_summary.NumTargeted
// NumRan is the number of hosts that actually ran the script successfully.
summary.NumRan = &temp_summary.NumSucceeded
// NumErrored is the number of hosts that errored out, which includes
// both failed and did not run.
summary.NumErrored = ptr.Uint(temp_summary.NumFailed + temp_summary.NumDidNotRun)
// NumFailed is the number of hosts that were canceled before execution.
summary.NumCanceled = &temp_summary.NumCancelled
// NumPending is the number of hosts that are pending execution.
summary.NumPending = ptr.Uint(temp_summary.NumTargeted - (temp_summary.NumSucceeded + temp_summary.NumFailed + temp_summary.NumDidNotRun + temp_summary.NumCancelled))
// Fill out the script details
if err := sqlx.GetContext(ctx, ds.reader(ctx), &summary, stmtScriptDetails, executionID); err != nil {
return nil, ctxerr.Wrap(ctx, err, "selecting script information for bulk execution summary")
}
if summary.TeamID == nil {
summary.TeamID = ptr.Uint(0)
}
return &summary, nil
}
func (ds *Datastore) ListBatchScriptExecutions(ctx context.Context, filter fleet.BatchExecutionStatusFilter) ([]fleet.BatchActivity, error) {
stmtExecutions := `
SELECT *
FROM (
-- If batch is finished, get the cached host result counts
SELECT
COALESCE(ba.num_targeted, 0) AS num_targeted,
COALESCE(ba.num_incompatible, 0) AS num_incompatible,
COALESCE(ba.num_ran, 0) AS num_ran,
COALESCE(ba.num_errored, 0) AS num_errored,
COALESCE(ba.num_canceled, 0) AS num_canceled,
COALESCE(ba.num_pending, 0) AS num_pending,
ba.execution_id,
ba.script_id,
ba.status,
ba.canceled,
ba.finished_at,
ba.started_at,
s.name AS script_name,
s.global_or_team_id AS team_id,
ba.created_at AS created_at,
j.not_before AS not_before,
ba.id AS id
FROM batch_activities ba
JOIN scripts s ON ba.script_id = s.id
LEFT JOIN jobs j ON j.id = ba.job_id
WHERE ( %s ) AND ba.status = 'finished'
UNION ALL
-- If batch is not finished, calculate the host result counts live.
SELECT
COUNT(bahr.host_id) AS num_targeted,
COUNT(bahr.error) AS num_incompatible,
COUNT(IF(hsr.exit_code = 0, 1, NULL)) AS num_ran,
COUNT(IF(hsr.exit_code <> 0, 1, NULL)) AS num_errored,
COUNT(IF((hsr.canceled = 1 AND hsr.exit_code IS NULL) OR (hsr.host_id IS NULL AND bahr.error is NULL AND ba.canceled = 1), 1, NULL)) AS num_cancelled,
(
COUNT(bahr.host_id)
- COUNT(bahr.error)
- COUNT(IF(hsr.exit_code = 0, 1, NULL))
- COUNT(IF(hsr.exit_code <> 0, 1, NULL))
- COUNT(IF((hsr.canceled = 1 AND hsr.exit_code IS NULL) OR (hsr.host_id IS NULL AND bahr.error is NULL AND ba.canceled = 1), 1, NULL))
) AS num_pending,
ba.execution_id,
ba.script_id,
ba.status,
ba.canceled,
ba.finished_at,
ba.started_at,
s.name AS script_name,
s.global_or_team_id AS team_id,
ba.created_at AS created_at,
j.not_before AS not_before,
ba.id AS id
FROM batch_activities ba
LEFT JOIN batch_activity_host_results bahr
ON ba.execution_id = bahr.batch_execution_id
LEFT JOIN host_script_results hsr
ON bahr.host_execution_id = hsr.execution_id
JOIN scripts s
ON ba.script_id = s.id
LEFT JOIN jobs j
ON j.id = ba.job_id
WHERE ( %s ) AND ba.status <> 'finished'
GROUP BY ba.id
) AS u
ORDER BY
%s
LIMIT %d OFFSET %d
`
limit := 10
offset := 0
args := []any{}
orderBy := []string{"u.created_at DESC", "u.id DESC"}
whereClauses := make([]string, 0, 2)
// If an execution ID is provided, use it to filter the results.
if filter.ExecutionID != nil && *filter.ExecutionID != "" {
whereClauses = append(whereClauses, "ba.execution_id = ?")
args = append(args, *filter.ExecutionID)
} else {
// Otherwise filter by status and/or team ID.
if filter.Status != nil && *filter.Status != "" {
whereClauses = append(whereClauses, "ba.status = ?")
args = append(args, *filter.Status)
switch *filter.Status {
case string(fleet.ScheduledBatchExecutionScheduled):
orderBy = append([]string{"u.not_before ASC"}, orderBy...)
case string(fleet.ScheduledBatchExecutionStarted):
orderBy = append([]string{"u.started_at DESC"}, orderBy...)
case string(fleet.ScheduledBatchExecutionFinished):
orderBy = append([]string{"u.finished_at DESC"}, orderBy...)
default:
// no additional ordering
}
}
if filter.TeamID != nil {
whereClauses = append(whereClauses, "s.global_or_team_id = ?")
args = append(args, *filter.TeamID)
}
}
// Double up the args to use them in both WHERE clauses.
args = append(args, args...)
// Use pagination parameters if provided.
if filter.Limit != nil {
limit = int(*filter.Limit) //nolint:gosec // dismiss G115
}
if filter.Offset != nil {
offset = int(*filter.Offset) //nolint:gosec // dismiss G115
}
where := strings.Join(whereClauses, " AND ")
stmtExecutions = fmt.Sprintf(stmtExecutions, where, where, strings.Join(orderBy, ", "), limit, offset)
var summary []fleet.BatchActivity
if err := sqlx.SelectContext(ctx, ds.reader(ctx), &summary, stmtExecutions, args...); err != nil {
return nil, ctxerr.Wrap(ctx, err, "selecting execution information for bulk execution summary")
}
return summary, nil
}
func (ds *Datastore) CountBatchScriptExecutions(ctx context.Context, filter fleet.BatchExecutionStatusFilter) (int64, error) {
stmtExecutions := `
SELECT
COUNT(*)
FROM
batch_activities ba
JOIN
scripts s
ON ba.script_id = s.id
WHERE
%s
`
args := []any{}
whereClauses := make([]string, 0, 2)
if filter.Status != nil && *filter.Status != "" {
whereClauses = append(whereClauses, "ba.status = ?")
args = append(args, *filter.Status)
}
if filter.TeamID != nil {
whereClauses = append(whereClauses, "s.global_or_team_id = ?")
args = append(args, *filter.TeamID)
}
where := strings.Join(whereClauses, " AND ")
stmtExecutions = fmt.Sprintf(stmtExecutions, where)
var count int64
if err := sqlx.GetContext(ctx, ds.reader(ctx), &count, stmtExecutions, args...); err != nil {
return 0, ctxerr.Wrap(ctx, err, "selecting execution information for bulk execution summary")
}
return count, nil
}
func (ds *Datastore) markActivitiesAsCompleted(ctx context.Context, tx sqlx.ExtContext) error {
const stmt = `
UPDATE batch_activities AS ba
JOIN (
SELECT
ba2.id AS batch_id,
COUNT(bahr.host_id) AS num_targeted,
COUNT(bahr.error) AS num_incompatible,
COUNT(IF(hsr.exit_code = 0, 1, NULL)) AS num_ran,
COUNT(IF(hsr.exit_code <> 0, 1, NULL)) AS num_errored,
COUNT(IF((hsr.canceled = 1 AND hsr.exit_code IS NULL) OR (hsr.host_id IS NULL AND bahr.error is NULL AND ba2.canceled = 1), 1, NULL)) AS num_canceled
FROM batch_activities AS ba2
LEFT JOIN batch_activity_host_results AS bahr
ON ba2.execution_id = bahr.batch_execution_id
LEFT JOIN host_script_results AS hsr
ON bahr.host_execution_id = hsr.execution_id
WHERE ba2.status = 'started'
GROUP BY ba2.id
HAVING (num_incompatible + num_ran + num_errored + num_canceled) >= num_targeted
) AS agg
ON agg.batch_id = ba.id
SET
ba.status = 'finished',
ba.finished_at = NOW(),
ba.num_targeted = agg.num_targeted,
ba.num_incompatible = agg.num_incompatible,
ba.num_ran = agg.num_ran,
ba.num_errored = agg.num_errored,
ba.num_canceled = agg.num_canceled,
ba.num_pending = 0
WHERE ba.status = 'started';
`
// TODO -- use `RETURNING` to return the IDs of the updated activities?
_, err := tx.ExecContext(ctx, stmt)
if err != nil {
return ctxerr.Wrap(ctx, err, "marking activities as completed")
}
return nil
}
func (ds *Datastore) MarkActivitiesAsCompleted(ctx context.Context) error {
return ds.markActivitiesAsCompleted(ctx, ds.writer(ctx))
}
Reference in a new issue View git blame Copy permalink