Fix `tsch create`

2 files changed, 222 insertions, 221 deletions

diff --git a/cmd/tsch.go b/cmd/tsch.go
index 2e8370e..a647f06 100644
--- a/cmd/tsch.go
+++ b/cmd/tsch.go
@@ -1,78 +1,78 @@
 package cmd
 
 import (
-	"context"
-	"fmt"
-	"github.com/FalconOpsLLC/goexec/internal/util"
-	"github.com/FalconOpsLLC/goexec/pkg/goexec"
-	tschexec "github.com/FalconOpsLLC/goexec/pkg/goexec/tsch"
-	"github.com/oiweiwei/go-msrpc/ssp/gssapi"
-	"github.com/spf13/cobra"
-	"time"
+  "context"
+  "fmt"
+  "github.com/FalconOpsLLC/goexec/internal/util"
+  "github.com/FalconOpsLLC/goexec/pkg/goexec"
+  tschexec "github.com/FalconOpsLLC/goexec/pkg/goexec/tsch"
+  "github.com/oiweiwei/go-msrpc/ssp/gssapi"
+  "github.com/spf13/cobra"
+  "time"
 )
 
 func tschCmdInit() {
-	registerRpcFlags(tschCmd)
+  registerRpcFlags(tschCmd)
 
-	tschDemandCmdInit()
-	tschCmd.AddCommand(tschDemandCmd)
+  tschDemandCmdInit()
+  tschCmd.AddCommand(tschDemandCmd)
 
-	tschCreateCmdInit()
-	tschCmd.AddCommand(tschCreateCmd)
+  tschCreateCmdInit()
+  tschCmd.AddCommand(tschCreateCmd)
 }
 
 func tschDemandCmdInit() {
-	tschDemandCmd.Flags().StringVarP(&tschTask, "task", "t", "", "Name or path of the new task")
-	tschDemandCmd.Flags().Uint32Var(&tschDemand.SessionId, "session", 0, "Hijack existing session given the session ID")
-	tschDemandCmd.Flags().BoolVar(&tschDemand.NoDelete, "no-delete", false, "Don't delete task after execution")
-	tschDemandCmd.Flags().StringVar(&tschDemand.UserSid, "sid", "S-1-5-18", "User SID to impersonate")
+  tschDemandCmd.Flags().StringVarP(&tschTask, "task", "t", "", "Name or path of the new task")
+  tschDemandCmd.Flags().Uint32Var(&tschDemand.SessionId, "session", 0, "Hijack existing session given the session ID")
+  tschDemandCmd.Flags().BoolVar(&tschDemand.NoDelete, "no-delete", false, "Don't delete task after execution")
+  tschDemandCmd.Flags().StringVar(&tschDemand.UserSid, "sid", "S-1-5-18", "User SID to impersonate")
 
-	registerProcessExecutionArgs(tschDemandCmd)
-	registerExecutionOutputArgs(tschDemandCmd)
+  registerProcessExecutionArgs(tschDemandCmd)
+  registerExecutionOutputArgs(tschDemandCmd)
 }
 
 func tschCreateCmdInit() {
-	tschCreateCmd.Flags().StringVarP(&tschTask, "task", "t", "", "Name or path of the new task")
-	tschCreateCmd.Flags().DurationVar(&tschCreate.StopDelay, "delay-stop", 5*time.Second, "Delay between task execution and termination. This won't stop the spawned process")
-	tschCreateCmd.Flags().DurationVar(&tschCreate.StartDelay, "start-delay", 5*time.Second, "Delay between task registration and execution")
-	tschCreateCmd.Flags().DurationVar(&tschCreate.DeleteDelay, "delete-delay", 0*time.Second, "Delay between task termination and deletion")
-	tschCreateCmd.Flags().BoolVar(&tschCreate.NoDelete, "no-delete", false, "Don't delete task after execution")
-	tschCreateCmd.Flags().BoolVar(&tschCreate.CallDelete, "call-delete", false, "Directly call SchRpcDelete to delete task")
-	tschCreateCmd.Flags().StringVar(&tschCreate.UserSid, "sid", "S-1-5-18", "User SID to impersonate")
-
-	registerProcessExecutionArgs(tschCreateCmd)
-	registerExecutionOutputArgs(tschCreateCmd)
+  tschCreateCmd.Flags().StringVarP(&tschTask, "task", "t", "", "Name or path of the new task")
+  tschCreateCmd.Flags().DurationVar(&tschCreate.StopDelay, "delay-stop", 5*time.Second, "Delay between task execution and termination. This won't stop the spawned process")
+  tschCreateCmd.Flags().DurationVar(&tschCreate.StartDelay, "start-delay", 5*time.Second, "Delay between task registration and execution")
+  tschCreateCmd.Flags().DurationVar(&tschCreate.DeleteDelay, "delete-delay", 0*time.Second, "Delay between task termination and deletion")
+  tschCreateCmd.Flags().BoolVar(&tschCreate.NoDelete, "no-delete", false, "Don't delete task after execution")
+  tschCreateCmd.Flags().BoolVar(&tschCreate.CallDelete, "call-delete", false, "Directly call SchRpcDelete to delete task")
+  tschCreateCmd.Flags().StringVar(&tschCreate.UserSid, "sid", "S-1-5-18", "User SID to impersonate")
+
+  registerProcessExecutionArgs(tschCreateCmd)
+  registerExecutionOutputArgs(tschCreateCmd)
 }
 
 func argsTask(*cobra.Command, []string) error {
-	switch {
-	case tschTask == "":
-		tschTask = `\` + util.RandomString()
-	case tschexec.ValidateTaskPath(tschTask) == nil:
-	case tschexec.ValidateTaskName(tschTask) == nil:
-		tschTask = `\` + tschTask
-	default:
-		return fmt.Errorf("invalid task name or path: %q", tschTask)
-	}
-	return nil
+  switch {
+  case tschTask == "":
+    tschTask = `\` + util.RandomString()
+  case tschexec.ValidateTaskPath(tschTask) == nil:
+  case tschexec.ValidateTaskName(tschTask) == nil:
+    tschTask = `\` + tschTask
+  default:
+    return fmt.Errorf("invalid task name or path: %q", tschTask)
+  }
+  return nil
 }
 
 var (
-	tschDemand tschexec.TschDemand
-	tschCreate tschexec.TschCreate
+  tschDemand tschexec.TschDemand
+  tschCreate tschexec.TschCreate
 
-	tschTask string
+  tschTask string
 
-	tschCmd = &cobra.Command{
-		Use:   "tsch",
-		Short: "Establish execution via Windows Task Scheduler (MS-TSCH)",
-		Args:  cobra.NoArgs,
-	}
+  tschCmd = &cobra.Command{
+    Use:   "tsch",
+    Short: "Establish execution via Windows Task Scheduler (MS-TSCH)",
+    Args:  cobra.NoArgs,
+  }
 
-	tschDemandCmd = &cobra.Command{
-		Use:   "demand [target]",
-		Short: "Register a remote scheduled task and demand immediate start",
-		Long: `Description:
+  tschDemandCmd = &cobra.Command{
+    Use:   "demand [target]",
+    Short: "Register a remote scheduled task and demand immediate start",
+    Long: `Description:
   Similar to the create method, the demand method will call SchRpcRegisterTask,
   But rather than setting a defined time when the task will start, it will
   additionally call SchRpcRun to forcefully start the task.
@@ -81,31 +81,31 @@ References:
   SchRpcRegisterTask - https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-tsch/849c131a-64e4-46ef-b015-9d4c599c5167
   SchRpcRun - https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-tsch/77f2250d-500a-40ee-be18-c82f7079c4f0
 `,
-		Args: args(
-			argsRpcClient("cifs"),
-			argsOutput("smb"),
-			argsTask,
-		),
-
-		Run: func(*cobra.Command, []string) {
-			tschDemand.IO = exec
-			tschDemand.Client = &rpcClient
-			tschDemand.TaskPath = tschTask
-
-			ctx := log.With().
-				Str("module", "tsch").
-				Str("method", "demand").
-				Logger().WithContext(gssapi.NewSecurityContext(context.TODO()))
-
-			if err := goexec.ExecuteCleanMethod(ctx, &tschDemand, &exec); err != nil {
-				log.Fatal().Err(err).Msg("Operation failed")
-			}
-		},
-	}
-	tschCreateCmd = &cobra.Command{
-		Use:   "create [target]",
-		Short: "Create a remote scheduled task with an automatic start time",
-		Long: `Description:
+    Args: args(
+      argsRpcClient("cifs"),
+      argsOutput("smb"),
+      argsTask,
+    ),
+
+    Run: func(*cobra.Command, []string) {
+      tschDemand.IO = exec
+      tschDemand.Client = &rpcClient
+      tschDemand.TaskPath = tschTask
+
+      ctx := log.With().
+        Str("module", "tsch").
+        Str("method", "demand").
+        Logger().WithContext(gssapi.NewSecurityContext(context.TODO()))
+
+      if err := goexec.ExecuteCleanMethod(ctx, &tschDemand, &exec); err != nil {
+        log.Fatal().Err(err).Msg("Operation failed")
+      }
+    },
+  }
+  tschCreateCmd = &cobra.Command{
+    Use:   "create [target]",
+    Short: "Create a remote scheduled task with an automatic start time",
+    Long: `Description:
   The create method calls SchRpcRegisterTask to register a scheduled task
   with an automatic start time.This method avoids directly calling SchRpcRun,
   and can even avoid calling SchRpcDelete by populating the DeleteExpiredTaskAfter
@@ -117,24 +117,25 @@ References:
   SchRpcDelete - https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-tsch/360bb9b1-dd2a-4b36-83ee-21f12cb97cff
   DeleteExpiredTaskAfter - https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-tsch/6bfde6fe-440e-4ddd-b4d6-c8fc0bc06fae
 `,
-		Args: args(
-			argsRpcClient("cifs"),
-			argsOutput("smb"),
-			argsTask,
-		),
-
-		Run: func(*cobra.Command, []string) {
-			tschCreate.Tsch.Client = &rpcClient
-			tschCreate.IO = exec
-
-			ctx := log.With().
-				Str("module", "tsch").
-				Str("method", "create").
-				Logger().WithContext(gssapi.NewSecurityContext(context.TODO()))
-
-			if err := goexec.ExecuteCleanMethod(ctx, &tschCreate, &exec); err != nil {
-				log.Fatal().Err(err).Msg("Operation failed")
-			}
-		},
-	}
+    Args: args(
+      argsRpcClient("cifs"),
+      argsOutput("smb"),
+      argsTask,
+    ),
+
+    Run: func(*cobra.Command, []string) {
+      tschCreate.Tsch.Client = &rpcClient
+      tschCreate.IO = exec
+      tschCreate.TaskPath = tschTask
+
+      ctx := log.With().
+        Str("module", "tsch").
+        Str("method", "create").
+        Logger().WithContext(gssapi.NewSecurityContext(context.TODO()))
+
+      if err := goexec.ExecuteCleanMethod(ctx, &tschCreate, &exec); err != nil {
+        log.Fatal().Err(err).Msg("Operation failed")
+      }
+    },
+  }
 )
diff --git a/pkg/goexec/tsch/module.go b/pkg/goexec/tsch/module.go
index 74ded2f..dd569cf 100644
--- a/pkg/goexec/tsch/module.go
+++ b/pkg/goexec/tsch/module.go
@@ -1,162 +1,162 @@
 package tschexec
 
 import (
-	"context"
-	"encoding/xml"
-	"errors"
-	"fmt"
-	"github.com/FalconOpsLLC/goexec/pkg/goexec"
-	"github.com/FalconOpsLLC/goexec/pkg/goexec/dce"
-	"github.com/oiweiwei/go-msrpc/msrpc/tsch/itaskschedulerservice/v1"
-	"github.com/rs/zerolog"
+  "context"
+  "encoding/xml"
+  "errors"
+  "fmt"
+  "github.com/FalconOpsLLC/goexec/pkg/goexec"
+  "github.com/FalconOpsLLC/goexec/pkg/goexec/dce"
+  "github.com/oiweiwei/go-msrpc/msrpc/tsch/itaskschedulerservice/v1"
+  "github.com/rs/zerolog"
 )
 
 const (
-	ModuleName = "TSCH"
+  ModuleName = "TSCH"
 )
 
 type Tsch struct {
-	goexec.Cleaner
+  goexec.Cleaner
 
-	Client *dce.Client
-	tsch   itaskschedulerservice.TaskSchedulerServiceClient
+  Client *dce.Client
+  tsch   itaskschedulerservice.TaskSchedulerServiceClient
 
-	TaskPath  string
-	UserSid   string
-	NotHidden bool
+  TaskPath  string
+  UserSid   string
+  NotHidden bool
 }
 
 type registerOptions struct {
-	AllowStartOnDemand bool
-	AllowHardTerminate bool
-	StartWhenAvailable bool
-	Hidden             bool
-	DeleteAfter        string
+  AllowStartOnDemand bool
+  AllowHardTerminate bool
+  StartWhenAvailable bool
+  Hidden             bool
+  DeleteAfter        string
 
-	triggers taskTriggers
+  triggers taskTriggers
 }
 
 func (m *Tsch) Connect(ctx context.Context) (err error) {
 
-	if err = m.Client.Connect(ctx); err == nil {
-		m.AddCleaner(m.Client.Close)
-	}
-	return
+  if err = m.Client.Connect(ctx); err == nil {
+    m.AddCleaner(m.Client.Close)
+  }
+  return
 }
 
 func (m *Tsch) Init(ctx context.Context) (err error) {
 
-	if m.Client.Dce() == nil {
-		return errors.New("DCE connection not initialized")
-	}
+  if m.Client.Dce() == nil {
+    return errors.New("DCE connection not initialized")
+  }
 
-	// Create ITaskSchedulerService Client
-	m.tsch, err = itaskschedulerservice.NewTaskSchedulerServiceClient(ctx, m.Client.Dce())
-	return
+  // Create ITaskSchedulerService Client
+  m.tsch, err = itaskschedulerservice.NewTaskSchedulerServiceClient(ctx, m.Client.Dce())
+  return
 }
 
 func (m *Tsch) registerTask(ctx context.Context, opts *registerOptions, in *goexec.ExecutionIO) (path string, err error) {
 
-	log := zerolog.Ctx(ctx).With().
-		Str("task", m.TaskPath).
-		Logger()
-
-	ctx = log.WithContext(ctx)
-
-	principalId := "1" // This value can be anything
-
-	settings := taskSettings{
-		MultipleInstancesPolicy: "IgnoreNew",
-		IdleSettings: taskIdleSettings{
-			StopOnIdleEnd: true,
-			RestartOnIdle: false,
-		},
-		Enabled:                true,
-		Priority:               7, // a pretty standard value for scheduled tasks
-		AllowHardTerminate:     opts.AllowHardTerminate,
-		AllowStartOnDemand:     opts.AllowStartOnDemand,
-		Hidden:                 opts.Hidden,
-		StartWhenAvailable:     opts.StartWhenAvailable,
-		DeleteExpiredTaskAfter: opts.DeleteAfter,
-	}
-
-	principals := taskPrincipals{
-		Principals: []taskPrincipal{
-			{
-				ID:       principalId,
-				UserID:   m.UserSid,
-				RunLevel: "HighestAvailable",
-			},
-		}}
-
-	cmdline := in.CommandLine()
-
-	actions := taskActions{
-		Context: principalId,
-		Exec: []taskActionExec{
-			{
-				Command:   cmdline[0],
-				Arguments: cmdline[1],
-			},
-		},
-	}
-
-	def := task{
-		TaskVersion:   TaskXmlVersion,
-		TaskNamespace: TaskXmlNamespace,
-		Triggers:      opts.triggers,
-		Actions:       actions,
-		Principals:    principals,
-		Settings:      settings,
-	}
-
-	// Generate task XML content. See https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-tsch/0d6383e4-de92-43e7-b0bb-a60cfa36379f
-
-	doc, err := xml.Marshal(def)
-
-	if err != nil {
-		log.Error().Err(err).Msg("failed to marshal task XML")
-		return "", fmt.Errorf("marshal task: %w", err)
-	}
-
-	taskXml := TaskXmlHeader + string(doc)
-
-	log.Debug().Str("content", taskXml).Msg("Generated task XML")
-
-	registerResponse, err := m.tsch.RegisterTask(ctx, &itaskschedulerservice.RegisterTaskRequest{
-		Path:       m.TaskPath,
-		XML:        taskXml,
-		Flags:      0, // FEATURE: dynamic
-		SDDL:       "",
-		LogonType:  0, // FEATURE: dynamic
-		CredsCount: 0,
-		Creds:      nil,
-	})
-
-	if err != nil {
-		log.Error().Err(err).Msg("Failed to register task")
-		return "", fmt.Errorf("register task: %w", err)
-	}
-	log.Info().Msg("Scheduled task registered")
-
-	return registerResponse.ActualPath, nil
+  log := zerolog.Ctx(ctx).With().
+    Str("task", m.TaskPath).
+    Logger()
+
+  ctx = log.WithContext(ctx)
+
+  principalId := "1" // This value can be anything
+
+  settings := taskSettings{
+    MultipleInstancesPolicy: "IgnoreNew",
+    IdleSettings: taskIdleSettings{
+      StopOnIdleEnd: true,
+      RestartOnIdle: false,
+    },
+    Enabled:                true,
+    Priority:               7, // a pretty standard value for scheduled tasks
+    AllowHardTerminate:     opts.AllowHardTerminate,
+    AllowStartOnDemand:     opts.AllowStartOnDemand,
+    Hidden:                 opts.Hidden,
+    StartWhenAvailable:     opts.StartWhenAvailable,
+    DeleteExpiredTaskAfter: opts.DeleteAfter,
+  }
+
+  principals := taskPrincipals{
+    Principals: []taskPrincipal{
+      {
+        ID:       principalId,
+        UserID:   m.UserSid,
+        RunLevel: "HighestAvailable",
+      },
+    }}
+
+  cmdline := in.CommandLine()
+
+  actions := taskActions{
+    Context: principalId,
+    Exec: []taskActionExec{
+      {
+        Command:   cmdline[0],
+        Arguments: cmdline[1],
+      },
+    },
+  }
+
+  def := task{
+    TaskVersion:   TaskXmlVersion,
+    TaskNamespace: TaskXmlNamespace,
+    Triggers:      opts.triggers,
+    Actions:       actions,
+    Principals:    principals,
+    Settings:      settings,
+  }
+
+  // Generate task XML content. See https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-tsch/0d6383e4-de92-43e7-b0bb-a60cfa36379f
+
+  doc, err := xml.Marshal(def)
+
+  if err != nil {
+    log.Error().Err(err).Msg("failed to marshal task XML")
+    return "", fmt.Errorf("marshal task: %w", err)
+  }
+
+  taskXml := TaskXmlHeader + string(doc)
+
+  log.Debug().Str("content", taskXml).Msg("Generated task XML")
+
+  registerResponse, err := m.tsch.RegisterTask(ctx, &itaskschedulerservice.RegisterTaskRequest{
+    Path:       m.TaskPath,
+    XML:        taskXml,
+    Flags:      0, // FEATURE: dynamic
+    SDDL:       "",
+    LogonType:  0, // FEATURE: dynamic
+    CredsCount: 0,
+    Creds:      nil,
+  })
+
+  if err != nil {
+    log.Error().Err(err).Msg("Failed to register task")
+    return "", fmt.Errorf("register task: %w", err)
+  }
+  log.Info().Msg("Scheduled task registered")
+
+  return registerResponse.ActualPath, nil
 }
 
 func (m *Tsch) deleteTask(ctx context.Context, taskPath string) (err error) {
 
-	log := zerolog.Ctx(ctx).With().
-		Str("path", taskPath).Logger()
+  log := zerolog.Ctx(ctx).With().
+    Str("path", taskPath).Logger()
 
-	_, err = m.tsch.Delete(ctx, &itaskschedulerservice.DeleteRequest{
-		Path: taskPath,
-	})
+  _, err = m.tsch.Delete(ctx, &itaskschedulerservice.DeleteRequest{
+    Path: taskPath,
+  })
 
-	if err != nil {
-		log.Error().Err(err).Msg("Failed to delete task")
-		return fmt.Errorf("delete task: %w", err)
-	}
+  if err != nil {
+    log.Error().Err(err).Msg("Failed to delete task")
+    return fmt.Errorf("delete task: %w", err)
+  }
 
-	log.Info().Msg("Task deleted")
+  log.Info().Msg("Task deleted")
 
-	return
+  return
 }