Implemented SCMR Clean*Method

2 files changed, 53 insertions, 109 deletions

diff --git a/cmd/scmr.go b/cmd/scmr.go
index 5dde36f..c527559 100644
--- a/cmd/scmr.go
+++ b/cmd/scmr.go
@@ -3,6 +3,7 @@ package cmd
 import (
   "context"
   "github.com/FalconOpsLLC/goexec/internal/util"
+  "github.com/FalconOpsLLC/goexec/pkg/goexec"
   "github.com/oiweiwei/go-msrpc/ssp/gssapi"
   "github.com/spf13/cobra"
 
@@ -54,7 +55,7 @@ func scmrCreateCmdInit() {
 
   // Constraints
   {
-    scmrCreateCmd.MarkFlagsMutuallyExclusive("no-delete", "no-start")
+    //scmrCreateCmd.MarkFlagsMutuallyExclusive("no-delete", "no-start")
     if err := scmrCreateCmd.MarkFlagRequired("executable-path"); err != nil {
       panic(err)
     }
@@ -98,7 +99,17 @@ func scmrChangeCmdInit() {
 }
 
 func scmrDeleteCmdInit() {
-  scmrDeleteCmd.Flags().StringVarP(&scmrDelete.ServiceName, "service-name", "s", scmrDelete.ServiceName, "Name of service to delete")
+  scmrDeleteFlags := newFlagSet("Service Control")
+  scmrDeleteFlags.Flags.StringVarP(&scmrDelete.ServiceName, "service-name", "s", scmrDelete.ServiceName, "Name of service to delete")
+
+  cmdFlags[scmrDeleteCmd] = []*flagSet{
+    scmrDeleteFlags,
+    defaultAuthFlags,
+    defaultLogFlags,
+    defaultNetRpcFlags,
+  }
+
+  scmrDeleteCmd.Flags().AddFlagSet(scmrDeleteFlags.Flags)
 
   if err := scmrDeleteCmd.MarkFlagRequired("service-name"); err != nil {
     panic(err)
@@ -106,9 +117,9 @@ func scmrDeleteCmdInit() {
 }
 
 var (
-  scmrCreate scmrexec.ScmrCreate
-  scmrChange scmrexec.ScmrChange
-  scmrDelete scmrexec.ScmrDelete
+  scmrCreate = scmrexec.ScmrCreate{}
+  scmrChange = scmrexec.ScmrChange{}
+  scmrDelete = scmrexec.ScmrDelete{}
 
   scmrCmd = &cobra.Command{
     Use:     "scmr",
@@ -127,60 +138,39 @@ var (
 References:
   - https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-scmr/6a8ca926-9477-4dd4-b766-692fab07227e
 `,
-    Args: argsRpcClient("cifs"),
+    Args: args(
+      argsRpcClient("cifs"),
+      argsSmbClient(),
+    ),
 
     Run: func(cmd *cobra.Command, args []string) {
-      var err error
-
-      ctx := gssapi.NewSecurityContext(context.Background())
+      scmrCreate.Client = &rpcClient
+      scmrCreate.IO = exec
 
-      ctx = log.With().
+      log = log.With().
         Str("module", "scmr").
         Str("method", "create").
-        Logger().
-        WithContext(ctx)
-
-      if scmrCreate.ServiceName == "" {
-        log.Warn().Msg("No service Label was provided. Using a random string")
-        scmrCreate.ServiceName = util.RandomString()
-      }
+        Logger()
 
-      if scmrCreate.NoDelete {
-        log.Warn().Msg("Service will not be deleted after execution")
-      }
-
-      if scmrCreate.DisplayName == "" {
-        log.Debug().Msg("No display Label specified, using service Label as display Label")
-        scmrCreate.DisplayName = scmrCreate.ServiceName
-      }
-
-      if err = rpcClient.Connect(ctx); err != nil {
-        log.Fatal().Err(err).Msg("Connection failed")
-      }
-
-      defer func() {
-        closeErr := rpcClient.Close(ctx)
-        if closeErr != nil {
-          log.Error().Err(closeErr).Msg("Failed to close connection")
+      // Warnings
+      {
+        if scmrCreate.ServiceName == "" {
+          log.Warn().Msg("No service name was provided. Using a random string")
+          scmrCreate.ServiceName = util.RandomString()
         }
-      }()
-
-      defer func() {
-        cleanErr := scmrCreate.Clean(ctx)
-        if cleanErr != nil {
-          log.Warn().Err(cleanErr).Msg("Clean operation failed")
+        if scmrCreate.NoDelete {
+          log.Warn().Msg("Service will not be deleted after execution")
+        }
+        if scmrCreate.DisplayName == "" {
+          log.Debug().Msg("No display name specified, using service name as display name")
+          scmrCreate.DisplayName = scmrCreate.ServiceName
         }
-      }()
-
-      if err = scmrCreate.Init(ctx, &rpcClient); err != nil {
-        log.Error().Err(err).Msg("Module initialization failed")
-        returnCode = 2
-        return
       }
 
-      if err = scmrCreate.Execute(ctx, exec.Input); err != nil {
-        log.Error().Err(err).Msg("Execution failed")
-        returnCode = 4
+      ctx := log.WithContext(gssapi.NewSecurityContext(context.Background()))
+
+      if err := goexec.ExecuteCleanMethod(ctx, &scmrCreate, &exec); err != nil {
+        log.Fatal().Err(err).Msg("Operation failed")
       }
     },
   }
@@ -189,44 +179,18 @@ References:
     Use:   "change [target]",
     Short: "Change an existing Windows service to gain execution",
     Args:  argsRpcClient("cifs"),
-    Run: func(cmd *cobra.Command, args []string) {
-      var err error
 
-      ctx := gssapi.NewSecurityContext(context.Background())
+    Run: func(cmd *cobra.Command, args []string) {
+      scmrChange.Client = &rpcClient
+      scmrChange.IO = exec
 
-      ctx = log.With().
+      ctx := log.With().
         Str("module", "scmr").
-        Str("method", "change").
-        Logger().
-        WithContext(ctx)
-
-      if err = rpcClient.Connect(ctx); err != nil {
-        log.Fatal().Err(err).Msg("Connection failed")
-      }
-
-      defer func() {
-        closeErr := rpcClient.Close(ctx)
-        if closeErr != nil {
-          log.Error().Err(closeErr).Msg("Failed to close connection")
-        }
-      }()
-
-      defer func() {
-        cleanErr := scmrChange.Clean(ctx)
-        if cleanErr != nil {
-          log.Warn().Err(cleanErr).Msg("Clean operation failed")
-        }
-      }()
-
-      if err = scmrChange.Init(ctx, &rpcClient); err != nil {
-        log.Error().Err(err).Msg("Module initialization failed")
-        returnCode = 2
-        return
-      }
+        Str("method", "create").
+        Logger().WithContext(gssapi.NewSecurityContext(context.Background()))
 
-      if err = scmrChange.Execute(ctx, exec.Input); err != nil {
-        log.Error().Err(err).Msg("Execution failed")
-        returnCode = 4
+      if err := goexec.ExecuteCleanMethod(ctx, &scmrChange, &exec); err != nil {
+        log.Fatal().Err(err).Msg("Operation failed")
       }
     },
   }
@@ -237,35 +201,15 @@ References:
 
     Args: argsRpcClient("cifs"),
     Run: func(cmd *cobra.Command, args []string) {
-      var err error
-
-      ctx := gssapi.NewSecurityContext(context.Background())
+      scmrDelete.Client = &rpcClient
 
-      ctx = log.With().
+      ctx := log.With().
         Str("module", "scmr").
         Str("method", "delete").
-        Logger().
-        WithContext(ctx)
-
-      if err = rpcClient.Connect(ctx); err != nil {
-        log.Fatal().Err(err).Msg("Connection failed")
-      }
-
-      defer func() {
-        closeErr := rpcClient.Close(ctx)
-        if closeErr != nil {
-          log.Error().Err(closeErr).Msg("Failed to close connection")
-        }
-      }()
-
-      if err = scmrDelete.Init(ctx, &rpcClient); err != nil {
-        log.Error().Err(err).Msg("Module initialization failed")
-        returnCode = 2
-      }
+        Logger().WithContext(gssapi.NewSecurityContext(context.Background()))
 
-      if err = scmrDelete.Clean(ctx); err != nil {
-        log.Warn().Err(err).Msg("Clean failed")
-        returnCode = 4
+      if err := goexec.ExecuteCleanAuxiliaryMethod(ctx, &scmrDelete); err != nil {
+        log.Fatal().Err(err).Msg("Operation failed")
       }
     },
   }
diff --git a/cmd/wmi.go b/cmd/wmi.go
index 78bcb8a..b2e463b 100644
--- a/cmd/wmi.go
+++ b/cmd/wmi.go
@@ -106,7 +106,7 @@ var (
         Str("method", "call").
         Logger().WithContext(gssapi.NewSecurityContext(context.Background()))
 
-      if err := goexec.ExecuteAuxiliaryMethod(ctx, &wmiCall); err != nil {
+      if err := goexec.ExecuteCleanAuxiliaryMethod(ctx, &wmiCall); err != nil {
         log.Fatal().Err(err).Msg("Operation failed")
       }
     },