rmmagent/agent/tasks_windows.go
2022-03-19 11:55:43 -07:00

355 lines
9.5 KiB
Go

/*
Copyright 2022 AmidaWare LLC.
Licensed under the Tactical RMM License Version 1.0 (the “License”).
You may only use the Licensed Software in accordance with the License.
A copy of the License is available at:
https://license.tacticalrmm.com
*/
package agent
import (
"encoding/json"
"fmt"
"os"
"path/filepath"
"strings"
"time"
rmm "github.com/amidaware/rmmagent/shared"
"github.com/amidaware/taskmaster"
"github.com/rickb777/date/period"
)
func (a *Agent) RunTask(id int) error {
data := rmm.AutomatedTask{}
url := fmt.Sprintf("/api/v3/%d/%s/taskrunner/", id, a.AgentID)
r1, gerr := a.rClient.R().Get(url)
if gerr != nil {
a.Logger.Debugln(gerr)
return gerr
}
if r1.IsError() {
a.Logger.Debugln("Run Task:", r1.String())
return nil
}
if err := json.Unmarshal(r1.Body(), &data); err != nil {
a.Logger.Debugln(err)
return err
}
start := time.Now()
type TaskResult struct {
Stdout string `json:"stdout"`
Stderr string `json:"stderr"`
RetCode int `json:"retcode"`
ExecTime float64 `json:"execution_time"`
}
payload := TaskResult{}
// loop through all task actions
for _, action := range data.TaskActions {
action_start := time.Now()
if action.ActionType == "script" {
stdout, stderr, retcode, err := a.RunScript(action.Code, action.Shell, action.Args, action.Timeout)
if err != nil {
a.Logger.Debugln(err)
}
// add text to stdout showing which script ran if more than 1 script
action_exec_time := time.Since(action_start).Seconds()
if len(data.TaskActions) > 1 {
payload.Stdout += fmt.Sprintf("\n------------\nRunning Script: %s. Execution Time: %f\n------------\n\n", action.ScriptName, action_exec_time)
}
// save results
payload.Stdout += stdout
payload.Stderr += stderr
payload.RetCode = retcode
if !data.ContinueOnError && stderr != "" {
break
}
} else if action.ActionType == "cmd" {
// out[0] == stdout, out[1] == stderr
out, err := CMDShell(action.Shell, []string{}, action.Command, action.Timeout, false)
if err != nil {
a.Logger.Debugln(err)
}
if len(data.TaskActions) > 1 {
action_exec_time := time.Since(action_start).Seconds()
// add text to stdout showing which script ran
payload.Stdout += fmt.Sprintf("\n------------\nRunning Command: %s. Execution Time: %f\n------------\n\n", action.Command, action_exec_time)
}
// save results
payload.Stdout += out[0]
payload.Stderr += out[1]
// no error
if out[1] == "" {
payload.RetCode = 0
} else {
payload.RetCode = 1
if !data.ContinueOnError {
break
}
}
} else {
a.Logger.Debugln("Invalid Action", action)
}
}
payload.ExecTime = time.Since(start).Seconds()
_, perr := a.rClient.R().SetBody(payload).Patch(url)
if perr != nil {
a.Logger.Debugln(perr)
return perr
}
return nil
}
type SchedTask struct {
PK int `json:"pk"`
Type string `json:"type"`
Name string `json:"name"`
Trigger string `json:"trigger"`
Enabled bool `json:"enabled"`
DayInterval taskmaster.DayInterval `json:"day_interval"`
WeekInterval taskmaster.WeekInterval `json:"week_interval"`
DaysOfWeek taskmaster.DayOfWeek `json:"days_of_week"`
DaysOfMonth taskmaster.DayOfMonth `json:"days_of_month"`
RunOnLastDayOfMonth bool `json:"run_on_last_day_of_month"`
MonthsOfYear taskmaster.Month `json:"months_of_year"`
WeeksOfMonth taskmaster.Week `json:"weeks_of_month"`
StartYear int `json:"start_year"`
StartMonth time.Month `json:"start_month"`
StartDay int `json:"start_day"`
StartHour int `json:"start_hour"`
StartMinute int `json:"start_min"`
ExpireYear int `json:"expire_year"`
ExpireMonth time.Month `json:"expire_month"`
ExpireDay int `json:"expire_day"`
ExpireHour int `json:"expire_hour"`
ExpireMinute int `json:"expire_min"`
RandomDelay period.Period `json:"random_delay"`
RepetitionInterval period.Period `json:"repetition_interval"`
RepetitionDuration period.Period `json:"repetition_duration"`
StopAtDurationEnd bool `json:"stop_at_duration_end"`
Path string `json:"path"`
WorkDir string `json:"workdir"`
Args string `json:"args"`
TaskPolicy taskmaster.TaskInstancesPolicy `json:"multiple_instances"`
RunASAPAfterMissed bool `json:"start_when_available"`
DeleteAfter bool `json:"delete_expired_task_after"`
Overwrite bool `json:"overwrite_task"`
}
func (a *Agent) CreateSchedTask(st SchedTask) (bool, error) {
a.Logger.Debugf("%+v\n", st)
conn, err := taskmaster.Connect()
if err != nil {
a.Logger.Errorln(err)
return false, err
}
defer conn.Disconnect()
var trigger taskmaster.Trigger
var action taskmaster.ExecAction
var tasktrigger taskmaster.TaskTrigger
var now = time.Now()
if st.Trigger == "manual" {
tasktrigger = taskmaster.TaskTrigger{
Enabled: st.Enabled,
StartBoundary: now,
}
} else {
tasktrigger = taskmaster.TaskTrigger{
Enabled: st.Enabled,
StartBoundary: time.Date(st.StartYear, st.StartMonth, st.StartDay, st.StartHour, st.StartMinute, 0, 0, now.Location()),
RepetitionPattern: taskmaster.RepetitionPattern{
RepetitionInterval: st.RepetitionInterval,
RepetitionDuration: st.RepetitionDuration,
StopAtDurationEnd: st.StopAtDurationEnd,
},
}
}
if st.ExpireMinute != 0 {
tasktrigger.EndBoundary = time.Date(st.ExpireYear, st.ExpireMonth, st.ExpireDay, st.ExpireHour, st.ExpireMinute, 0, 0, now.Location())
}
var path, workdir, args string
def := conn.NewTaskDefinition()
switch st.Trigger {
case "runonce":
trigger = taskmaster.TimeTrigger{
TaskTrigger: tasktrigger,
RandomDelay: st.RandomDelay,
}
case "daily":
trigger = taskmaster.DailyTrigger{
TaskTrigger: tasktrigger,
DayInterval: st.DayInterval,
RandomDelay: st.RandomDelay,
}
case "weekly":
trigger = taskmaster.WeeklyTrigger{
TaskTrigger: tasktrigger,
DaysOfWeek: st.DaysOfWeek,
WeekInterval: st.WeekInterval,
RandomDelay: st.RandomDelay,
}
case "monthly":
trigger = taskmaster.MonthlyTrigger{
TaskTrigger: tasktrigger,
DaysOfMonth: st.DaysOfMonth,
MonthsOfYear: st.MonthsOfYear,
RandomDelay: st.RandomDelay,
RunOnLastDayOfMonth: st.RunOnLastDayOfMonth,
}
case "monthlydow":
trigger = taskmaster.MonthlyDOWTrigger{
TaskTrigger: tasktrigger,
DaysOfWeek: st.DaysOfWeek,
MonthsOfYear: st.MonthsOfYear,
WeeksOfMonth: st.WeeksOfMonth,
RandomDelay: st.RandomDelay,
}
case "manual":
trigger = taskmaster.TimeTrigger{
TaskTrigger: tasktrigger,
}
}
def.AddTrigger(trigger)
switch st.Type {
case "rmm":
path = winExeName
workdir = a.ProgramDir
args = fmt.Sprintf("-m taskrunner -p %d", st.PK)
case "schedreboot":
path = "shutdown.exe"
workdir = filepath.Join(os.Getenv("SYSTEMROOT"), "System32")
args = "/r /t 5 /f"
case "custom":
path = st.Path
workdir = st.WorkDir
args = st.Args
}
action = taskmaster.ExecAction{
Path: path,
WorkingDir: workdir,
Args: args,
}
def.AddAction(action)
def.Principal.RunLevel = taskmaster.TASK_RUNLEVEL_HIGHEST
def.Principal.LogonType = taskmaster.TASK_LOGON_SERVICE_ACCOUNT
def.Principal.UserID = "SYSTEM"
def.Settings.AllowDemandStart = true
def.Settings.AllowHardTerminate = true
def.Settings.DontStartOnBatteries = false
def.Settings.Enabled = st.Enabled
def.Settings.StopIfGoingOnBatteries = false
def.Settings.WakeToRun = true
def.Settings.MultipleInstances = st.TaskPolicy
if st.DeleteAfter {
def.Settings.DeleteExpiredTaskAfter = "PT15M"
}
if st.RunASAPAfterMissed {
def.Settings.StartWhenAvailable = true
}
_, success, err := conn.CreateTask(fmt.Sprintf("\\%s", st.Name), def, st.Overwrite)
if err != nil {
a.Logger.Errorln(err)
return false, err
}
return success, nil
}
func DeleteSchedTask(name string) error {
conn, err := taskmaster.Connect()
if err != nil {
return err
}
defer conn.Disconnect()
err = conn.DeleteTask(fmt.Sprintf("\\%s", name))
if err != nil {
return err
}
return nil
}
// CleanupSchedTasks removes all tacticalrmm sched tasks during uninstall
func CleanupSchedTasks() {
conn, err := taskmaster.Connect()
if err != nil {
return
}
defer conn.Disconnect()
tasks, err := conn.GetRegisteredTasks()
if err != nil {
return
}
for _, task := range tasks {
if strings.HasPrefix(task.Name, "TacticalRMM_") {
conn.DeleteTask(fmt.Sprintf("\\%s", task.Name))
}
}
tasks.Release()
}
func ListSchedTasks() []string {
ret := make([]string, 0)
conn, err := taskmaster.Connect()
if err != nil {
return ret
}
defer conn.Disconnect()
tasks, err := conn.GetRegisteredTasks()
if err != nil {
return ret
}
for _, task := range tasks {
ret = append(ret, task.Name)
}
tasks.Release()
return ret
}