Validate output flag before method execution

2 files changed, 20 insertions, 33 deletions

diff --git a/cmd/args.go b/cmd/args.go
index 50e7c74..b2b6fe2 100644
--- a/cmd/args.go
+++ b/cmd/args.go
@@ -6,6 +6,7 @@ import (
   "fmt"
   "github.com/spf13/cobra"
   "github.com/spf13/pflag"
+  "os"
 )
 
 func registerRpcFlags(cmd *cobra.Command) {
@@ -122,5 +123,20 @@ func argsOutput(methods ...string) func(cmd *cobra.Command, args []string) error
       as = append(as, argsSmbClient())
     }
   }
-  return args(as...)
+
+  return args(append(as, func(*cobra.Command, []string) (err error) {
+
+    if outputPath != "" {
+      if outputPath == "-" {
+        exec.Output.Writer = os.Stdout
+
+      } else if outputPath != "" {
+
+        if exec.Output.Writer, err = os.OpenFile(outputPath, os.O_WRONLY|os.O_CREATE, 0644); err != nil {
+          log.Fatal().Err(err).Msg("Failed to open output file")
+        }
+      }
+    }
+    return
+  })...)
 }
diff --git a/cmd/tsch.go b/cmd/tsch.go
index d08126b..9192b0a 100644
--- a/cmd/tsch.go
+++ b/cmd/tsch.go
@@ -7,7 +7,6 @@ import (
   tschexec "github.com/FalconOpsLLC/goexec/pkg/goexec/tsch"
   "github.com/oiweiwei/go-msrpc/ssp/gssapi"
   "github.com/spf13/cobra"
-  "os"
   "time"
 )
 
@@ -64,6 +63,7 @@ func tschCreateCmdInit() {
   tschCreateCmd.Flags().StringVar(&tschCreate.UserSid, "sid", "S-1-5-18", "User SID to impersonate")
 
   registerProcessExecutionArgs(tschCreateCmd)
+  registerExecutionOutputArgs(tschCreateCmd)
 
   tschCreateCmd.MarkFlagsMutuallyExclusive("name", "path")
 }
@@ -92,7 +92,7 @@ References:
 `,
     Args: args(
       argsRpcClient("cifs"),
-      argsSmbClient(),
+      argsOutput("smb"),
       argsTschDemand,
     ),
 
@@ -108,17 +108,6 @@ References:
 
       ctx := log.WithContext(gssapi.NewSecurityContext(context.TODO()))
 
-      if outputPath == "-" {
-        exec.Output.Writer = os.Stdout
-
-      } else if outputPath != "" {
-
-        if exec.Output.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 exec.Output.Writer.Close()
-      }
-
       if err = goexec.ExecuteCleanMethod(ctx, &tschDemand, &exec); err != nil {
         log.Fatal().Err(err).Msg("Operation failed")
       }
@@ -141,7 +130,7 @@ References:
 `,
     Args: args(
       argsRpcClient("cifs"),
-      argsSmbClient(),
+      argsOutput("smb"),
       argsTschCreate,
     ),
 
@@ -157,27 +146,9 @@ References:
 
       ctx := log.WithContext(gssapi.NewSecurityContext(context.TODO()))
 
-      if outputPath == "-" {
-        exec.Output.Writer = os.Stdout
-
-      } else if outputPath != "" {
-
-        if exec.Output.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 exec.Output.Writer.Close()
-      }
-
       if err = goexec.ExecuteCleanMethod(ctx, &tschCreate, &exec); err != nil {
         log.Fatal().Err(err).Msg("Operation failed")
       }
-
-      if outputPath != "" {
-        if err = tschCreate.IO.GetOutput(ctx); err != nil {
-          log.Error().Err(err).Msg("Failed to get process execution output")
-          returnCode = 4
-        }
-      }
     },
   }
 )