1 files changed, 53 insertions, 49 deletions

diff --git a/cmd/dcom.go b/cmd/dcom.go
index d105b0c..5bcec78 100644
--- a/cmd/dcom.go
+++ b/cmd/dcom.go
@@ -1,38 +1,37 @@
 package cmd
 
 import (
-	"github.com/FalconOpsLLC/goexec/internal/exec"
-	dcomexec "github.com/FalconOpsLLC/goexec/internal/exec/dcom"
-	"github.com/spf13/cobra"
+  "context"
+  dcomexec "github.com/FalconOpsLLC/goexec/pkg/goexec/dcom"
+  "github.com/oiweiwei/go-msrpc/ssp/gssapi"
+  "github.com/spf13/cobra"
 )
 
 func dcomCmdInit() {
-	registerRpcFlags(dcomCmd)
-	dcomMmcCmdInit()
-	dcomCmd.AddCommand(dcomMmcCmd)
+  registerRpcFlags(dcomCmd)
+  dcomMmcCmdInit()
+  dcomCmd.AddCommand(dcomMmcCmd)
 }
 
 func dcomMmcCmdInit() {
-	dcomMmcCmd.Flags().StringVarP(&executable, "executable", "e", "", "Remote Windows executable to invoke")
-	dcomMmcCmd.Flags().StringVarP(&workingDirectory, "directory", "d", `C:\`, "Working directory")
-	dcomMmcCmd.Flags().StringVarP(&executableArgs, "args", "a", "", "Process command line")
-	dcomMmcCmd.Flags().StringVar(&windowState, "window", "Minimized", "Window state")
-	dcomMmcCmd.Flags().StringVarP(&command, "command", "c", ``, "Windows executable & arguments to run")
+  dcomMmcCmd.Flags().StringVarP(&dcomMmc.WorkingDirectory, "directory", "d", `C:\`, "Working directory")
+  dcomMmcCmd.Flags().StringVar(&dcomMmc.WindowState, "window", "Minimized", "Window state")
 
-	dcomMmcCmd.MarkFlagsOneRequired("executable", "command")
-	dcomMmcCmd.MarkFlagsMutuallyExclusive("executable", "command")
+  registerProcessExecutionArgs(dcomMmcCmd)
 }
 
 var (
-	dcomCmd = &cobra.Command{
-		Use:   "dcom",
-		Short: "Establish execution via DCOM",
-		Args:  cobra.NoArgs,
-	}
-	dcomMmcCmd = &cobra.Command{
-		Use:   "mmc [target]",
-		Short: "Establish execution via the DCOM MMC20.Application object",
-		Long: `Description:
+  dcomMmc dcomexec.DcomMmc
+
+  dcomCmd = &cobra.Command{
+    Use:   "dcom",
+    Short: "Establish execution via DCOM",
+    Args:  cobra.NoArgs,
+  }
+  dcomMmcCmd = &cobra.Command{
+    Use:   "mmc [target]",
+    Short: "Establish execution via the DCOM MMC20.Application object",
+    Long: `Description:
   The mmc method uses the exposed MMC20.Application object to call Document.ActiveView.ShellExec,
   and ultimately execute system commands.
 
@@ -42,34 +41,39 @@ References:
   https://github.com/fortra/impacket/blob/master/examples/dcomexec.py
   https://learn.microsoft.com/en-us/previous-versions/windows/desktop/mmc/view-executeshellcommand
 `,
-		Args: needsRpcTarget("host"),
-		Run: func(cmd *cobra.Command, args []string) {
+    Args: argsRpcClient("host"),
+    Run: func(cmd *cobra.Command, args []string) {
+      var err error
+
+      ctx := gssapi.NewSecurityContext(context.Background())
+
+      ctx = log.With().
+        Str("module", "dcom").
+        Str("method", "mmc").
+        Logger().
+        WithContext(ctx)
+
+      if err = rpcClient.Connect(ctx); err != nil {
+        log.Fatal().Err(err).Msg("Connection failed")
+      }
 
-			ctx = log.With().
-				Str("module", "dcom").
-				Str("method", "mmc").
-				Logger().WithContext(ctx)
+      defer func() {
+        closeErr := rpcClient.Close(ctx)
+        if closeErr != nil {
+          log.Error().Err(closeErr).Msg("Failed to close connection")
+        }
+      }()
 
-			module := dcomexec.Module{}
-			connCfg := &exec.ConnectionConfig{
-				ConnectionMethod:       exec.ConnectionMethodDCE,
-				ConnectionMethodConfig: dceConfig,
-			}
-			execCfg := &exec.ExecutionConfig{
-				ExecutableName:  executable,
-				ExecutableArgs:  executableArgs,
-				ExecutionMethod: dcomexec.MethodMmc,
+      if err = dcomMmc.Init(ctx, &rpcClient); err != nil {
+        log.Error().Err(err).Msg("Module initialization failed")
+        returnCode = 1
+        return
+      }
 
-				ExecutionMethodConfig: dcomexec.MethodMmcConfig{
-					WorkingDirectory: workingDirectory,
-					WindowState:      windowState,
-				},
-			}
-			if err := module.Connect(ctx, creds, target, connCfg); err != nil {
-				log.Fatal().Err(err).Msg("Connection failed")
-			} else if err = module.Exec(ctx, execCfg); err != nil {
-				log.Fatal().Err(err).Msg("Execution failed")
-			}
-		},
-	}
+      if err = dcomMmc.Execute(ctx, exec.Input); err != nil {
+        log.Error().Err(err).Msg("Execution failed")
+        returnCode = 1
+      }
+    },
+  }
 )