Remote output file deletion after

3 files changed, 15 insertions, 35 deletions

diff --git a/cmd/root.go b/cmd/root.go
index 733bb75..0f8a17a 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -59,9 +59,10 @@ var (
           exec.Output.RemotePath = util.RandomWindowsTempFile()
         }
         exec.Output.Provider = &smb.OutputFileFetcher{
-          Client: &smbClient,
-          Share:  `C$`,
-          File:   exec.Output.RemotePath,
+          Client:           &smbClient,
+          Share:            `C$`,
+          File:             exec.Output.RemotePath,
+          DeleteOutputFile: true, // TEMP
         }
       }
     },
diff --git a/cmd/tsch.go b/cmd/tsch.go
index 7c0fdde..d08126b 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"
-  "io"
   "os"
   "time"
 )
@@ -109,32 +108,20 @@ References:
 
       ctx := log.WithContext(gssapi.NewSecurityContext(context.TODO()))
 
-      var writer io.WriteCloser
-
       if outputPath == "-" {
-        writer = os.Stdout
+        exec.Output.Writer = os.Stdout
 
       } else if outputPath != "" {
 
-        if writer, err = os.OpenFile(outputPath, os.O_WRONLY|os.O_CREATE, 0644); err != nil {
+        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 writer.Close()
+        defer exec.Output.Writer.Close()
       }
 
       if err = goexec.ExecuteCleanMethod(ctx, &tschDemand, &exec); err != nil {
         log.Fatal().Err(err).Msg("Operation failed")
       }
-
-      if outputPath != "" {
-        if reader, err := tschDemand.GetOutput(ctx); err == nil {
-          _, err = io.Copy(writer, reader)
-
-        } else {
-          log.Error().Err(err).Msg("Failed to get process execution output")
-          returnCode = 2
-        }
-      }
     },
   }
   tschCreateCmd = &cobra.Command{
@@ -170,30 +157,25 @@ References:
 
       ctx := log.WithContext(gssapi.NewSecurityContext(context.TODO()))
 
-      var writer io.WriteCloser
-
       if outputPath == "-" {
-        writer = os.Stdout
+        exec.Output.Writer = os.Stdout
 
       } else if outputPath != "" {
 
-        if writer, err = os.OpenFile(outputPath, os.O_WRONLY|os.O_CREATE, 0644); err != nil {
+        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 writer.Close()
+        defer exec.Output.Writer.Close()
       }
 
-      if err = goexec.ExecuteCleanMethod(ctx, &tschDemand, &exec); err != nil {
+      if err = goexec.ExecuteCleanMethod(ctx, &tschCreate, &exec); err != nil {
         log.Fatal().Err(err).Msg("Operation failed")
       }
 
       if outputPath != "" {
-        if reader, err := tschDemand.GetOutput(ctx); err == nil {
-          _, err = io.Copy(writer, reader)
-
-        } else {
+        if err = tschCreate.IO.GetOutput(ctx); err != nil {
           log.Error().Err(err).Msg("Failed to get process execution output")
-          returnCode = 2
+          returnCode = 4
         }
       }
     },
diff --git a/cmd/wmi.go b/cmd/wmi.go
index 196ff82..a9563a8 100644
--- a/cmd/wmi.go
+++ b/cmd/wmi.go
@@ -149,12 +149,9 @@ References:
       }
 
       if outputPath != "" {
-        if reader, err := wmiProc.GetOutput(ctx); err == nil {
-          _, err = io.Copy(writer, reader)
-
-        } else {
+        if err = wmiProc.IO.GetOutput(ctx); err != nil {
           log.Error().Err(err).Msg("Failed to get process execution output")
-          returnCode = 2
+          returnCode = 4
         }
       }
     },