Remodel WMI module; Update TODO; Update deps

4 files changed, 238 insertions, 175 deletions

diff --git a/internal/exec/dcom/exec.go b/internal/exec/dcom/exec.go
index b33791d..a34baf4 100644
--- a/internal/exec/dcom/exec.go
+++ b/internal/exec/dcom/exec.go
@@ -1,186 +1,181 @@
 package dcomexec
 
 import (
-	"context"
-	"errors"
-	"fmt"
-	"github.com/FalconOpsLLC/goexec/internal/client/dce"
-	"github.com/FalconOpsLLC/goexec/internal/exec"
-	"github.com/RedTeamPentesting/adauth"
-	guuid "github.com/google/uuid"
-	"github.com/oiweiwei/go-msrpc/dcerpc"
-	"github.com/oiweiwei/go-msrpc/midl/uuid"
-	"github.com/oiweiwei/go-msrpc/msrpc/dcom"
-	"github.com/oiweiwei/go-msrpc/msrpc/dcom/iremotescmactivator/v0"
-	"github.com/oiweiwei/go-msrpc/msrpc/dcom/oaut"
-	"github.com/oiweiwei/go-msrpc/msrpc/dcom/oaut/idispatch/v0"
-	"github.com/oiweiwei/go-msrpc/msrpc/dtyp"
-	"github.com/rs/zerolog"
+  "context"
+  "errors"
+  "fmt"
+  "github.com/FalconOpsLLC/goexec/internal/client/dce"
+  "github.com/FalconOpsLLC/goexec/internal/exec"
+  "github.com/RedTeamPentesting/adauth"
+  guuid "github.com/google/uuid"
+  "github.com/oiweiwei/go-msrpc/dcerpc"
+  "github.com/oiweiwei/go-msrpc/midl/uuid"
+  "github.com/oiweiwei/go-msrpc/msrpc/dcom"
+  "github.com/oiweiwei/go-msrpc/msrpc/dcom/iremotescmactivator/v0"
+  "github.com/oiweiwei/go-msrpc/msrpc/dcom/oaut"
+  "github.com/oiweiwei/go-msrpc/msrpc/dcom/oaut/idispatch/v0"
+  "github.com/oiweiwei/go-msrpc/msrpc/dtyp"
+  "github.com/rs/zerolog"
 )
 
 const (
-	DefaultDcomEndpoint = "ncacn_ip_tcp:[135]"
+  DefaultDcomEndpoint = "ncacn_ip_tcp:[135]"
 )
 
 var (
-	MmcUuid          = uuid.MustParse("49B2791A-B1AE-4C90-9B8E-E860BA07F889")
-	ShellWindowsUuid = uuid.MustParse("9BA05972-F6A8-11CF-A442-00A0C90A8F39")
-	RandCid          = dcom.CID(*dtyp.GUIDFromUUID(uuid.MustParse(guuid.NewString())))
-	IDispatchIID     = &dcom.IID{
-		Data1: 0x20400,
-		Data2: 0x0,
-		Data3: 0x0,
-		Data4: []byte{0xc0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x46},
-	}
-	ComVersion = &dcom.COMVersion{
-		MajorVersion: 5,
-		MinorVersion: 7,
-	}
-	MmcClsid = dcom.ClassID(*dtyp.GUIDFromUUID(MmcUuid))
-	ORPCThis = &dcom.ORPCThis{
-		Version: ComVersion,
-		CID:     &RandCid,
-	}
+  MmcUuid          = uuid.MustParse("49B2791A-B1AE-4C90-9B8E-E860BA07F889")
+  ShellWindowsUuid = uuid.MustParse("9BA05972-F6A8-11CF-A442-00A0C90A8F39")
+  RandCid          = dcom.CID(*dtyp.GUIDFromUUID(uuid.MustParse(guuid.NewString())))
+  IDispatchIID     = &dcom.IID{
+    Data1: 0x20400,
+    Data2: 0x0,
+    Data3: 0x0,
+    Data4: []byte{0xc0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x46},
+  }
+  ComVersion = &dcom.COMVersion{
+    MajorVersion: 5,
+    MinorVersion: 7,
+  }
+  MmcClsid = dcom.ClassID(*dtyp.GUIDFromUUID(MmcUuid))
+  ORPCThis = &dcom.ORPCThis{
+    Version: ComVersion,
+    CID:     &RandCid,
+  }
 )
 
 func (mod *Module) Connect(ctx context.Context, creds *adauth.Credential, target *adauth.Target, ccfg *exec.ConnectionConfig) (err error) {
 
-	log := zerolog.Ctx(ctx).With().
-		Str("method", ccfg.ConnectionMethod).
-		Str("func", "Exec").Logger()
-
-	if ccfg.ConnectionMethod == exec.ConnectionMethodDCE {
-		if cfg, ok := ccfg.ConnectionMethodConfig.(dce.ConnectionMethodDCEConfig); !ok {
-			return fmt.Errorf("invalid configuration for DCE connection method")
-		} else {
-			opts := []dcerpc.Option{dcerpc.WithSign(), dcerpc.WithSecurityLevel(0)}
-
-			// Fetch target hostname
-			if mod.hostname, err = target.Hostname(ctx); err != nil {
-				log.Debug().Err(err).Msg("Failed to get target hostname")
-				opts = append(opts, dcerpc.WithTargetName(mod.hostname))
-			}
-			// Create DCE connection
-			if mod.dce, err = cfg.GetDce(ctx, creds, target, DefaultDcomEndpoint, "", opts...); err != nil {
-				log.Error().Err(err).Msg("Failed to initialize DCE dialer")
-				return fmt.Errorf("create DCE dialer: %w", err)
-			}
-
-			inst := &dcom.InstantiationInfoData{
-				ClassID:          &MmcClsid,
-				IID:              []*dcom.IID{IDispatchIID},
-				ClientCOMVersion: ComVersion,
-			}
-			scm := &dcom.SCMRequestInfoData{
-				RemoteRequest: &dcom.CustomRemoteRequestSCMInfo{
-					RequestedProtocolSequences: []uint16{7},
-				},
-			}
-			loc := &dcom.LocationInfoData{}
-			ac := &dcom.ActivationContextInfoData{}
-			ap := &dcom.ActivationProperties{
-				DestinationContext: 2,
-				Properties:         []dcom.ActivationProperty{inst, ac, loc, scm},
-			}
-			apin, err := ap.ActivationPropertiesIn()
-			if err != nil {
-				return err
-			}
-			act, err := iremotescmactivator.NewRemoteSCMActivatorClient(ctx, mod.dce)
-			if err != nil {
-				return err
-			}
-			cr, err := act.RemoteCreateInstance(ctx, &iremotescmactivator.RemoteCreateInstanceRequest{
-				ORPCThis: &dcom.ORPCThis{
-					Version: ComVersion,
-					Flags:   1,
-					CID:     &RandCid,
-				},
-				ActPropertiesIn: apin,
-			})
-			if err != nil {
-				return err
-			}
-			log.Info().Msg("RemoteCreateInstance succeeded")
-
-			apout := &dcom.ActivationProperties{}
-			if err = apout.Parse(cr.ActPropertiesOut); err != nil {
-				return err
-			}
-			si := apout.SCMReplyInfoData()
-			pi := apout.PropertiesOutInfo()
-
-			if si == nil {
-				return fmt.Errorf("remote create instance response: SCMReplyInfoData is nil")
-			}
-			if pi == nil {
-				return fmt.Errorf("remote create instance response: PropertiesOutInfo is nil")
-			}
-			oIPID := pi.InterfaceData[0].IPID()
-
-			opts = append(opts, si.RemoteReply.OXIDBindings.EndpointsByProtocol("ncacn_ip_tcp")...)
-			mod.dce, err = cfg.GetDce(ctx, creds, target, DefaultDcomEndpoint, "", opts...)
-			if err != nil {
-				return err
-			}
-			log.Info().Msg("created new DCERPC dialer")
-
-			mod.dc, err = idispatch.NewDispatchClient(ctx, mod.dce, dcom.WithIPID(oIPID))
-			if err != nil {
-				return err
-			}
-			log.Info().Msg("created IDispatch client")
-		}
-	}
-	return
+  log := zerolog.Ctx(ctx).With().
+    Str("method", ccfg.ConnectionMethod).
+    Str("func", "Exec").Logger()
+
+  if ccfg.ConnectionMethod == exec.ConnectionMethodDCE {
+    if cfg, ok := ccfg.ConnectionMethodConfig.(dce.ConnectionMethodDCEConfig); !ok {
+      return errors.New("invalid configuration for DCE connection method")
+    } else {
+      opts := []dcerpc.Option{dcerpc.WithSign(), dcerpc.WithSecurityLevel(0)}
+
+      // Create DCE connection
+      if mod.dce, err = cfg.GetDce(ctx, creds, target, DefaultDcomEndpoint, "", opts...); err != nil {
+        log.Error().Err(err).Msg("Failed to initialize DCE dialer")
+        return fmt.Errorf("create DCE dialer: %w", err)
+      }
+
+      inst := &dcom.InstantiationInfoData{
+        ClassID:          &MmcClsid,
+        IID:              []*dcom.IID{IDispatchIID},
+        ClientCOMVersion: ComVersion,
+      }
+      scm := &dcom.SCMRequestInfoData{
+        RemoteRequest: &dcom.CustomRemoteRequestSCMInfo{
+          RequestedProtocolSequences: []uint16{7},
+        },
+      }
+      loc := &dcom.LocationInfoData{}
+      ac := &dcom.ActivationContextInfoData{}
+      ap := &dcom.ActivationProperties{
+        DestinationContext: 2,
+        Properties:         []dcom.ActivationProperty{inst, ac, loc, scm},
+      }
+      apin, err := ap.ActivationPropertiesIn()
+      if err != nil {
+        return err
+      }
+      act, err := iremotescmactivator.NewRemoteSCMActivatorClient(ctx, mod.dce)
+      if err != nil {
+        return err
+      }
+      cr, err := act.RemoteCreateInstance(ctx, &iremotescmactivator.RemoteCreateInstanceRequest{
+        ORPCThis: &dcom.ORPCThis{
+          Version: ComVersion,
+          Flags:   1,
+          CID:     &RandCid,
+        },
+        ActPropertiesIn: apin,
+      })
+      if err != nil {
+        return err
+      }
+      log.Info().Msg("RemoteCreateInstance succeeded")
+
+      apout := &dcom.ActivationProperties{}
+      if err = apout.Parse(cr.ActPropertiesOut); err != nil {
+        return err
+      }
+      si := apout.SCMReplyInfoData()
+      pi := apout.PropertiesOutInfo()
+
+      if si == nil {
+        return fmt.Errorf("remote create instance response: SCMReplyInfoData is nil")
+      }
+      if pi == nil {
+        return fmt.Errorf("remote create instance response: PropertiesOutInfo is nil")
+      }
+      oIPID := pi.InterfaceData[0].IPID()
+
+      opts = append(opts, si.RemoteReply.OXIDBindings.EndpointsByProtocol("ncacn_ip_tcp")...) // TODO
+      mod.dce, err = cfg.GetDce(ctx, creds, target, DefaultDcomEndpoint, "", opts...)
+      if err != nil {
+        return err
+      }
+      log.Info().Msg("created new DCERPC dialer")
+
+      mod.dc, err = idispatch.NewDispatchClient(ctx, mod.dce, dcom.WithIPID(oIPID))
+      if err != nil {
+        return err
+      }
+      log.Info().Msg("created IDispatch client")
+    }
+  }
+  return
 }
 
 func (mod *Module) Exec(ctx context.Context, ecfg *exec.ExecutionConfig) (err error) {
 
-	log := zerolog.Ctx(ctx).With().
-		Str("method", ecfg.ExecutionMethod).
-		Str("func", "Exec").Logger()
-
-	if ecfg.ExecutionMethod == MethodMmc {
-		if cfg, ok := ecfg.ExecutionMethodConfig.(MethodMmcConfig); !ok {
-			return errors.New("invalid configuration")
-
-		} else {
-			// https://learn.microsoft.com/en-us/previous-versions/windows/desktop/mmc/view-executeshellcommand
-			method := "Document.ActiveView.ExecuteShellCommand"
-			log = log.With().Str("classMethod", method).Logger()
-
-			log.Info().
-				Str("executable", ecfg.ExecutableName).
-				Str("arguments", ecfg.ExecutableArgs).Msg("Attempting execution")
-
-			command := &oaut.Variant{
-				Size:     5,
-				VT:       8,
-				VarUnion: &oaut.Variant_VarUnion{Value: &oaut.Variant_VarUnion_BSTR{BSTR: &oaut.String{Data: ecfg.ExecutableName}}},
-			}
-			directory := &oaut.Variant{
-				Size:     5,
-				VT:       8,
-				VarUnion: &oaut.Variant_VarUnion{Value: &oaut.Variant_VarUnion_BSTR{BSTR: &oaut.String{Data: cfg.WorkingDirectory}}},
-			}
-			parameters := &oaut.Variant{
-				Size:     5,
-				VT:       8,
-				VarUnion: &oaut.Variant_VarUnion{Value: &oaut.Variant_VarUnion_BSTR{BSTR: &oaut.String{Data: ecfg.ExecutableArgs}}},
-			}
-			windowState := &oaut.Variant{
-				Size:     5,
-				VT:       8,
-				VarUnion: &oaut.Variant_VarUnion{Value: &oaut.Variant_VarUnion_BSTR{BSTR: &oaut.String{Data: cfg.WindowState}}},
-			}
-			// Arguments must be passed in reverse order
-			if _, err := callMethod(ctx, mod.dc, method, windowState, parameters, directory, command); err != nil {
-				log.Error().Err(err).Msg("Failed to call method")
-				return fmt.Errorf("call %q: %w", method, err)
-			}
-			log.Info().Msg("Method call successful")
-		}
-	}
-	return nil
+  log := zerolog.Ctx(ctx).With().
+    Str("method", ecfg.ExecutionMethod).
+    Str("func", "Exec").Logger()
+
+  if ecfg.ExecutionMethod == MethodMmc {
+    if cfg, ok := ecfg.ExecutionMethodConfig.(MethodMmcConfig); !ok {
+      return errors.New("invalid configuration")
+
+    } else {
+      // https://learn.microsoft.com/en-us/previous-versions/windows/desktop/mmc/view-executeshellcommand
+      method := "Document.ActiveView.ExecuteShellCommand"
+      log = log.With().Str("classMethod", method).Logger()
+
+      log.Info().
+        Str("executable", ecfg.ExecutableName).
+        Str("arguments", ecfg.ExecutableArgs).Msg("Attempting execution")
+
+      command := &oaut.Variant{
+        Size:     5,
+        VT:       8,
+        VarUnion: &oaut.Variant_VarUnion{Value: &oaut.Variant_VarUnion_BSTR{BSTR: &oaut.String{Data: ecfg.ExecutableName}}},
+      }
+      directory := &oaut.Variant{
+        Size:     5,
+        VT:       8,
+        VarUnion: &oaut.Variant_VarUnion{Value: &oaut.Variant_VarUnion_BSTR{BSTR: &oaut.String{Data: cfg.WorkingDirectory}}},
+      }
+      parameters := &oaut.Variant{
+        Size:     5,
+        VT:       8,
+        VarUnion: &oaut.Variant_VarUnion{Value: &oaut.Variant_VarUnion_BSTR{BSTR: &oaut.String{Data: ecfg.ExecutableArgs}}},
+      }
+      windowState := &oaut.Variant{
+        Size:     5,
+        VT:       8,
+        VarUnion: &oaut.Variant_VarUnion{Value: &oaut.Variant_VarUnion_BSTR{BSTR: &oaut.String{Data: cfg.WindowState}}},
+      }
+      // Arguments must be passed in reverse order
+      if _, err := callMethod(ctx, mod.dc, method, windowState, parameters, directory, command); err != nil {
+        log.Error().Err(err).Msg("Failed to call method")
+        return fmt.Errorf("call %q: %w", method, err)
+      }
+      log.Info().Msg("Method call successful")
+    }
+  }
+  return nil
 }
diff --git a/internal/exec/exec.go b/internal/exec/exec.go
index f3fe7ae..db83d91 100644
--- a/internal/exec/exec.go
+++ b/internal/exec/exec.go
@@ -28,8 +28,7 @@ type ExecutionConfig struct {
 
   ExecutionMethod       string // ExecutionMethod represents the specific execution strategy used by the module.
   ExecutionMethodConfig interface{}
-  //ExecutionOutput       string      // not implemented
-  //ExecutionOutputConfig interface{} // not implemented
+  ReturnOutput          bool
 }
 
 type ShellConfig struct {
diff --git a/internal/exec/wmi/exec.go b/internal/exec/wmi/exec.go
index 5f47738..b61bf5c 100644
--- a/internal/exec/wmi/exec.go
+++ b/internal/exec/wmi/exec.go
@@ -5,6 +5,7 @@ import (
   "encoding/json"
   "errors"
   "fmt"
+  "github.com/FalconOpsLLC/goexec/internal/client/dce"
   "github.com/FalconOpsLLC/goexec/internal/exec"
   "github.com/RedTeamPentesting/adauth"
   "github.com/RedTeamPentesting/adauth/dcerpcauth"
@@ -21,6 +22,7 @@ import (
 const (
   ProtocolSequenceRPC uint16 = 7
   ProtocolSequenceNP  uint16 = 15
+  DefaultWmiEndpoint  string = "ncacn_ip_tcp:[135]"
 )
 
 var (
@@ -45,7 +47,74 @@ func (mod *Module) Cleanup(ctx context.Context, _ *exec.CleanupConfig) (err erro
   return
 }
 
-func (mod *Module) Connect(ctx context.Context, creds *adauth.Credential, target *adauth.Target, _ *exec.ConnectionConfig) (err error) {
+func (mod *Module) Connect(ctx context.Context, creds *adauth.Credential, target *adauth.Target, ccfg *exec.ConnectionConfig) (err error) {
+
+  log := zerolog.Ctx(ctx).With().
+    Str("method", ccfg.ConnectionMethod).
+    Str("func", "Connect").Logger()
+
+  if cfg, ok := ccfg.ConnectionMethodConfig.(dce.ConnectionMethodDCEConfig); !ok {
+    return errors.New("invalid configuration for DCE connection method")
+  } else {
+    var dceOpts []dcerpc.Option
+
+    // Create DCE connection
+    if mod.dce, err = cfg.GetDce(ctx, creds, target, DefaultWmiEndpoint, "", dceOpts...); err != nil {
+      log.Error().Err(err).Msg("Failed to initialize DCE dialer")
+      return fmt.Errorf("create DCE dialer: %w", err)
+    }
+    ia, err := iactivation.NewActivationClient(ctx, mod.dce)
+    if err != nil {
+      log.Error().Err(err).Msg("Failed to create activation client")
+      return fmt.Errorf("create activation client: %w", err)
+    }
+    act, err := ia.RemoteActivation(ctx, &iactivation.RemoteActivationRequest{
+      ORPCThis:                   ORPCThis,
+      ClassID:                    wmi.Level1LoginClassID.GUID(),
+      IIDs:                       []*dcom.IID{iwbemlevel1login.Level1LoginIID},
+      RequestedProtocolSequences: []uint16{ProtocolSequenceRPC, ProtocolSequenceNP}, // TODO: dynamic
+    })
+    if err != nil {
+      return fmt.Errorf("request remote activation: %w", err)
+    }
+    if act.HResult != 0 {
+      return fmt.Errorf("remote activation failed with code %d", act.HResult)
+    }
+    retBinds := act.OXIDBindings.GetStringBindings()
+    if len(act.InterfaceData) < 1 || len(retBinds) < 1 {
+      return errors.New("remote activation failed")
+    }
+    ipid := act.InterfaceData[0].GetStandardObjectReference().Std.IPID
+    for _, sb := range retBinds {
+      //sb.NetworkAddr = target.AddressWithoutPort() // TODO: check if sb.NetworkAddr contains the port
+      dceOpts = append(dceOpts, dcerpc.WithEndpoint(sb.String()))
+    }
+
+    if mod.dce, err = cfg.GetDce(ctx, creds, target, DefaultWmiEndpoint, "", dceOpts...); err != nil {
+      log.Error().Err(err).Msg("Failed to initialize secondary DCE dialer")
+    }
+    loginClient, err := iwbemlevel1login.NewLevel1LoginClient(ctx, mod.dce, dcom.WithIPID(ipid))
+    if err != nil {
+      return fmt.Errorf("initialize wbem login client: %w", err)
+    }
+    login, err := loginClient.NTLMLogin(ctx, &iwbemlevel1login.NTLMLoginRequest{
+      This:            ORPCThis,
+      NetworkResource: "//./root/cimv2", // TODO: make this dynamic
+    })
+    if err != nil {
+      return fmt.Errorf("ntlm login: %w", err)
+    }
+
+    mod.sc, err = iwbemservices.NewServicesClient(ctx, mod.dce, dcom.WithIPID(login.Namespace.InterfacePointer().IPID()))
+    if err != nil {
+      return fmt.Errorf("iwbemservices.NewServicesClient: %w", err)
+    }
+  }
+
+  return
+}
+
+func (mod *Module) _Connect(ctx context.Context, creds *adauth.Credential, target *adauth.Target, _ *exec.ConnectionConfig) (err error) {
 
   var baseOpts, authOpts []dcerpc.Option
   var ipid *dcom.IPID // This will store the IPID of the remote instance
@@ -157,8 +226,8 @@ func (mod *Module) Exec(ctx context.Context, ecfg *exec.ExecutionConfig) (err er
     Str("module", "tsch").
     Str("method", ecfg.ExecutionMethod).Logger()
 
-  if ecfg.ExecutionMethod == MethodCustom {
-    if cfg, ok := ecfg.ExecutionMethodConfig.(MethodCustomConfig); !ok {
+  if ecfg.ExecutionMethod == MethodCall {
+    if cfg, ok := ecfg.ExecutionMethodConfig.(MethodCallConfig); !ok {
       return errors.New("invalid execution configuration")
 
     } else {
diff --git a/internal/exec/wmi/module.go b/internal/exec/wmi/module.go
index f90af42..0e83aa8 100644
--- a/internal/exec/wmi/module.go
+++ b/internal/exec/wmi/module.go
@@ -16,7 +16,7 @@ type Module struct {
   sc  iwbemservices.ServicesClient
 }
 
-type MethodCustomConfig struct {
+type MethodCallConfig struct {
   Class     string
   Method    string
   Arguments map[string]any
@@ -28,6 +28,6 @@ type MethodProcessConfig struct {
 }
 
 const (
-  MethodCustom  = "custom"
+  MethodCall    = "call"
   MethodProcess = "process"
 )