Remote output file deletion after

6 files changed, 57 insertions, 53 deletions

diff --git a/pkg/goexec/io.go b/pkg/goexec/io.go
index 846a1fd..c26fc7f 100644
--- a/pkg/goexec/io.go
+++ b/pkg/goexec/io.go
@@ -3,14 +3,14 @@ package goexec
 import (
   "bytes"
   "context"
-  "errors"
   "fmt"
   "io"
   "os"
 )
 
 type OutputProvider interface {
-  GetOutput(ctx context.Context) (out io.ReadCloser, err error)
+  GetOutput(ctx context.Context, writer io.Writer) (err error)
+  Clean(ctx context.Context) (err error)
 }
 
 type ExecutionIO struct {
@@ -24,6 +24,7 @@ type ExecutionOutput struct {
   NoDelete   bool
   RemotePath string
   Provider   OutputProvider
+  Writer     io.WriteCloser
 }
 
 type ExecutionInput struct {
@@ -34,18 +35,24 @@ type ExecutionInput struct {
   CommandLine    string
 }
 
-func (execIO *ExecutionIO) GetOutput(ctx context.Context) (out io.ReadCloser, err error) {
-
+func (execIO *ExecutionIO) GetOutput(ctx context.Context) (err error) {
   if execIO.Output.Provider != nil {
-    return execIO.Output.Provider.GetOutput(ctx)
+    return execIO.Output.Provider.GetOutput(ctx, execIO.Output.Writer)
   }
-  return nil, errors.New("output provider not set")
+  return nil
 }
 
 func (execIO *ExecutionIO) CommandLine() string {
   return execIO.Input.Command()
 }
 
+func (execIO *ExecutionIO) Clean(ctx context.Context) (err error) {
+  if execIO.Output.Provider != nil {
+    return execIO.Output.Provider.Clean(ctx)
+  }
+  return nil
+}
+
 func (execIO *ExecutionIO) String() (cmd string) {
 
   cmd = execIO.Input.Command()
diff --git a/pkg/goexec/method.go b/pkg/goexec/method.go
index 7ea0ddc..bc6039e 100644
--- a/pkg/goexec/method.go
+++ b/pkg/goexec/method.go
@@ -18,21 +18,11 @@ type RemoteExecuteMethod interface {
   Execute(ctx context.Context, io *ExecutionIO) error
 }
 
-type RemoteCleanMethod interface {
-  RemoteMethod
-  Clean(ctx context.Context) error
-}
-
 type RemoteExecuteCleanMethod interface {
   RemoteExecuteMethod
   Clean(ctx context.Context) error
 }
 
-type RemoteExecuteCleanMethodWithOutput interface {
-  RemoteExecuteCleanMethod
-  OutputProvider
-}
-
 func ExecuteMethod(ctx context.Context, module RemoteExecuteMethod, execIO *ExecutionIO) (err error) {
 
   log := zerolog.Ctx(ctx)
@@ -51,6 +41,7 @@ func ExecuteMethod(ctx context.Context, module RemoteExecuteMethod, execIO *Exec
     log.Error().Err(err).Msg("Execution failed")
     return fmt.Errorf("execute: %w", err)
   }
+
   return
 }
 
@@ -61,12 +52,18 @@ func ExecuteCleanMethod(ctx context.Context, module RemoteExecuteCleanMethod, ex
   defer func() {
     if err = module.Clean(ctx); err != nil {
       log.Error().Err(err).Msg("Module cleanup failed")
+      err = nil
     }
   }()
 
-  return ExecuteMethod(ctx, module, execIO)
-}
+  if err = ExecuteMethod(ctx, module, execIO); err != nil {
+    return
+  }
 
-func ExecuteCleanMethodWithOutput(ctx context.Context, module RemoteExecuteCleanMethodWithOutput, execIO *ExecutionIO) (err error) {
-  return ExecuteCleanMethod(ctx, module, execIO)
+  if execIO.Output != nil && execIO.Output.Provider != nil {
+    defer execIO.Output.Provider.Clean(ctx)
+
+    execIO.Output.Provider.GetOutput(ctx, execIO.Output.Writer)
+  }
+  return
 }
diff --git a/pkg/goexec/smb/client.go b/pkg/goexec/smb/client.go
index d9e8772..384f4bc 100644
--- a/pkg/goexec/smb/client.go
+++ b/pkg/goexec/smb/client.go
@@ -29,13 +29,15 @@ func (c *Client) Logger(ctx context.Context) zerolog.Logger {
   return zerolog.Ctx(ctx).With().Str("client", c.String()).Logger()
 }
 
-func (c *Client) Mount(_ context.Context, share string) (err error) {
+func (c *Client) Mount(ctx context.Context, share string) (err error) {
 
   if c.sess == nil {
     return errors.New("SMB session not initialized")
   }
 
   c.mount, err = c.sess.Mount(share)
+  zerolog.Ctx(ctx).Debug().Str("share", share).Msg("Mounted SMB share")
+
   return
 }
 
@@ -82,7 +84,7 @@ func (c *Client) Close(ctx context.Context) (err error) {
   if c.conn != nil {
     defer func() {
       if err = c.conn.Close(); err != nil {
-        log.Error().Err(err).Msgf("Failed to close %s connection", c.String())
+        log.Debug().Err(err).Msgf("Failed to close %s connection", c.String())
       }
       log.Debug().Msgf("Closed %s connection", c.String())
     }()
@@ -92,7 +94,7 @@ func (c *Client) Close(ctx context.Context) (err error) {
   if c.sess != nil {
     defer func() {
       if err = c.sess.Logoff(); err != nil {
-        log.Error().Err(err).Msgf("Failed to discard %s session", c.String())
+        log.Debug().Err(err).Msgf("Failed to discard %s session", c.String())
       }
       log.Debug().Msgf("Discarded %s session", c.String())
     }()
@@ -102,7 +104,7 @@ func (c *Client) Close(ctx context.Context) (err error) {
   if c.mount != nil {
     defer func() {
       if err = c.mount.Umount(); err != nil {
-        log.Error().Err(err).Msg("Failed to unmount share")
+        log.Debug().Err(err).Msg("Failed to unmount share")
       }
       log.Debug().Msg("Unmounted file share")
     }()
diff --git a/pkg/goexec/smb/output.go b/pkg/goexec/smb/output.go
index 35c40c8..25768d1 100644
--- a/pkg/goexec/smb/output.go
+++ b/pkg/goexec/smb/output.go
@@ -2,6 +2,7 @@ package smb
 
 import (
   "context"
+  "errors"
   "github.com/FalconOpsLLC/goexec/pkg/goexec"
   "io"
   "os"
@@ -12,22 +13,24 @@ import (
 var (
   DefaultOutputPollInterval = 1 * time.Second
   DefaultOutputPollTimeout  = 60 * time.Second
-  pathPrefix                = regexp.MustCompile(`^([a-zA-Z]:)?\\*`)
+  pathPrefix                = regexp.MustCompile(`^([a-zA-Z]:)?[\\/]*`)
 )
 
 type OutputFileFetcher struct {
   goexec.Cleaner
 
-  Client       *Client
-  Share        string
-  File         string
-  PollInterval time.Duration
-  PollTimeout  time.Duration
+  Client *Client
+
+  Share            string
+  File             string
+  DeleteOutputFile bool
+  PollInterval     time.Duration
+  PollTimeout      time.Duration
 
   relativePath string
 }
 
-func (o *OutputFileFetcher) GetOutput(ctx context.Context) (reader io.ReadCloser, err error) {
+func (o *OutputFileFetcher) GetOutput(ctx context.Context, writer io.Writer) (err error) {
 
   if o.PollInterval == 0 {
     o.PollInterval = DefaultOutputPollInterval
@@ -42,6 +45,7 @@ func (o *OutputFileFetcher) GetOutput(ctx context.Context) (reader io.ReadCloser
   if err != nil {
     return
   }
+  defer o.AddCleaner(o.Client.Close)
 
   err = o.Client.Mount(ctx, o.Share)
   if err != nil {
@@ -49,15 +53,29 @@ func (o *OutputFileFetcher) GetOutput(ctx context.Context) (reader io.ReadCloser
   }
 
   stopAt := time.Now().Add(o.PollTimeout)
+  var reader io.ReadCloser
 
   for {
     if time.Now().After(stopAt) {
-      return
+      return errors.New("output timeout")
     }
     if reader, err = o.Client.mount.OpenFile(o.relativePath, os.O_RDONLY, 0); err == nil {
-      return
+      break
     }
     time.Sleep(o.PollInterval)
   }
+
+  if _, err = io.Copy(writer, reader); err != nil {
+    return
+  }
+
+  o.AddCleaner(func(_ context.Context) error { return reader.Close() })
+
+  if o.DeleteOutputFile {
+    o.AddCleaner(func(_ context.Context) error {
+      return o.Client.mount.Remove(o.relativePath)
+    })
+  }
+
   return
 }
diff --git a/pkg/goexec/tsch/demand.go b/pkg/goexec/tsch/demand.go
index cd2ccd5..11522dd 100644
--- a/pkg/goexec/tsch/demand.go
+++ b/pkg/goexec/tsch/demand.go
@@ -2,12 +2,10 @@ package tschexec
 
 import (
   "context"
-  "errors"
   "fmt"
   "github.com/FalconOpsLLC/goexec/pkg/goexec"
   "github.com/oiweiwei/go-msrpc/msrpc/tsch/itaskschedulerservice/v1"
   "github.com/rs/zerolog"
-  "io"
 )
 
 const (
@@ -81,11 +79,3 @@ func (m *TschDemand) Execute(ctx context.Context, in *goexec.ExecutionIO) (err e
   }
   return
 }
-
-func (m *TschDemand) GetOutput(ctx context.Context) (reader io.ReadCloser, err error) {
-
-  if m.IO.Output != nil {
-    return m.IO.GetOutput(ctx)
-  }
-  return nil, errors.New("no available output provider")
-}
diff --git a/pkg/goexec/wmi/proc.go b/pkg/goexec/wmi/proc.go
index 25e515c..444643b 100644
--- a/pkg/goexec/wmi/proc.go
+++ b/pkg/goexec/wmi/proc.go
@@ -5,7 +5,6 @@ import (
   "errors"
   "github.com/FalconOpsLLC/goexec/pkg/goexec"
   "github.com/rs/zerolog"
-  "io"
 )
 
 const (
@@ -53,12 +52,3 @@ func (m *WmiProc) Execute(ctx context.Context, execIO *goexec.ExecutionIO) (err
   }
   return
 }
-
-func (m *WmiProc) GetOutput(ctx context.Context) (reader io.ReadCloser, err error) {
-
-  if m.IO.Output != nil {
-
-    return m.IO.GetOutput(ctx)
-  }
-  return nil, errors.New("no available output provider")
-}