package cmd

import (
	"fmt"
	"github.com/FalconOpsLLC/goexec/pkg/exec"
	tschexec "github.com/FalconOpsLLC/goexec/pkg/exec/tsch"
	"github.com/spf13/cobra"
	"time"
)

func tschCmdInit() {
	tschDeleteCmdInit()
	tschCmd.AddCommand(tschDeleteCmd)

	tschRegisterCmdInit()
	tschCmd.AddCommand(tschRegisterCmd)

	tschDemandCmdInit()
	tschCmd.AddCommand(tschDemandCmd)
}

func tschDeleteCmdInit() {
	tschDeleteCmd.Flags().StringVarP(&tschTaskPath, "path", "t", "", "Scheduled task path")
	tschDeleteCmd.MarkFlagRequired("path")
}

func tschDemandCmdInit() {
	tschDemandCmd.Flags().StringVarP(&executable, "executable", "e", "", "Remote Windows executable to invoke")
	tschDemandCmd.Flags().StringVarP(&executableArgs, "args", "a", "", "Arguments to pass to executable")
	tschDemandCmd.Flags().StringVarP(&tschName, "name", "n", "", "Target task name")
	tschDemandCmd.Flags().BoolVar(&tschNoDelete, "no-delete", false, "Don't delete task after execution")
	tschDemandCmd.MarkFlagRequired("executable")
}

func tschRegisterCmdInit() {
	tschRegisterCmd.Flags().StringVarP(&executable, "executable", "e", "", "Remote Windows executable to invoke")
	tschRegisterCmd.Flags().StringVarP(&executableArgs, "args", "a", "", "Arguments to pass to executable")
	tschRegisterCmd.Flags().StringVarP(&tschName, "name", "n", "", "Target task name")
	tschRegisterCmd.Flags().DurationVar(&tschStopDelay, "delay-stop", time.Duration(5*time.Second), "Delay between task execution and termination. This will not stop the process spawned by the task")
	tschRegisterCmd.Flags().DurationVarP(&tschDelay, "delay-start", "d", time.Duration(5*time.Second), "Delay between task registration and execution")
	tschRegisterCmd.Flags().DurationVarP(&tschDeleteDelay, "delay-delete", "D", time.Duration(0*time.Second), "Delay between task termination and deletion")
	tschRegisterCmd.Flags().BoolVar(&tschNoDelete, "no-delete", false, "Don't delete task after execution")
	tschRegisterCmd.Flags().BoolVar(&tschCallDelete, "call-delete", false, "Directly call SchRpcDelete to delete task")

	tschRegisterCmd.MarkFlagsMutuallyExclusive("no-delete", "delay-delete")
	tschRegisterCmd.MarkFlagsMutuallyExclusive("no-delete", "call-delete")
	tschRegisterCmd.MarkFlagsMutuallyExclusive("delay-delete", "call-delete")
	tschRegisterCmd.MarkFlagRequired("executable")
}

var (
	tschNoDelete    bool
	tschCallDelete  bool
	tschDeleteDelay time.Duration
	tschStopDelay   time.Duration
	tschDelay       time.Duration
	tschName        string
	tschTaskPath    string

	tschCmd = &cobra.Command{
		Use:   "tsch",
		Short: "Establish execution via TSCH (ITaskSchedulerService)",
		Args: func(cmd *cobra.Command, args []string) error {
			return fmt.Errorf("command not set. Choose from (delete, register, demand)")
		},
	}
	tschRegisterCmd = &cobra.Command{
		Use:   "register [target]",
		Short: "Register a remote scheduled task with an automatic start time",
		Long: `Description:
  The register 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
  Setting.

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
  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: needsTarget,
		Run: func(cmd *cobra.Command, args []string) {
			if tschNoDelete {
				log.Warn().Msg("Task will not be deleted after execution")
			}
			module := tschexec.Module{}
			execCfg := &exec.ExecutionConfig{
				ExecutableName:  executable,
				ExecutableArgs:  executableArgs,
				ExecutionMethod: tschexec.MethodRegister,

				ExecutionMethodConfig: tschexec.MethodRegisterConfig{
					NoDelete:    tschNoDelete,
					CallDelete:  tschCallDelete,
					StartDelay:  tschDelay,
					StopDelay:   tschStopDelay,
					DeleteDelay: tschDeleteDelay,
					TaskName:    tschName,
				},
			}
			if err := module.Exec(log.WithContext(ctx), creds, target, execCfg); err != nil {
				log.Fatal().Err(err).Msg("TSCH execution failed")
			}
		},
	}
	tschDemandCmd = &cobra.Command{
		Use:   "demand [target]",
		Short: "Register a remote scheduled task and demand immediate start",
		Long: `Description:
  Similar to the register 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.

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: needsTarget,
		Run: func(cmd *cobra.Command, args []string) {
			if tschNoDelete {
				log.Warn().Msg("Task will not be deleted after execution")
			}
			module := tschexec.Module{}
			execCfg := &exec.ExecutionConfig{
				ExecutableName:  executable,
				ExecutableArgs:  executableArgs,
				ExecutionMethod: tschexec.MethodDemand,

				ExecutionMethodConfig: tschexec.MethodDemandConfig{
					NoDelete: tschNoDelete,
					TaskName: tschName,
				},
			}
			if err := module.Exec(log.WithContext(ctx), creds, target, execCfg); err != nil {
				log.Fatal().Err(err).Msg("TSCH execution failed")
			}
		},
	}
	tschDeleteCmd = &cobra.Command{
		Use:   "delete [target]",
		Short: "Manually delete a scheduled task",
		Long: `Description:
  The delete method manually deletes a scheduled task by calling SchRpcDelete

References:
  SchRpcDelete - https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-tsch/360bb9b1-dd2a-4b36-83ee-21f12cb97cff
`,
		Args: needsTarget,
		Run: func(cmd *cobra.Command, args []string) {
			module := tschexec.Module{}
			cleanCfg := &exec.CleanupConfig{
				CleanupMethod:       tschexec.MethodDelete,
				CleanupMethodConfig: tschexec.MethodDeleteConfig{TaskPath: tschTaskPath},
			}
			if err := module.Cleanup(log.WithContext(ctx), creds, target, cleanCfg); err != nil {
				log.Fatal().Err(err).Msg("TSCH cleanup failed")
			}
		},
	}
)