Update output,IO; add output support to WMI

1 files changed, 39 insertions, 56 deletions

diff --git a/cmd/dcom.go b/cmd/dcom.go
index 5bcec78..c671252 100644
--- a/cmd/dcom.go
+++ b/cmd/dcom.go
@@ -1,37 +1,40 @@
 package cmd
 
 import (
-  "context"
-  dcomexec "github.com/FalconOpsLLC/goexec/pkg/goexec/dcom"
-  "github.com/oiweiwei/go-msrpc/ssp/gssapi"
-  "github.com/spf13/cobra"
+	"context"
+	"github.com/FalconOpsLLC/goexec/pkg/goexec"
+	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(&dcomMmc.WorkingDirectory, "directory", "d", `C:\`, "Working directory")
-  dcomMmcCmd.Flags().StringVar(&dcomMmc.WindowState, "window", "Minimized", "Window state")
+	dcomMmcCmd.Flags().StringVarP(&dcomMmc.WorkingDirectory, "directory", "d", `C:\`, "Working directory")
+	dcomMmcCmd.Flags().StringVar(&dcomMmc.WindowState, "window", "Minimized", "Window state")
 
-  registerProcessExecutionArgs(dcomMmcCmd)
+	registerProcessExecutionArgs(dcomMmcCmd)
+	registerExecutionOutputArgs(dcomMmcCmd)
 }
 
 var (
-  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:
+	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.
 
@@ -41,39 +44,19 @@ 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: 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")
-      }
-
-      defer func() {
-        closeErr := rpcClient.Close(ctx)
-        if closeErr != nil {
-          log.Error().Err(closeErr).Msg("Failed to close connection")
-        }
-      }()
-
-      if err = dcomMmc.Init(ctx, &rpcClient); err != nil {
-        log.Error().Err(err).Msg("Module initialization failed")
-        returnCode = 1
-        return
-      }
-
-      if err = dcomMmc.Execute(ctx, exec.Input); err != nil {
-        log.Error().Err(err).Msg("Execution failed")
-        returnCode = 1
-      }
-    },
-  }
+		Args: args(
+			argsRpcClient("host"),
+			argsOutput("smb"),
+		),
+		Run: func(cmd *cobra.Command, args []string) {
+			dcomMmc.Dcom.Client = &rpcClient
+			dcomMmc.IO = exec
+
+			ctx := log.WithContext(gssapi.NewSecurityContext(context.TODO()))
+
+			if err := goexec.ExecuteCleanMethod(ctx, &dcomMmc, &exec); err != nil {
+				log.Fatal().Err(err).Msg("Operation failed")
+			}
+		},
+	}
 )