package cmd

import (
  "context"
  "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"
  "io"
  "os"
  "time"
)

func tschCmdInit() {
  registerRpcFlags(tschCmd)

  tschDemandCmdInit()
  tschCmd.AddCommand(tschDemandCmd)

  tschCreateCmdInit()
  tschCmd.AddCommand(tschCreateCmd)
}

func argsTschTaskIdentifiers(name, path string) error {
  switch {
  case path != "":
    return tschexec.ValidateTaskPath(path)
  case name != "":
    return tschexec.ValidateTaskName(name)
  default:
  }
  return nil
}

func argsTschDemand(_ *cobra.Command, _ []string) error {
  return argsTschTaskIdentifiers(tschDemand.TaskName, tschDemand.TaskPath)
}

func argsTschCreate(_ *cobra.Command, _ []string) error {
  return argsTschTaskIdentifiers(tschCreate.TaskName, tschCreate.TaskPath)
}

func tschDemandCmdInit() {
  tschDemandCmd.Flags().StringVarP(&tschDemand.TaskName, "name", "t", "", "Name of task to register")
  tschDemandCmd.Flags().StringVarP(&tschDemand.TaskPath, "path", "P", "", "Path of task to register")
  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)

  tschDemandCmd.MarkFlagsMutuallyExclusive("name", "path")
}

func tschCreateCmdInit() {
  tschCreateCmd.Flags().StringVarP(&tschCreate.TaskName, "name", "t", "", "Name of task to register")
  tschCreateCmd.Flags().StringVarP(&tschCreate.TaskPath, "path", "P", "", "Path of task to register")
  tschCreateCmd.Flags().DurationVar(&tschCreate.StopDelay, "delay-stop", 5*time.Second, "Delay between task execution and termination. This will not stop the process spawned by the task")
  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)

  tschCreateCmd.MarkFlagsMutuallyExclusive("name", "path")
}

var (
  tschDemand tschexec.TschDemand
  tschCreate tschexec.TschCreate

  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:
  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.

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"),
      argsSmbClient(),
      argsTschDemand,
    ),

    Run: func(cmd *cobra.Command, args []string) {
      var err error

      tschDemand.Client = &rpcClient
      tschDemand.IO = exec

      if tschDemand.TaskName == "" && tschDemand.TaskPath == "" {
        tschDemand.TaskPath = `\` + util.RandomString()
      }

      ctx := log.WithContext(gssapi.NewSecurityContext(context.TODO()))

      var writer io.WriteCloser

      if outputPath == "-" {
        writer = os.Stdout

      } else if outputPath != "" {

        if writer, err = os.OpenFile(outputPath, os.O_WRONLY|os.O_CREATE, 0644); err != nil {
          log.Fatal().Err(err).Msg("Failed to open output file")
        }
        defer writer.Close()
      }

      if err = goexec.ExecuteCleanMethod(ctx, &tschDemand, &exec); err != nil {
        log.Fatal().Err(err).Msg("Operation failed")
      }

      if outputPath != "" {
        if reader, err := tschDemand.GetOutput(ctx); err == nil {
          _, err = io.Copy(writer, reader)

        } else {
          log.Error().Err(err).Msg("Failed to get process execution output")
          returnCode = 2
        }
      }
    },
  }
  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
  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: args(
      argsRpcClient("cifs"),
      argsSmbClient(),
      argsTschCreate,
    ),

    Run: func(cmd *cobra.Command, args []string) {
      var err error

      tschCreate.Tsch.Client = &rpcClient
      tschCreate.IO = exec

      if tschCreate.TaskName == "" && tschDemand.TaskPath == "" {
        tschCreate.TaskPath = `\` + util.RandomString()
      }

      ctx := log.WithContext(gssapi.NewSecurityContext(context.TODO()))

      var writer io.WriteCloser

      if outputPath == "-" {
        writer = os.Stdout

      } else if outputPath != "" {

        if writer, err = os.OpenFile(outputPath, os.O_WRONLY|os.O_CREATE, 0644); err != nil {
          log.Fatal().Err(err).Msg("Failed to open output file")
        }
        defer writer.Close()
      }

      if err = goexec.ExecuteCleanMethod(ctx, &tschDemand, &exec); err != nil {
        log.Fatal().Err(err).Msg("Operation failed")
      }

      if outputPath != "" {
        if reader, err := tschDemand.GetOutput(ctx); err == nil {
          _, err = io.Copy(writer, reader)

        } else {
          log.Error().Err(err).Msg("Failed to get process execution output")
          returnCode = 2
        }
      }
    },
  }
)