Major update to the CLI cosmetics

2 files changed, 94 insertions, 15 deletions

diff --git a/pkg/goexec/method.go b/pkg/goexec/method.go
index bc6039e..97f6f3b 100644
--- a/pkg/goexec/method.go
+++ b/pkg/goexec/method.go
@@ -6,36 +6,54 @@ import (
   "github.com/rs/zerolog"
 )
 
-type Method interface{}
-
-type RemoteMethod interface {
+type Method interface {
   Connect(ctx context.Context) error
   Init(ctx context.Context) error
 }
 
-type RemoteExecuteMethod interface {
-  RemoteMethod
+type Clean interface {
+  Clean(ctx context.Context) error
+}
+
+type CleanMethod interface {
+  Method
+  Clean
+}
+
+type ExecutionMethod interface {
+  Method
   Execute(ctx context.Context, io *ExecutionIO) error
 }
 
-type RemoteExecuteCleanMethod interface {
-  RemoteExecuteMethod
-  Clean(ctx context.Context) error
+type CleanExecutionMethod interface {
+  ExecutionMethod
+  Clean
 }
 
-func ExecuteMethod(ctx context.Context, module RemoteExecuteMethod, execIO *ExecutionIO) (err error) {
+type AuxiliaryMethod interface {
+  Method
+  Call(ctx context.Context) error
+}
+
+type CleanAuxiliaryMethod interface {
+  AuxiliaryMethod
+  Clean
+}
 
+func ExecuteMethod(ctx context.Context, module ExecutionMethod, execIO *ExecutionIO) (err error) {
   log := zerolog.Ctx(ctx)
 
   if err = module.Connect(ctx); err != nil {
     log.Error().Err(err).Msg("Connection failed")
     return fmt.Errorf("connect: %w", err)
   }
+  log.Debug().Msg("Module connected")
 
   if err = module.Init(ctx); err != nil {
     log.Error().Err(err).Msg("Module initialization failed")
     return fmt.Errorf("init module: %w", err)
   }
+  log.Debug().Msg("Module initialized")
 
   if err = module.Execute(ctx, execIO); err != nil {
     log.Error().Err(err).Msg("Execution failed")
@@ -45,8 +63,48 @@ func ExecuteMethod(ctx context.Context, module RemoteExecuteMethod, execIO *Exec
   return
 }
 
-func ExecuteCleanMethod(ctx context.Context, module RemoteExecuteCleanMethod, execIO *ExecutionIO) (err error) {
+func ExecuteAuxiliaryMethod(ctx context.Context, module AuxiliaryMethod) (err error) {
+  log := zerolog.Ctx(ctx)
+
+  if err = module.Connect(ctx); err != nil {
+    log.Error().Err(err).Msg("Connection failed")
+    return fmt.Errorf("connect: %w", err)
+  }
+  log.Debug().Msg("Auxiliary module connected")
+
+  if err = module.Init(ctx); err != nil {
+    log.Error().Err(err).Msg("Module initialization failed")
+    return fmt.Errorf("init module: %w", err)
+  }
+  log.Debug().Msg("Auxiliary module initialized")
+
+  if err = module.Call(ctx); err != nil {
+    log.Error().Err(err).Msg("Auxiliary method failed")
+    return fmt.Errorf("call: %w", err)
+  }
+  log.Debug().Msg("Auxiliary method succeeded")
+
+  return nil
+}
+
+func ExecuteCleanAuxiliaryMethod(ctx context.Context, module CleanAuxiliaryMethod) (err error) {
+  log := zerolog.Ctx(ctx)
+
+  defer func() {
+    if err = module.Clean(ctx); err != nil {
+      log.Error().Err(err).Msg("Module cleanup failed")
+      err = nil
+    }
+  }()
+
+  if err = ExecuteAuxiliaryMethod(ctx, module); err != nil {
+    log.Error().Err(err).Msg("Auxiliary method failed")
+    return fmt.Errorf("execute auxiliary method: %w", err)
+  }
+  return
+}
 
+func ExecuteCleanMethod(ctx context.Context, module CleanExecutionMethod, execIO *ExecutionIO) (err error) {
   log := zerolog.Ctx(ctx)
 
   defer func() {
@@ -61,9 +119,19 @@ func ExecuteCleanMethod(ctx context.Context, module RemoteExecuteCleanMethod, ex
   }
 
   if execIO.Output != nil && execIO.Output.Provider != nil {
-    defer execIO.Output.Provider.Clean(ctx)
+    log.Info().Msg("Collecting output")
 
-    execIO.Output.Provider.GetOutput(ctx, execIO.Output.Writer)
+    defer func() {
+      if cleanErr := execIO.Output.Provider.Clean(ctx); cleanErr != nil {
+        log.Debug().Err(cleanErr).Msg("Output provider cleanup failed")
+      }
+    }()
+
+    if err := execIO.Output.Provider.GetOutput(ctx, execIO.Output.Writer); err != nil {
+      log.Error().Err(err).Msg("Output collection failed")
+      return fmt.Errorf("get output: %w", err)
+    }
+    log.Debug().Msg("Output collection succeeded")
   }
   return
 }
diff --git a/pkg/goexec/wmi/call.go b/pkg/goexec/wmi/call.go
index 3ca7596..7bd2cda 100644
--- a/pkg/goexec/wmi/call.go
+++ b/pkg/goexec/wmi/call.go
@@ -3,7 +3,9 @@ package wmiexec
 import (
   "context"
   "encoding/json"
+  "fmt"
   "github.com/rs/zerolog"
+  "io"
 )
 
 type WmiCall struct {
@@ -12,16 +14,25 @@ type WmiCall struct {
   Class  string
   Method string
   Args   map[string]any
+
+  Out io.Writer
 }
 
-func (m *WmiCall) Call(ctx context.Context) (out []byte, err error) {
+func (m *WmiCall) Call(ctx context.Context) (err error) {
   var outMap map[string]any
 
   if outMap, err = m.query(ctx, m.Class, m.Method, m.Args); err != nil {
     return
   }
-  zerolog.Ctx(ctx).Info().Msg("Call succeeded")
+  zerolog.Ctx(ctx).Info().Msg("WMI call successful")
+
+  out, err := json.Marshal(outMap)
 
-  out, err = json.Marshal(outMap)
+  if m.Out != nil {
+    // Write output with a trailing line feed
+    if _, err = m.Out.Write(append(out, 0x0a)); err != nil {
+      return fmt.Errorf("write output: %w", err)
+    }
+  }
   return
 }