package cmd

import (
  "context"
  "encoding/json"
  "github.com/FalconOpsLLC/goexec/pkg/goexec"
  wmiexec "github.com/FalconOpsLLC/goexec/pkg/goexec/wmi"
  "github.com/oiweiwei/go-msrpc/ssp/gssapi"
  "github.com/spf13/cobra"
  "os"
)

func wmiCmdInit() {
  cmdFlags[wmiCmd] = []*flagSet{
    defaultAuthFlags,
    defaultLogFlags,
    defaultNetRpcFlags,
  }
  wmiCallCmdInit()
  wmiProcCmdInit()

  wmiCmd.PersistentFlags().AddFlagSet(defaultAuthFlags.Flags)
  wmiCmd.PersistentFlags().AddFlagSet(defaultLogFlags.Flags)
  wmiCmd.PersistentFlags().AddFlagSet(defaultNetRpcFlags.Flags)
  wmiCmd.AddCommand(wmiProcCmd, wmiCallCmd)
}

func wmiCallCmdInit() {
  wmiCallFlags := newFlagSet("WMI")

  wmiCallFlags.Flags.StringVarP(&wmiCall.Resource, "namespace", "n", "//./root/cimv2", "WMI namespace")
  wmiCallFlags.Flags.StringVarP(&wmiCall.Class, "class", "C", "", `WMI class to instantiate (i.e. "Win32_Process")`)
  wmiCallFlags.Flags.StringVarP(&wmiCall.Method, "method", "m", "", `WMI Method to call (i.e. "Create")`)
  wmiCallFlags.Flags.StringVarP(&wmiArguments, "args", "A", "{}", `WMI Method argument(s) in JSON dictionary format (i.e. {"Command":"calc.exe"})`)

  wmiCallCmd.Flags().AddFlagSet(wmiCallFlags.Flags)

  cmdFlags[wmiCallCmd] = []*flagSet{
    wmiCallFlags,
    defaultAuthFlags,
    defaultLogFlags,
    defaultNetRpcFlags,
  }
  if err := wmiCallCmd.MarkFlagRequired("class"); err != nil {
    panic(err)
  }
  if err := wmiCallCmd.MarkFlagRequired("method"); err != nil {
    panic(err)
  }
}

func wmiProcCmdInit() {
  wmiProcExecFlags := newFlagSet("Execution")

  registerExecutionFlags(wmiProcExecFlags.Flags)
  registerExecutionOutputFlags(wmiProcExecFlags.Flags)

  wmiProcExecFlags.Flags.StringVarP(&wmiProc.WorkingDirectory, "directory", "d", `C:\`, "Working directory")

  cmdFlags[wmiProcCmd] = []*flagSet{
    wmiProcExecFlags,
    defaultAuthFlags,
    defaultLogFlags,
    defaultNetRpcFlags,
  }

  wmiProcCmd.Flags().AddFlagSet(wmiProcExecFlags.Flags)
}

var (
  wmiCall = wmiexec.WmiCall{}
  wmiProc = wmiexec.WmiProc{}

  wmiArguments string

  wmiCmd = &cobra.Command{
    Use:   "wmi",
    Short: "Execute with Windows Management Instrumentation (MS-WMI)",
    Long: `Description:
  The wmi module uses remote Windows Management Instrumentation (WMI) to
  perform various operations including process creation.`,
    GroupID: "module",
    Args:    cobra.NoArgs,
  }

  wmiCallCmd = &cobra.Command{
    Use:   "call [target]",
    Short: "Execute specified WMI method",
    Long: `Description:
  The call method creates an instance of the specified WMI class (-c),
  then calls the provided method (-m) with the provided arguments (-A).`,
    Args: args(
      argsRpcClient("cifs", ""),
      func(cmd *cobra.Command, args []string) error {
        return json.Unmarshal([]byte(wmiArguments), &wmiCall.Args)
      }),

    Run: func(cmd *cobra.Command, args []string) {
      wmiCall.Client = &rpcClient
      wmiCall.Out = os.Stdout

      ctx := log.With().
        Str("module", "wmi").
        Str("method", "call").
        Logger().WithContext(gssapi.NewSecurityContext(context.Background()))

      if err := goexec.ExecuteCleanAuxiliaryMethod(ctx, &wmiCall); err != nil {
        log.Fatal().Err(err).Msg("Operation failed")
      }
    },
  }

  wmiProcCmd = &cobra.Command{
    Use:   "proc [target]",
    Short: "Start a Windows process",
    Long: `Description:
  The proc method creates an instance of the Win32_Process WMI class, then
  calls the Win32_Process.Create method with the provided command (-c),
  and optional working directory (-d).`,
    Args: args(
      argsRpcClient("cifs", ""),
      argsOutput("smb"),
    ),

    Run: func(cmd *cobra.Command, args []string) {
      wmiProc.Client = &rpcClient
      wmiProc.IO = exec
      wmiProc.Resource = "//./root/cimv2"

      ctx := log.With().
        Str("module", "wmi").
        Str("method", "proc").
        Logger().WithContext(gssapi.NewSecurityContext(context.Background()))

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