massively improved SCMR module + new DCE client

11 files changed, 606 insertions, 820 deletions

diff --git a/internal/client/dce/dce.go b/internal/client/dce/dce.go
index 2be5a1e..1ddec65 100644
--- a/internal/client/dce/dce.go
+++ b/internal/client/dce/dce.go
@@ -1,21 +1,53 @@
 package dce
 
-import "github.com/oiweiwei/go-msrpc/dcerpc"
+import (
+  "context"
+  "fmt"
+  "github.com/RedTeamPentesting/adauth"
+  "github.com/RedTeamPentesting/adauth/dcerpcauth"
+  "github.com/oiweiwei/go-msrpc/dcerpc"
+  "github.com/oiweiwei/go-msrpc/msrpc/epm/epm/v3"
+  "github.com/oiweiwei/go-msrpc/ssp/gssapi"
+  "github.com/rs/zerolog"
+)
 
 var (
-  NP           = "ncacn_np"
-  TCP          = "ncacn_ip_tcp"
-  HTTP         = "ncacn_http"
-  DefaultPorts = map[string]uint16{
-    NP:   445,
-    TCP:  135,
-    HTTP: 593,
-  }
+  NP   = "ncacn_np"
+  TCP  = "ncacn_ip_tcp"
+  HTTP = "ncacn_http"
 )
 
 type ConnectionMethodDCEConfig struct {
-  NoEpm    bool // NoEpm disables EPM
-  EpmAuto  bool // EpmAuto will find any suitable endpoint, without any filter
-  Endpoint *dcerpc.StringBinding
-  Options  []dcerpc.Option
+  NoEpm      bool                  // NoEpm disables EPM
+  EpmAuto    bool                  // EpmAuto will find any suitable endpoint, without any filter
+  Endpoint   *dcerpc.StringBinding // Endpoint is the endpoint passed to dcerpc.WithEndpoint. ignored if EpmAuto is false
+  DceOptions []dcerpc.Option       // DceOptions are the options passed to dcerpc.Dial
+  EpmOptions []dcerpc.Option       // EpmOptions are the options passed to epm.EndpointMapper
+}
+
+func (cfg *ConnectionMethodDCEConfig) GetDce(ctx context.Context, creds *adauth.Credential, target *adauth.Target, opts ...dcerpc.Option) (cc dcerpc.Conn, err error) {
+  dceOpts := append(opts, cfg.DceOptions...)
+  epmOpts := append(opts, cfg.EpmOptions...)
+
+  log := zerolog.Ctx(ctx).With().
+    Str("client", "DCERPC").Logger()
+
+  // Mandatory logging
+  dceOpts = append(dceOpts, dcerpc.WithLogger(log))
+  epmOpts = append(epmOpts, dcerpc.WithLogger(log))
+
+  ctx = gssapi.NewSecurityContext(ctx)
+  ao, err := dcerpcauth.AuthenticationOptions(ctx, creds, target, &dcerpcauth.Options{})
+  if err != nil {
+    log.Error().Err(err).Msg("Failed to parse authentication options")
+    return nil, fmt.Errorf("parse auth options: %w", err)
+  }
+  if cfg.Endpoint != nil && !cfg.EpmAuto {
+    dceOpts = append(dceOpts, dcerpc.WithEndpoint(cfg.Endpoint.String()))
+  }
+  if !cfg.NoEpm {
+    dceOpts = append(dceOpts,
+      epm.EndpointMapper(ctx, target.AddressWithoutPort(), append(epmOpts, ao...)...))
+  }
+  return dcerpc.Dial(ctx, target.AddressWithoutPort(), append(dceOpts, ao...)...)
 }
diff --git a/internal/client/dcerpc/client.go b/internal/client/dcerpc/client.go
deleted file mode 100644
index d9d8d71..0000000
--- a/internal/client/dcerpc/client.go
+++ /dev/null
@@ -1,11 +0,0 @@
-package dcerpc
-
-import (
-	"context"
-	"github.com/RedTeamPentesting/adauth"
-)
-
-type Client interface {
-	Connect(ctx context.Context, creds *adauth.Credential, target *adauth.Target) error
-	Close(ctx context.Context) error
-}
diff --git a/internal/client/dcerpc/dcerpc.go b/internal/client/dcerpc/dcerpc.go
deleted file mode 100644
index 5c9a734..0000000
--- a/internal/client/dcerpc/dcerpc.go
+++ /dev/null
@@ -1,106 +0,0 @@
-package dcerpc
-
-import (
-	"context"
-	"errors"
-	"fmt"
-	"github.com/RedTeamPentesting/adauth"
-	"github.com/RedTeamPentesting/adauth/dcerpcauth"
-	"github.com/oiweiwei/go-msrpc/dcerpc"
-	"github.com/oiweiwei/go-msrpc/msrpc/epm/epm/v3"
-	"github.com/oiweiwei/go-msrpc/msrpc/scmr/svcctl/v2"
-	"github.com/oiweiwei/go-msrpc/smb2"
-	"github.com/oiweiwei/go-msrpc/ssp/gssapi"
-	"github.com/rs/zerolog"
-)
-
-const (
-	DceDefaultProto string = "ncacn_np"
-	DceDefaultPort  uint16 = 445
-)
-
-type DCEClient struct {
-	Port  uint16
-	Proto string
-
-	log      zerolog.Logger
-	conn     dcerpc.Conn
-	opts     []dcerpc.Option
-	authOpts dcerpcauth.Options
-}
-
-func NewDCEClient(ctx context.Context, insecure bool, smbConfig *SmbConfig) (client *DCEClient) {
-	client = &DCEClient{
-		Port:     DceDefaultPort,
-		Proto:    DceDefaultProto,
-		log:      zerolog.Ctx(ctx).With().Str("client", "DCE").Logger(),
-		authOpts: dcerpcauth.Options{},
-	}
-	client.opts = []dcerpc.Option{dcerpc.WithLogger(client.log)}
-	client.authOpts = dcerpcauth.Options{Debug: client.log.Trace().Msgf}
-
-	if smbConfig != nil {
-		if smbConfig.Port != 0 {
-			client.Port = smbConfig.Port
-			client.opts = append(client.opts, dcerpc.WithSMBPort(int(smbConfig.Port)))
-		}
-	}
-	if insecure {
-		client.log.Debug().Msg("Using insecure DCERPC connection")
-		client.opts = append(client.opts, dcerpc.WithInsecure())
-	} else {
-		client.log.Debug().Msg("Using secure DCERPC connection")
-		client.authOpts.SMBOptions = append(client.authOpts.SMBOptions, smb2.WithSeal())
-	}
-	return
-}
-
-func (client *DCEClient) OpenSvcctl(ctx context.Context) (ctl svcctl.SvcctlClient, err error) {
-	if client.conn == nil {
-		return nil, errors.New("DCE connection not open")
-	}
-	if ctl, err = svcctl.NewSvcctlClient(ctx, client.conn, dcerpc.WithInsecure()); err != nil {
-		client.log.Debug().Err(err).Msg("Failed to open Svcctl client")
-	}
-	return
-}
-
-func (client *DCEClient) DCE() dcerpc.Conn {
-	return client.conn
-}
-
-func (client *DCEClient) Connect(ctx context.Context, creds *adauth.Credential, target *adauth.Target, dialOpts ...dcerpc.Option) (err error) {
-	if creds != nil && target != nil {
-		authCtx := gssapi.NewSecurityContext(ctx)
-
-		binding := fmt.Sprintf(`%s:%s`, client.Proto, target.AddressWithoutPort())
-		mapper := epm.EndpointMapper(ctx, fmt.Sprintf("%s:%d", target.AddressWithoutPort(), client.Port), dcerpc.WithLogger(client.log))
-		dceOpts := []dcerpc.Option{dcerpc.WithLogger(client.log), dcerpc.WithSeal()}
-
-		if dceOpts, err = dcerpcauth.AuthenticationOptions(authCtx, creds, target, &client.authOpts); err == nil {
-			dceOpts = append(dceOpts, mapper)
-			dceOpts = append(dceOpts, client.opts...)
-			dceOpts = append(dceOpts, dialOpts...)
-
-			if client.conn, err = dcerpc.Dial(authCtx, binding, dceOpts...); err == nil {
-				client.log.Debug().Msg("Bind successful")
-				return nil
-			}
-			client.log.Debug().Err(err).Msg("DCERPC bind failed")
-			return errors.New("bind failed")
-		}
-		return errors.New("unable to parse DCE authentication options")
-	}
-	return errors.New("invalid arguments")
-}
-
-func (client *DCEClient) Close(ctx context.Context) (err error) {
-	if client.conn == nil {
-		client.log.Debug().Msg("Connection already closed")
-	} else if err = client.conn.Close(ctx); err == nil {
-		client.log.Debug().Msg("Connection closed successfully")
-	} else {
-		client.log.Error().Err(err).Msg("Failed to close connection")
-	}
-	return
-}
diff --git a/internal/client/dcerpc/smb.go b/internal/client/dcerpc/smb.go
deleted file mode 100644
index cab82eb..0000000
--- a/internal/client/dcerpc/smb.go
+++ /dev/null
@@ -1,13 +0,0 @@
-package dcerpc
-
-import "github.com/oiweiwei/go-msrpc/smb2"
-
-const (
-	SmbDefaultPort = 445
-)
-
-type SmbConfig struct {
-	Port         uint16
-	FullSecurity bool
-	ForceDialect smb2.Dialect
-}
diff --git a/internal/exec/scmr/exec.go b/internal/exec/scmr/exec.go
index 41c5d8a..c47fcee 100644
--- a/internal/exec/scmr/exec.go
+++ b/internal/exec/scmr/exec.go
@@ -1,226 +1,347 @@
 package scmrexec
 
 import (
-	"context"
-	"errors"
-	"fmt"
-	dcerpc2 "github.com/FalconOpsLLC/goexec/internal/client/dcerpc"
-	"github.com/FalconOpsLLC/goexec/internal/exec"
-	"github.com/FalconOpsLLC/goexec/internal/windows"
-	"github.com/RedTeamPentesting/adauth"
-	"github.com/rs/zerolog"
+  "context"
+  "errors"
+  "fmt"
+  "github.com/FalconOpsLLC/goexec/internal/client/dce"
+  "github.com/FalconOpsLLC/goexec/internal/exec"
+  "github.com/FalconOpsLLC/goexec/internal/util"
+  "github.com/FalconOpsLLC/goexec/internal/windows"
+  "github.com/RedTeamPentesting/adauth"
+  "github.com/oiweiwei/go-msrpc/dcerpc"
+  "github.com/oiweiwei/go-msrpc/midl/uuid"
+  "github.com/oiweiwei/go-msrpc/msrpc/scmr/svcctl/v2"
+  "github.com/rs/zerolog"
 )
 
 const (
-	MethodCreate string = "create"
-	MethodModify string = "modify"
+  DefaultEndpoint = "ncacn_np:[srvsvc]"
+)
 
-	ServiceModifyAccess uint32 = windows.SERVICE_QUERY_CONFIG | windows.SERVICE_CHANGE_CONFIG | windows.SERVICE_STOP | windows.SERVICE_START | windows.SERVICE_DELETE
-	ServiceCreateAccess uint32 = windows.SC_MANAGER_CREATE_SERVICE | windows.SERVICE_START | windows.SERVICE_STOP | windows.SERVICE_DELETE
-	ServiceAllAccess    uint32 = ServiceCreateAccess | ServiceModifyAccess
+var (
+  ScmrRpcUuid = uuid.MustParse("367ABB81-9844-35F1-AD32-98F038001003")
 )
 
-func (mod *Module) createClients(ctx context.Context) (cleanup func(cCtx context.Context), err error) {
-
-	cleanup = func(context.Context) {
-		if mod.dce != nil {
-			mod.log.Debug().Msg("Cleaning up clients")
-			if err := mod.dce.Close(ctx); err != nil {
-				mod.log.Error().Err(err).Msg("Failed to destroy DCE connection")
-			}
-		}
-	}
-	cleanup(ctx)
-	mod.dce = dcerpc2.NewDCEClient(ctx, false, &dcerpc2.SmbConfig{Port: 445})
-	cleanup = func(context.Context) {}
-
-	if err = mod.dce.Connect(ctx, mod.creds, mod.target); err != nil {
-		return nil, fmt.Errorf("connection to DCERPC failed: %w", err)
-	}
-	mod.ctl, err = mod.dce.OpenSvcctl(ctx)
-	return
+func (mod *Module) Connect(ctx context.Context, creds *adauth.Credential, target *adauth.Target, ccfg *exec.ConnectionConfig) (err error) {
+
+  log := zerolog.Ctx(ctx).With().
+    Str("func", "Connect").Logger()
+
+  if ccfg.ConnectionMethod == exec.ConnectionMethodDCE {
+    if cfg, ok := ccfg.ConnectionMethodConfig.(dce.ConnectionMethodDCEConfig); !ok {
+      return fmt.Errorf("invalid configuration for DCE connection method")
+    } else {
+
+      // Fetch target hostname - for opening SCM handle
+      if mod.hostname, err = target.Hostname(ctx); err != nil {
+        log.Debug().Err(err).Msg("Failed to get target hostname")
+        mod.hostname = util.RandomHostname()
+        err = nil
+      }
+      connect := func(ctx context.Context) error {
+        // Create DCE connection
+        if mod.dce, err = cfg.GetDce(ctx, creds, target, dcerpc.WithObjectUUID(ScmrRpcUuid)); err != nil {
+          log.Error().Err(err).Msg("Failed to initialize DCE dialer")
+          return fmt.Errorf("create DCE dialer: %w", err)
+        }
+        log.Info().Msg("DCE dialer initialized")
+
+        // Create SVCCTL client
+        mod.ctl, err = svcctl.NewSvcctlClient(ctx, mod.dce)
+        if err != nil {
+          log.Error().Err(err).Msg("Failed to initialize SCMR client")
+          return fmt.Errorf("init SCMR client: %w", err)
+        }
+        log.Info().Msg("DCE connection successful")
+        return nil
+      }
+      mod.reconnect = func(c context.Context) error {
+        mod.dce = nil
+        mod.ctl = nil
+        return connect(c)
+      }
+      return connect(ctx)
+    }
+  } else {
+    return errors.New("unsupported connection method")
+  }
 }
 
-func (mod *Module) Exec(ctx context.Context, creds *adauth.Credential, target *adauth.Target, ecfg *exec.ExecutionConfig) (err error) {
-
-	vctx := context.WithoutCancel(ctx)
-	mod.log = zerolog.Ctx(ctx).With().
-		Str("module", "scmr").
-		Str("method", ecfg.ExecutionMethod).Logger()
-	mod.creds = creds
-	mod.target = target
-
-	if ecfg.ExecutionMethod == MethodCreate {
-		if cfg, ok := ecfg.ExecutionMethodConfig.(MethodCreateConfig); !ok || cfg.ServiceName == "" {
-			return errors.New("invalid configuration")
-		} else {
-			if cleanup, err := mod.createClients(ctx); err != nil {
-				return fmt.Errorf("failed to create client: %w", err)
-			} else {
-				mod.log.Debug().Msg("Created clients")
-				defer cleanup(ctx)
-			}
-			svc := &service{
-				createConfig: &cfg,
-				name:         cfg.ServiceName,
-			}
-			scm, code, err := mod.openSCM(ctx)
-			if err != nil {
-				return fmt.Errorf("failed to open SCM with code %d: %w", code, err)
-			}
-			mod.log.Debug().Msg("Opened handle to SCM")
-			code, err = mod.createService(ctx, scm, svc, ecfg)
-			if err != nil {
-				return fmt.Errorf("failed to create service with code %d: %w", code, err)
-			}
-			mod.log.Info().Str("service", svc.name).Msg("Service created")
-			// From here on out, make sure the service is properly deleted, even if the connection drops or something fails.
-			if !cfg.NoDelete {
-				defer func() {
-					// TODO: stop service?
-					if code, err = mod.deleteService(ctx, scm, svc); err != nil {
-						mod.log.Error().Err(err).Msg("Failed to delete service") // TODO
-					}
-					mod.log.Info().Str("service", svc.name).Msg("Service deleted successfully")
-				}()
-			}
-			if code, err = mod.startService(ctx, scm, svc); err != nil {
-				if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
-					// In case of timeout or cancel, try to reestablish a connection to restore the service
-					mod.log.Info().Msg("Service start timeout/cancelled. Execution likely successful")
-					mod.log.Info().Msg("Reconnecting for cleanup procedure")
-					ctx = vctx
-
-					if _, err = mod.createClients(ctx); err != nil {
-						mod.log.Error().Err(err).Msg("Reconnect failed")
-
-					} else if scm, code, err = mod.openSCM(ctx); err != nil {
-						mod.log.Error().Err(err).Msg("Failed to reopen SCM")
-
-					} else if svc.handle, code, err = mod.openService(ctx, scm, svc.name); err != nil {
-						mod.log.Error().Str("service", svc.name).Err(err).Msg("Failed to reopen service handle")
-
-					} else {
-						mod.log.Debug().Str("service", svc.name).Msg("Reconnection successful")
-					}
-				} else {
-					mod.log.Error().Err(err).Msg("Failed to start service")
-				}
-			} else {
-				mod.log.Info().Str("service", svc.name).Msg("Execution successful")
-			}
-		}
-	} else if ecfg.ExecutionMethod == MethodModify {
-		// Use service modification method
-		if cfg, ok := ecfg.ExecutionMethodConfig.(MethodModifyConfig); !ok || cfg.ServiceName == "" {
-			return errors.New("invalid configuration")
-
-		} else {
-			// Ensure that a command (executable full path + args) is supplied
-			cmd := ecfg.GetRawCommand()
-			if cmd == "" {
-				return errors.New("no command provided")
-			}
-
-			// Initialize protocol clients
-			if cleanup, err := mod.createClients(ctx); err != nil {
-				return fmt.Errorf("failed to create client: %w", err)
-			} else {
-				mod.log.Debug().Msg("Created clients")
-				defer cleanup(ctx)
-			}
-			svc := &service{modifyConfig: &cfg, name: cfg.ServiceName}
-
-			// Open SCM handle
-			scm, code, err := mod.openSCM(ctx)
-			if err != nil {
-				return fmt.Errorf("failed to create service with code %d: %w", code, err)
-			}
-			mod.log.Debug().Msg("Opened handle to SCM")
-
-			// Open service handle
-			if svc.handle, code, err = mod.openService(ctx, scm, svc.name); err != nil {
-				return fmt.Errorf("failed to open service with code %d: %w", code, err)
-			}
-			mod.log.Debug().Str("service", svc.name).Msg("Opened service")
-
-			// Stop service before editing
-			if !cfg.NoStart {
-				if code, err = mod.stopService(ctx, scm, svc); err != nil {
-					mod.log.Warn().Err(err).Msg("Failed to stop existing service")
-				} else if code == windows.ERROR_SERVICE_NOT_ACTIVE {
-					mod.log.Debug().Str("service", svc.name).Msg("Service is not running")
-				} else {
-					mod.log.Info().Str("service", svc.name).Msg("Stopped existing service")
-					defer func() {
-						if code, err = mod.startService(ctx, scm, svc); err != nil {
-							mod.log.Error().Err(err).Msg("Failed to restore service state to running")
-						}
-					}()
-				}
-			}
-			if code, err = mod.queryServiceConfig(ctx, svc); err != nil {
-				return fmt.Errorf("failed to query service configuration with code %d: %w", code, err)
-			}
-			mod.log.Debug().
-				Str("service", svc.name).
-				Str("command", svc.svcConfig.BinaryPathName).Msg("Fetched existing service configuration")
-
-			// Change service configuration
-			if code, err = mod.changeServiceConfigBinary(ctx, svc, cmd); err != nil {
-				return fmt.Errorf("failed to edit service configuration with code %d: %w", code, err)
-			}
-			defer func() {
-				// Revert configuration
-				if code, err = mod.changeServiceConfigBinary(ctx, svc, svc.svcConfig.BinaryPathName); err != nil {
-					mod.log.Error().Err(err).Msg("Failed to restore service configuration")
-				} else {
-					mod.log.Info().Str("service", svc.name).Msg("Restored service configuration")
-				}
-			}()
-			mod.log.Info().
-				Str("service", svc.name).
-				Str("command", cmd).Msg("Changed service configuration")
-
-			// Start service
-			if !cfg.NoStart {
-				if code, err = mod.startService(ctx, scm, svc); err != nil {
-					if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
-						// In case of timeout or cancel, try to reestablish a connection to restore the service
-						mod.log.Info().Msg("Service start timeout/cancelled. Execution likely successful")
-						mod.log.Info().Msg("Reconnecting for cleanup procedure")
-						ctx = vctx
-
-						if _, err = mod.createClients(ctx); err != nil {
-							mod.log.Error().Err(err).Msg("Reconnect failed")
-
-						} else if scm, code, err = mod.openSCM(ctx); err != nil {
-							mod.log.Error().Err(err).Msg("Failed to reopen SCM")
-
-						} else if svc.handle, code, err = mod.openService(ctx, scm, svc.name); err != nil {
-							mod.log.Error().Str("service", svc.name).Err(err).Msg("Failed to reopen service handle")
-
-						} else {
-							mod.log.Debug().Str("service", svc.name).Msg("Reconnection successful")
-						}
-					} else {
-						mod.log.Error().Err(err).Msg("Failed to start service")
-					}
-				} else {
-					mod.log.Info().Str("service", svc.name).Msg("Started service")
-				}
-				defer func() {
-					// Stop service
-					if code, err = mod.stopService(ctx, scm, svc); err != nil {
-						mod.log.Error().Err(err).Msg("Failed to stop service")
-					} else {
-						mod.log.Info().Str("service", svc.name).Msg("Stopped service")
-					}
-				}()
-			}
-		}
-	} else {
-		return fmt.Errorf("invalid method: %s", ecfg.ExecutionMethod)
-	}
-	return err
+func (mod *Module) Cleanup(ctx context.Context, ccfg *exec.CleanupConfig) (err error) {
+
+  log := zerolog.Ctx(ctx).With().
+    Str("method", ccfg.CleanupMethod).
+    Str("func", "Cleanup").Logger()
+
+  if len(mod.services) == 0 {
+    return nil
+  }
+
+  if mod.dce == nil || mod.ctl == nil {
+    // Try to reconnect
+    if err := mod.reconnect(ctx); err != nil {
+      log.Error().Err(err).Msg("Reconnect failed")
+      return err
+    }
+    log.Info().Msg("Reconnect successful")
+  }
+  if mod.scm == nil {
+    // Open a handle to SCM (again)
+    if resp, err := mod.ctl.OpenSCMW(ctx, &svcctl.OpenSCMWRequest{
+      MachineName:   util.CheckNullString(mod.hostname),
+      DatabaseName:  "ServicesActive\x00",
+      DesiredAccess: ServiceAllAccess, // TODO: Replace
+    }); err != nil {
+      log.Error().Err(err).Msg("Failed to reopen an SCM handle")
+      return err
+    } else {
+      mod.scm = resp.SCM
+      log.Info().Msg("Reopened SCM handle")
+    }
+  }
+
+  for _, rsvc := range mod.services {
+    log = log.With().Str("service", rsvc.name).Logger()
+
+    if rsvc.handle == nil {
+      // Open a handle to the service in question
+      if or, err := mod.ctl.OpenServiceW(ctx, &svcctl.OpenServiceWRequest{
+        ServiceManager: mod.scm,
+        ServiceName:    rsvc.name,
+        DesiredAccess:  windows.SERVICE_DELETE | windows.SERVICE_CHANGE_CONFIG,
+      }); err != nil {
+        log.Error().Err(err).Msg("Failed to reopen a service handle")
+        continue
+      } else {
+        rsvc.handle = or.Service
+      }
+      log.Info().Msg("Service handle reopened")
+    }
+    if ccfg.CleanupMethod == CleanupMethodDelete {
+      // Delete the service
+      if _, err = mod.ctl.DeleteService(ctx, &svcctl.DeleteServiceRequest{Service: rsvc.handle}); err != nil {
+        log.Error().Err(err).Msg("Failed to delete service")
+        continue
+      }
+      log.Info().Msg("Service deleted successfully")
+
+    } else if ccfg.CleanupMethod == CleanupMethodRevert {
+      // Revert the service configuration & state
+      log.Info().Msg("Attempting to revert service configuration")
+      if _, err = mod.ctl.ChangeServiceConfigW(ctx, &svcctl.ChangeServiceConfigWRequest{
+        Service: rsvc.handle,
+        //Dependencies:     []byte(rsvc.originalConfig.Dependencies), // TODO: ensure this works
+        ServiceType:      rsvc.originalConfig.ServiceType,
+        StartType:        rsvc.originalConfig.StartType,
+        ErrorControl:     rsvc.originalConfig.ErrorControl,
+        BinaryPathName:   rsvc.originalConfig.BinaryPathName,
+        LoadOrderGroup:   rsvc.originalConfig.LoadOrderGroup,
+        ServiceStartName: rsvc.originalConfig.ServiceStartName,
+        DisplayName:      rsvc.originalConfig.DisplayName,
+        TagID:            rsvc.originalConfig.TagID,
+      }); err != nil {
+        log.Error().Err(err).Msg("Failed to revert service configuration")
+        continue
+      }
+      log.Info().Msg("Service configuration reverted")
+    }
+    if _, err = mod.ctl.CloseService(ctx, &svcctl.CloseServiceRequest{ServiceObject: rsvc.handle}); err != nil {
+      log.Warn().Err(err).Msg("Failed to close service handle")
+      return nil
+    }
+    log.Info().Msg("Closed service handle")
+  }
+  return
+}
+
+func (mod *Module) Exec(ctx context.Context, ecfg *exec.ExecutionConfig) (err error) {
+
+  //vctx := context.WithoutCancel(ctx)
+  log := zerolog.Ctx(ctx).With().
+    Str("method", ecfg.ExecutionMethod).
+    Str("func", "Exec").Logger()
+
+  if ecfg.ExecutionMethod == MethodCreate {
+    if cfg, ok := ecfg.ExecutionMethodConfig.(MethodCreateConfig); !ok {
+      return errors.New("invalid configuration")
+
+    } else {
+      svc := remoteService{
+        name: cfg.ServiceName,
+      }
+      defer func() { // TODO: relocate this?
+        mod.services = append(mod.services, svc)
+      }()
+      // Open a handle to SCM
+      if resp, err := mod.ctl.OpenSCMW(ctx, &svcctl.OpenSCMWRequest{
+        MachineName:   util.CheckNullString(mod.hostname),
+        DatabaseName:  "ServicesActive\x00",
+        DesiredAccess: ServiceAllAccess, // TODO: Replace
+      }); err != nil {
+        log.Debug().Err(err).Msg("Failed to open SCM handle")
+        return fmt.Errorf("open SCM handle: %w", err)
+      } else {
+        mod.scm = resp.SCM
+        log.Info().Msg("Opened SCM handle")
+      }
+      // Create service
+      serviceName := util.RandomStringIfBlank(svc.name)
+      resp, err := mod.ctl.CreateServiceW(ctx, &svcctl.CreateServiceWRequest{
+        ServiceManager: mod.scm,
+        ServiceName:    serviceName,
+        DisplayName:    util.RandomStringIfBlank(cfg.DisplayName),
+        BinaryPathName: util.CheckNullString(ecfg.GetRawCommand()),
+        ServiceType:    windows.SERVICE_WIN32_OWN_PROCESS,
+        StartType:      windows.SERVICE_DEMAND_START,
+        DesiredAccess:  ServiceAllAccess, // TODO: Replace
+      })
+      if err != nil || resp == nil || resp.Return != 0 {
+        log.Error().Err(err).Msg("Failed to create service")
+        return fmt.Errorf("create service: %w", err)
+      }
+      svc.handle = resp.Service
+
+      log = log.With().
+        Str("service", serviceName).Logger()
+      log.Info().Msg("Service created")
+
+      // Start the service
+      sr, err := mod.ctl.StartServiceW(ctx, &svcctl.StartServiceWRequest{Service: svc.handle})
+      if err != nil {
+
+        if errors.Is(err, context.DeadlineExceeded) { // Check if execution timed out (execute "cmd.exe /c notepad" for test case)
+          log.Warn().Err(err).Msg("Service execution deadline exceeded")
+          // Connection closes, so we nullify the client variables and handles
+          mod.dce = nil
+          mod.ctl = nil
+          mod.scm = nil
+          svc.handle = nil
+
+        } else if sr != nil && sr.Return == windows.ERROR_SERVICE_REQUEST_TIMEOUT { // Check for request timeout
+          log.Info().Msg("Received request timeout. Execution was likely successful")
+
+        } else {
+          log.Error().Err(err).Msg("Failed to start service")
+          return fmt.Errorf("start service: %w", err)
+        }
+        // Inform the caller that execution was likely successful despite error
+        err = nil
+      } else {
+        log.Info().Msg("Started service")
+      }
+    }
+  } else if ecfg.ExecutionMethod == MethodChange {
+    if cfg, ok := ecfg.ExecutionMethodConfig.(MethodChangeConfig); !ok {
+      return errors.New("invalid configuration")
+
+    } else {
+      svc := remoteService{
+        name: cfg.ServiceName,
+      }
+      defer func() { // TODO: relocate this?
+        mod.services = append(mod.services, svc)
+      }()
+
+      // Open a handle to SCM
+      if resp, err := mod.ctl.OpenSCMW(ctx, &svcctl.OpenSCMWRequest{
+        MachineName:   util.CheckNullString(mod.hostname),
+        DatabaseName:  "ServicesActive\x00",
+        DesiredAccess: ServiceAllAccess, // TODO: Replace
+      }); err != nil {
+        log.Debug().Err(err).Msg("Failed to open SCM handle")
+        return fmt.Errorf("open SCM handle: %w", err)
+      } else {
+        mod.scm = resp.SCM
+        log.Info().Msg("Opened SCM handle")
+      }
+
+      // Open a handle to the desired service
+      if resp, err := mod.ctl.OpenServiceW(ctx, &svcctl.OpenServiceWRequest{
+        ServiceManager: mod.scm,
+        ServiceName:    svc.name,
+        DesiredAccess:  ServiceAllAccess, // TODO: Replace
+      }); err != nil {
+        log.Error().Err(err).Msg("Failed to open service handle")
+        return fmt.Errorf("open service: %w", err)
+      } else {
+        svc.handle = resp.Service
+      }
+
+      // Note original service status
+      if resp, err := mod.ctl.QueryServiceStatus(ctx, &svcctl.QueryServiceStatusRequest{
+        Service: svc.handle,
+      }); err != nil {
+        log.Warn().Err(err).Msg("Failed to get service status")
+      } else {
+        svc.originalState = resp.ServiceStatus
+      }
+
+      // Note original service configuration
+      if resp, err := mod.ctl.QueryServiceConfigW(ctx, &svcctl.QueryServiceConfigWRequest{
+        Service:      svc.handle,
+        BufferLength: 8 * 1024,
+      }); err != nil {
+        log.Error().Err(err).Msg("Failed to fetch service configuration")
+        return fmt.Errorf("get service config: %w", err)
+      } else {
+        log.Info().Str("binaryPath", resp.ServiceConfig.BinaryPathName).Msg("Fetched original service configuration")
+        svc.originalConfig = resp.ServiceConfig
+      }
+
+      // Stop service if its running
+      if svc.originalState == nil || svc.originalState.CurrentState != windows.SERVICE_STOPPED {
+        if resp, err := mod.ctl.ControlService(ctx, &svcctl.ControlServiceRequest{
+          Service: svc.handle,
+          Control: windows.SERVICE_STOPPED,
+        }); err != nil {
+          if resp != nil && resp.Return == windows.ERROR_SERVICE_NOT_ACTIVE {
+            log.Info().Msg("Service is already stopped")
+          } else {
+            log.Error().Err(err).Msg("Failed to stop service")
+          }
+        } else {
+          log.Info().Msg("Service stopped")
+        }
+      }
+      // Change service configuration
+      if _, err = mod.ctl.ChangeServiceConfigW(ctx, &svcctl.ChangeServiceConfigWRequest{
+        Service:        svc.handle,
+        BinaryPathName: ecfg.GetRawCommand(),
+        //Dependencies:     []byte(svc.originalConfig.Dependencies), // TODO: ensure this works
+        ServiceType:      svc.originalConfig.ServiceType,
+        StartType:        windows.SERVICE_DEMAND_START,
+        ErrorControl:     svc.originalConfig.ErrorControl,
+        LoadOrderGroup:   svc.originalConfig.LoadOrderGroup,
+        ServiceStartName: svc.originalConfig.ServiceStartName,
+        DisplayName:      svc.originalConfig.DisplayName,
+        TagID:            svc.originalConfig.TagID,
+      }); err != nil {
+        log.Error().Err(err).Msg("Failed to change service configuration")
+        return fmt.Errorf("change service configuration: %w", err)
+      }
+      log.Info().Msg("Successfully altered service configuration")
+
+      if !cfg.NoStart {
+        if resp, err := mod.ctl.StartServiceW(ctx, &svcctl.StartServiceWRequest{Service: svc.handle}); err != nil {
+
+          if errors.Is(err, context.DeadlineExceeded) { // Check if execution timed out (execute "cmd.exe /c notepad" for test case)
+            log.Warn().Err(err).Msg("Service execution deadline exceeded")
+            // Connection closes, so we nullify the client variables and handles
+            mod.dce = nil
+            mod.ctl = nil
+            mod.scm = nil
+            svc.handle = nil
+
+          } else if resp != nil && resp.Return == windows.ERROR_SERVICE_REQUEST_TIMEOUT { // Check for request timeout
+            log.Info().Err(err).Msg("Received request timeout. Execution was likely successful")
+          } else {
+            log.Error().Err(err).Msg("Failed to start service")
+            return fmt.Errorf("start service: %w", err)
+          }
+        }
+      }
+    }
+  }
+  return
 }
diff --git a/internal/exec/scmr/module.go b/internal/exec/scmr/module.go
index 95977b8..2a2d378 100644
--- a/internal/exec/scmr/module.go
+++ b/internal/exec/scmr/module.go
@@ -1,31 +1,38 @@
 package scmrexec
 
 import (
-	"github.com/FalconOpsLLC/goexec/internal/client/dcerpc"
-	"github.com/RedTeamPentesting/adauth"
-	"github.com/oiweiwei/go-msrpc/msrpc/scmr/svcctl/v2"
-	"github.com/rs/zerolog"
+  "context"
+  "github.com/oiweiwei/go-msrpc/dcerpc"
+  "github.com/oiweiwei/go-msrpc/msrpc/scmr/svcctl/v2"
+)
+
+const (
+  MethodCreate = "create"
+  MethodChange = "change"
+
+  CleanupMethodDelete = "delete"
+  CleanupMethodRevert = "revert"
 )
 
 type Module struct {
-	creds    *adauth.Credential
-	target   *adauth.Target
-	hostname string
+  hostname  string // The target hostname
+  dce       dcerpc.Conn
+  reconnect func(context.Context) error
 
-	log zerolog.Logger
-	dce *dcerpc.DCEClient
-	ctl svcctl.SvcctlClient
+  ctl      svcctl.SvcctlClient
+  scm      *svcctl.Handle
+  services []remoteService
 }
 
 type MethodCreateConfig struct {
-	NoDelete    bool
-	ServiceName string
-	DisplayName string
-	ServiceType uint32
-	StartType   uint32
+  NoDelete    bool
+  ServiceName string
+  DisplayName string
+  ServiceType uint32
+  StartType   uint32
 }
 
-type MethodModifyConfig struct {
-	NoStart     bool
-	ServiceName string
+type MethodChangeConfig struct {
+  NoStart     bool
+  ServiceName string
 }
diff --git a/internal/exec/scmr/scmr.go b/internal/exec/scmr/scmr.go
deleted file mode 100644
index 3a6ae08..0000000
--- a/internal/exec/scmr/scmr.go
+++ /dev/null
@@ -1,238 +0,0 @@
-package scmrexec
-
-import (
-	"context"
-	"errors"
-	"fmt"
-	"github.com/FalconOpsLLC/goexec/internal/exec"
-	"github.com/FalconOpsLLC/goexec/internal/util"
-	"github.com/FalconOpsLLC/goexec/internal/windows"
-	"github.com/oiweiwei/go-msrpc/msrpc/scmr/svcctl/v2"
-)
-
-type service struct {
-	name         string
-	exec         string
-	createConfig *MethodCreateConfig
-	modifyConfig *MethodModifyConfig
-
-	svcState  uint32
-	svcConfig *svcctl.QueryServiceConfigW
-	handle    *svcctl.Handle
-}
-
-// openSCM opens a handle to SCM via ROpenSCManagerW
-// https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-scmr/dc84adb3-d51d-48eb-820d-ba1c6ca5faf2
-func (mod *Module) openSCM(ctx context.Context) (scm *svcctl.Handle, code uint32, err error) {
-	if mod.ctl != nil {
-
-		hostname := mod.hostname
-		if hostname == "" {
-			hostname = util.RandomHostname()
-		}
-		if response, err := mod.ctl.OpenSCMW(ctx, &svcctl.OpenSCMWRequest{
-			MachineName:   hostname + "\x00",    // lpMachineName; The server's name (i.e. DC01, dc01.domain.local)
-			DatabaseName:  "ServicesActive\x00", // lpDatabaseName; must be "ServicesActive" or "ServicesFailed"
-			DesiredAccess: ServiceModifyAccess,  // dwDesiredAccess; requested access - appears to be ignored?
-		}); err != nil {
-
-			if response != nil {
-				return nil, response.Return, fmt.Errorf("open scm response: %w", err)
-			}
-			return nil, 0, fmt.Errorf("open scm: %w", err)
-		} else {
-			return response.SCM, 0, nil
-		}
-	}
-	return nil, 0, errors.New("invalid arguments")
-}
-
-// createService creates a service with the provided configuration via RCreateServiceW
-// https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-scmr/6a8ca926-9477-4dd4-b766-692fab07227e
-func (mod *Module) createService(ctx context.Context, scm *svcctl.Handle, scfg *service, ecfg *exec.ExecutionConfig) (code uint32, err error) {
-	if mod.ctl != nil && scm != nil && scfg != nil && scfg.createConfig != nil {
-		cfg := scfg.createConfig
-		if response, err := mod.ctl.CreateServiceW(ctx, &svcctl.CreateServiceWRequest{
-			ServiceManager: scm,
-			ServiceName:    cfg.ServiceName + "\x00",
-			DisplayName:    cfg.DisplayName + "\x00",
-			BinaryPathName: ecfg.GetRawCommand() + "\x00",
-			ServiceType:    cfg.ServiceType,
-			StartType:      cfg.StartType,
-			DesiredAccess:  ServiceCreateAccess,
-		}); err != nil {
-
-			if response != nil {
-				return response.Return, fmt.Errorf("create service response: %w", err)
-			}
-			return 0, fmt.Errorf("create service: %w", err)
-		} else {
-			scfg.handle = response.Service
-			return response.Return, err
-		}
-	}
-	return 0, errors.New("invalid arguments")
-}
-
-// openService opens a handle to a service given the service name (lpServiceName)
-// https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-scmr/6d0a4225-451b-4132-894d-7cef7aecfd2d
-func (mod *Module) openService(ctx context.Context, scm *svcctl.Handle, svcName string) (*svcctl.Handle, uint32, error) {
-	if mod.ctl != nil && scm != nil {
-		if openResponse, err := mod.ctl.OpenServiceW(ctx, &svcctl.OpenServiceWRequest{
-			ServiceManager: scm,
-			ServiceName:    svcName,
-			DesiredAccess:  ServiceAllAccess,
-		}); err != nil {
-			if openResponse != nil {
-				if openResponse.Return == windows.ERROR_SERVICE_DOES_NOT_EXIST {
-					return nil, openResponse.Return, fmt.Errorf("remote service does not exist: %s", svcName)
-				}
-				return nil, openResponse.Return, fmt.Errorf("open service response: %w", err)
-			}
-			return nil, 0, fmt.Errorf("open service: %w", err)
-		} else {
-			return openResponse.Service, 0, nil
-		}
-	}
-	return nil, 0, errors.New("invalid arguments")
-}
-
-// deleteService deletes an existing service with RDeleteService
-// https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-scmr/6744cdb8-f162-4be0-bb31-98996b6495be
-func (mod *Module) deleteService(ctx context.Context, scm *svcctl.Handle, svc *service) (code uint32, err error) {
-	if mod.ctl != nil && scm != nil && svc != nil {
-		if deleteResponse, err := mod.ctl.DeleteService(ctx, &svcctl.DeleteServiceRequest{Service: svc.handle}); err != nil {
-			defer func() {}()
-			if deleteResponse != nil {
-				return deleteResponse.Return, fmt.Errorf("delete service response: %w", err)
-			}
-			return 0, fmt.Errorf("delete service: %w", err)
-		}
-		return 0, nil
-	}
-	return 0, errors.New("invalid arguments")
-}
-
-// controlService sets the state of the provided process
-// https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-scmr/e1c478be-117f-4512-9b67-17c20a48af97
-func (mod *Module) controlService(ctx context.Context, scm *svcctl.Handle, svc *service, control uint32) (code uint32, err error) {
-	if mod.ctl != nil && scm != nil && svc != nil {
-		if controlResponse, err := mod.ctl.ControlService(ctx, &svcctl.ControlServiceRequest{
-			Service: svc.handle,
-			Control: control,
-		}); err != nil {
-			if controlResponse != nil {
-				return controlResponse.Return, fmt.Errorf("control service response: %w", err)
-			}
-			return 0, fmt.Errorf("control service: %w", err)
-		}
-		return 0, nil
-	}
-	return 0, errors.New("invalid arguments")
-}
-
-// stopService sends stop signal to existing service using controlService
-// https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-scmr/e1c478be-117f-4512-9b67-17c20a48af97
-func (mod *Module) stopService(ctx context.Context, scm *svcctl.Handle, svc *service) (code uint32, err error) {
-	if code, err = mod.controlService(ctx, scm, svc, windows.SERVICE_CONTROL_STOP); code == windows.ERROR_SERVICE_NOT_ACTIVE {
-		err = nil
-	}
-	return
-}
-
-// startService starts the specified service with RStartServiceW
-// https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-scmr/d9be95a2-cf01-4bdc-b30f-6fe4b37ada16
-func (mod *Module) startService(ctx context.Context, scm *svcctl.Handle, svc *service) (code uint32, err error) {
-	if mod.ctl != nil && scm != nil && svc != nil {
-		if startResponse, err := mod.ctl.StartServiceW(ctx, &svcctl.StartServiceWRequest{Service: svc.handle}); err != nil {
-			if startResponse != nil {
-				// TODO: check if service is already running, return nil error if so
-				if startResponse.Return == windows.ERROR_SERVICE_REQUEST_TIMEOUT {
-					return 0, nil
-				}
-				return startResponse.Return, fmt.Errorf("start service response: %w", err)
-			}
-			return 0, fmt.Errorf("start service: %w", err)
-		}
-		return 0, nil
-	}
-	return 0, errors.New("invalid arguments")
-}
-
-// closeService closes the specified service handle using RCloseServiceHandle
-// https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-scmr/a2a4e174-09fb-4e55-bad3-f77c4b13245c
-func (mod *Module) closeService(ctx context.Context, svc *svcctl.Handle) (code uint32, err error) {
-	if mod.ctl != nil && svc != nil {
-		if closeResponse, err := mod.ctl.CloseService(ctx, &svcctl.CloseServiceRequest{ServiceObject: svc}); err != nil {
-			if closeResponse != nil {
-				return closeResponse.Return, fmt.Errorf("close service response: %w", err)
-			}
-			return 0, fmt.Errorf("close service: %w", err)
-		}
-		return 0, nil
-	}
-	return 0, errors.New("invalid arguments")
-}
-
-// getServiceConfig fetches the configuration details of a service given a handle passed in a service{} structure.
-// https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-scmr/89e2d5b1-19cf-44ca-969f-38eea9fe7f3c
-func (mod *Module) queryServiceConfig(ctx context.Context, svc *service) (code uint32, err error) {
-	if mod.ctl != nil && svc != nil && svc.handle != nil {
-		if getResponse, err := mod.ctl.QueryServiceConfigW(ctx, &svcctl.QueryServiceConfigWRequest{
-			Service:      svc.handle,
-			BufferLength: 1024 * 8,
-		}); err != nil {
-			if getResponse != nil {
-				return getResponse.Return, fmt.Errorf("get service config response: %w", err)
-			}
-			return 0, fmt.Errorf("get service config: %w", err)
-		} else {
-			svc.svcConfig = getResponse.ServiceConfig
-			return code, err
-		}
-	}
-	return 0, errors.New("invalid arguments")
-}
-
-// queryServiceStatus fetches the state of the specified service
-// https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-scmr/cf94d915-b4e1-40e5-872b-a9cb3ad09b46
-func (mod *Module) queryServiceStatus(ctx context.Context, svc *service) (uint32, error) {
-	if mod.ctl != nil && svc != nil {
-		if queryResponse, err := mod.ctl.QueryServiceStatus(ctx, &svcctl.QueryServiceStatusRequest{Service: svc.handle}); err != nil {
-			if queryResponse != nil {
-				return queryResponse.Return, fmt.Errorf("query service status response: %w", err)
-			}
-			return 0, fmt.Errorf("query service status: %w", err)
-		} else {
-			svc.svcState = queryResponse.ServiceStatus.CurrentState
-			return 0, nil
-		}
-	}
-	return 0, errors.New("invalid arguments")
-}
-
-// changeServiceConfigBinary edits the provided service's lpBinaryPathName
-// https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-scmr/61ea7ed0-c49d-4152-a164-b4830f16c8a4
-func (mod *Module) changeServiceConfigBinary(ctx context.Context, svc *service, bin string) (code uint32, err error) {
-	if mod.ctl != nil && svc != nil && svc.handle != nil {
-		if changeResponse, err := mod.ctl.ChangeServiceConfigW(ctx, &svcctl.ChangeServiceConfigWRequest{
-			Service:        svc.handle,
-			ServiceType:    svc.svcConfig.ServiceType,
-			StartType:      svc.svcConfig.StartType,
-			ErrorControl:   svc.svcConfig.ErrorControl,
-			BinaryPathName: bin + "\x00",
-			LoadOrderGroup: svc.svcConfig.LoadOrderGroup,
-			TagID:          svc.svcConfig.TagID,
-			// Dependencies: svc.svcConfig.Dependencies // TODO
-			ServiceStartName: svc.svcConfig.ServiceStartName,
-			DisplayName:      svc.svcConfig.DisplayName,
-		}); err != nil {
-			if changeResponse != nil {
-				return changeResponse.Return, fmt.Errorf("change service config response: %w", err)
-			}
-			return 0, fmt.Errorf("change service config: %w", err)
-		}
-		return
-	}
-	return 0, errors.New("invalid arguments")
-}
diff --git a/internal/exec/scmr/service.go b/internal/exec/scmr/service.go
new file mode 100644
index 0000000..9a580cb
--- /dev/null
+++ b/internal/exec/scmr/service.go
@@ -0,0 +1,25 @@
+package scmrexec
+
+import (
+  "context"
+  "github.com/FalconOpsLLC/goexec/internal/windows"
+  "github.com/oiweiwei/go-msrpc/msrpc/scmr/svcctl/v2"
+)
+
+const (
+  ServiceDeleteAccess uint32 = windows.SERVICE_DELETE
+  ServiceModifyAccess uint32 = windows.SERVICE_QUERY_CONFIG | windows.SERVICE_CHANGE_CONFIG | windows.SERVICE_STOP | windows.SERVICE_START | windows.SERVICE_DELETE
+  ServiceCreateAccess uint32 = windows.SC_MANAGER_CREATE_SERVICE | windows.SERVICE_START | windows.SERVICE_STOP | windows.SERVICE_DELETE
+  ServiceAllAccess    uint32 = ServiceCreateAccess | ServiceModifyAccess
+)
+
+type remoteService struct {
+  name           string
+  handle         *svcctl.Handle
+  originalConfig *svcctl.QueryServiceConfigW
+  originalState  *svcctl.ServiceStatus
+}
+
+func (mod *Module) parseServiceDependencies(ctx context.Context, ) (err error) {
+  return nil
+}
diff --git a/internal/exec/tsch/exec.go b/internal/exec/tsch/exec.go
index 45af819..44f11d1 100644
--- a/internal/exec/tsch/exec.go
+++ b/internal/exec/tsch/exec.go
@@ -1,227 +1,183 @@
 package tschexec
 
 import (
-	"context"
-	"errors"
-	"fmt"
-	"github.com/FalconOpsLLC/goexec/internal/client/dce"
-	"github.com/FalconOpsLLC/goexec/internal/exec"
-	"github.com/FalconOpsLLC/goexec/internal/util"
-	"github.com/RedTeamPentesting/adauth"
-	"github.com/RedTeamPentesting/adauth/dcerpcauth"
-	"github.com/oiweiwei/go-msrpc/dcerpc"
-	"github.com/oiweiwei/go-msrpc/midl/uuid"
-	"github.com/oiweiwei/go-msrpc/msrpc/epm/epm/v3"
-	"github.com/oiweiwei/go-msrpc/msrpc/tsch/itaskschedulerservice/v1"
-	"github.com/oiweiwei/go-msrpc/ssp/gssapi"
-	"github.com/rs/zerolog"
-	"time"
+  "context"
+  "errors"
+  "fmt"
+  "github.com/FalconOpsLLC/goexec/internal/client/dce"
+  "github.com/FalconOpsLLC/goexec/internal/exec"
+  "github.com/FalconOpsLLC/goexec/internal/util"
+  "github.com/RedTeamPentesting/adauth"
+  "github.com/oiweiwei/go-msrpc/dcerpc"
+  "github.com/oiweiwei/go-msrpc/midl/uuid"
+  "github.com/oiweiwei/go-msrpc/msrpc/tsch/itaskschedulerservice/v1"
+  "github.com/rs/zerolog"
+  "time"
 )
 
 const (
-	DefaultEndpoint = "ncacn_np:[atsvc]"
+  DefaultEndpoint = "ncacn_np:[atsvc]"
 )
 
 var (
-	TschRpcUuid                = uuid.MustParse("86D35949-83C9-4044-B424-DB363231FD0C")
-	SupportedEndpointProtocols = []string{"ncacn_np", "ncacn_ip_tcp"}
+  TschRpcUuid = uuid.MustParse("86D35949-83C9-4044-B424-DB363231FD0C")
 )
 
 // Connect to the target & initialize DCE & TSCH clients
 func (mod *Module) Connect(ctx context.Context, creds *adauth.Credential, target *adauth.Target, ccfg *exec.ConnectionConfig) (err error) {
 
-	//var port uint16
-	var endpoint string = DefaultEndpoint
-	var epmOpts []dcerpc.Option
-	var dceOpts []dcerpc.Option
-
-	log := zerolog.Ctx(ctx).With().
-		Str("func", "Connect").Logger()
-
-	if mod.dce == nil {
-		if ccfg.ConnectionMethod == exec.ConnectionMethodDCE {
-			if cfg, ok := ccfg.ConnectionMethodConfig.(dce.ConnectionMethodDCEConfig); !ok {
-				return fmt.Errorf("invalid configuration for DCE connection method")
-			} else {
-				// Connect to ITaskSchedulerService
-				{
-					// Parse target & creds
-					ctx = gssapi.NewSecurityContext(ctx)
-					ao, err := dcerpcauth.AuthenticationOptions(ctx, creds, target, &dcerpcauth.Options{})
-					if err != nil {
-						log.Error().Err(err).Msg("Failed to parse authentication options")
-						return fmt.Errorf("parse auth options: %w", err)
-					}
-					dceOpts = append(cfg.Options,
-						dcerpc.WithLogger(log),
-						dcerpc.WithSecurityLevel(dcerpc.AuthLevelPktPrivacy), // AuthLevelPktPrivacy is required for TSCH/ATSVC
-						dcerpc.WithObjectUUID(TschRpcUuid))
-
-					if cfg.Endpoint != nil {
-						endpoint = cfg.Endpoint.String()
-					}
-					if cfg.NoEpm {
-						dceOpts = append(dceOpts, dcerpc.WithEndpoint(endpoint))
-					} else {
-						epmOpts = append(epmOpts, dceOpts...)
-						dceOpts = append(dceOpts,
-							epm.EndpointMapper(ctx, target.AddressWithoutPort(), append(epmOpts, ao...)...))
-						if !cfg.EpmAuto {
-							dceOpts = append(dceOpts, dcerpc.WithEndpoint(endpoint))
-						}
-					}
-					log = log.With().Str("endpoint", endpoint).Logger()
-					log.Info().Msg("Connecting to target")
-
-					// Create DCERPC dialer
-					mod.dce, err = dcerpc.Dial(ctx, target.AddressWithoutPort(), append(dceOpts, ao...)...)
-					if err != nil {
-						log.Error().Err(err).Msg("Failed to create DCERPC dialer")
-						return fmt.Errorf("create DCERPC dialer: %w", err)
-					}
-
-					// Create ITaskSchedulerService
-					mod.tsch, err = itaskschedulerservice.NewTaskSchedulerServiceClient(ctx, mod.dce)
-					if err != nil {
-						log.Error().Err(err).Msg("Failed to initialize TSCH client")
-						return fmt.Errorf("init TSCH client: %w", err)
-					}
-					log.Info().Msg("DCE connection successful")
-				}
-			}
-		} else {
-			return errors.New("unsupported connection method")
-		}
-	}
-	return
+  log := zerolog.Ctx(ctx).With().
+    Str("func", "Connect").Logger()
+
+  if ccfg.ConnectionMethod == exec.ConnectionMethodDCE {
+    if cfg, ok := ccfg.ConnectionMethodConfig.(dce.ConnectionMethodDCEConfig); !ok {
+      return fmt.Errorf("invalid configuration for DCE connection method")
+    } else {
+      // Create DCERPC dialer
+      mod.dce, err = cfg.GetDce(ctx, creds, target, dcerpc.WithObjectUUID(TschRpcUuid))
+      if err != nil {
+        log.Error().Err(err).Msg("Failed to create DCERPC dialer")
+        return fmt.Errorf("create DCERPC dialer: %w", err)
+      }
+      // Create ITaskSchedulerService client
+      mod.tsch, err = itaskschedulerservice.NewTaskSchedulerServiceClient(ctx, mod.dce)
+      if err != nil {
+        log.Error().Err(err).Msg("Failed to initialize TSCH client")
+        return fmt.Errorf("init TSCH client: %w", err)
+      }
+      log.Info().Msg("DCE connection successful")
+    }
+  } else {
+    return errors.New("unsupported connection method")
+  }
+  return
 }
 
 func (mod *Module) Cleanup(ctx context.Context, ccfg *exec.CleanupConfig) (err error) {
-	log := zerolog.Ctx(ctx).With().
-		Str("method", ccfg.CleanupMethod).
-		Str("func", "Cleanup").Logger()
-
-	if ccfg.CleanupMethod == MethodDelete {
-		if cfg, ok := ccfg.CleanupMethodConfig.(MethodDeleteConfig); !ok {
-			return errors.New("invalid configuration")
-
-		} else {
-			log = log.With().Str("task", cfg.TaskPath).Logger()
-			log.Info().Msg("Manually deleting task")
-
-			if err = mod.deleteTask(ctx, cfg.TaskPath); err == nil {
-				log.Info().Msg("Task deleted successfully")
-			}
-		}
-	} else if ccfg.CleanupMethod == "" {
-		return nil
-	} else {
-		return fmt.Errorf("unsupported cleanup method")
-	}
-	return
+  log := zerolog.Ctx(ctx).With().
+    Str("method", ccfg.CleanupMethod).
+    Str("func", "Cleanup").Logger()
+
+  if ccfg.CleanupMethod == MethodDelete {
+    if cfg, ok := ccfg.CleanupMethodConfig.(MethodDeleteConfig); !ok {
+      return errors.New("invalid configuration")
+
+    } else {
+      log = log.With().Str("task", cfg.TaskPath).Logger()
+      log.Info().Msg("Manually deleting task")
+
+      if err = mod.deleteTask(ctx, cfg.TaskPath); err == nil {
+        log.Info().Msg("Task deleted successfully")
+      }
+    }
+  } else if ccfg.CleanupMethod == "" {
+    return nil
+  } else {
+    return fmt.Errorf("unsupported cleanup method")
+  }
+  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 == MethodRegister {
-		if cfg, ok := ecfg.ExecutionMethodConfig.(MethodRegisterConfig); !ok {
-			return errors.New("invalid configuration")
-
-		} else {
-			startTime := time.Now().UTC().Add(cfg.StartDelay)
-			stopTime := startTime.Add(cfg.StopDelay)
-
-			tr := taskTimeTrigger{
-				StartBoundary: startTime.Format(TaskXMLDurationFormat),
-				//EndBoundary:   stopTime.Format(TaskXMLDurationFormat),
-				Enabled: true,
-			}
-			tk := newTask(nil, nil, triggers{TimeTriggers: []taskTimeTrigger{tr}}, ecfg.ExecutableName, ecfg.ExecutableArgs)
-
-			if !cfg.NoDelete && !cfg.CallDelete {
-				if cfg.StopDelay == 0 {
-					cfg.StopDelay = time.Second
-				}
-				tk.Settings.DeleteExpiredTaskAfter = xmlDuration(cfg.DeleteDelay)
-				tk.Triggers.TimeTriggers[0].EndBoundary = stopTime.Format(TaskXMLDurationFormat)
-			}
-			taskPath := cfg.TaskPath
-			if taskPath == "" {
-				log.Debug().Msg("Task path not defined. Using random path")
-				taskPath = `\` + util.RandomString()
-			}
-			// The taskPath is changed here to the *actual path returned by SchRpcRegisterTask
-			taskPath, err = mod.registerTask(ctx, *tk, taskPath)
-			if err != nil {
-				return fmt.Errorf("call registerTask: %w", err)
-			}
-
-			if !cfg.NoDelete {
-				if cfg.CallDelete {
-					defer mod.deleteTask(ctx, taskPath)
-
-					log.Info().Dur("ms", cfg.StartDelay).Msg("Waiting for task to run")
-					select {
-					case <-ctx.Done():
-						log.Warn().Msg("Cancelling execution")
-						return err
-					case <-time.After(cfg.StartDelay + time.Second): // + one second for good measure
-						for {
-							if stat, err := mod.tsch.GetLastRunInfo(ctx, &itaskschedulerservice.GetLastRunInfoRequest{Path: taskPath}); err != nil {
-								log.Warn().Err(err).Msg("Failed to get last run info. Assuming task was executed")
-								break
-							} else if stat.LastRuntime.AsTime().IsZero() {
-								log.Warn().Msg("Task was not yet run. Waiting 10 additional seconds")
-								time.Sleep(10 * time.Second)
-							} else {
-								break
-							}
-						}
-						break
-					}
-				} else {
-					log.Info().Time("when", stopTime).Msg("Task is scheduled to delete")
-				}
-			}
-		}
-	} else if ecfg.ExecutionMethod == MethodDemand {
-		if cfg, ok := ecfg.ExecutionMethodConfig.(MethodDemandConfig); !ok {
-			return errors.New("invalid configuration")
-
-		} else {
-			taskPath := cfg.TaskPath
-			if taskPath == "" {
-				log.Debug().Msg("Task path not defined. Using random path")
-				taskPath = `\` + util.RandomString()
-			}
-			st := newSettings(true, true, false)
-			tk := newTask(st, nil, triggers{}, ecfg.ExecutableName, ecfg.ExecutableArgs)
-
-			// The taskPath is changed here to the *actual path returned by SchRpcRegisterTask
-			taskPath, err = mod.registerTask(ctx, *tk, taskPath)
-			if err != nil {
-				return fmt.Errorf("call registerTask: %w", err)
-			}
-			if !cfg.NoDelete {
-				defer mod.deleteTask(ctx, taskPath)
-			}
-			_, err := mod.tsch.Run(ctx, &itaskschedulerservice.RunRequest{
-				Path:  taskPath,
-				Flags: 0, // Maybe we want to use these?
-			})
-			if err != nil {
-				log.Error().Str("task", taskPath).Err(err).Msg("Failed to run task")
-				return fmt.Errorf("force run task: %w", err)
-			}
-			log.Info().Str("task", taskPath).Msg("Started task")
-		}
-	} else {
-		return fmt.Errorf("method '%s' not implemented", ecfg.ExecutionMethod)
-	}
-
-	return nil
+  log := zerolog.Ctx(ctx).With().
+    Str("method", ecfg.ExecutionMethod).
+    Str("func", "Exec").Logger()
+
+  if ecfg.ExecutionMethod == MethodRegister {
+    if cfg, ok := ecfg.ExecutionMethodConfig.(MethodRegisterConfig); !ok {
+      return errors.New("invalid configuration")
+
+    } else {
+      startTime := time.Now().UTC().Add(cfg.StartDelay)
+      stopTime := startTime.Add(cfg.StopDelay)
+
+      tr := taskTimeTrigger{
+        StartBoundary: startTime.Format(TaskXMLDurationFormat),
+        Enabled:       true,
+      }
+      tk := newTask(nil, nil, triggers{TimeTriggers: []taskTimeTrigger{tr}}, ecfg.ExecutableName, ecfg.ExecutableArgs)
+
+      if !cfg.NoDelete && !cfg.CallDelete {
+        if cfg.StopDelay == 0 {
+          cfg.StopDelay = time.Second
+        }
+        tk.Settings.DeleteExpiredTaskAfter = xmlDuration(cfg.DeleteDelay)
+        tk.Triggers.TimeTriggers[0].EndBoundary = stopTime.Format(TaskXMLDurationFormat)
+      }
+      taskPath := cfg.TaskPath
+      if taskPath == "" {
+        log.Debug().Msg("Task path not defined. Using random path")
+        taskPath = `\` + util.RandomString()
+      }
+      // The taskPath is changed here to the *actual path returned by SchRpcRegisterTask
+      taskPath, err = mod.registerTask(ctx, *tk, taskPath)
+      if err != nil {
+        return fmt.Errorf("call registerTask: %w", err)
+      }
+
+      if !cfg.NoDelete {
+        if cfg.CallDelete {
+          defer mod.deleteTask(ctx, taskPath)
+
+          log.Info().Dur("ms", cfg.StartDelay).Msg("Waiting for task to run")
+          select {
+          case <-ctx.Done():
+            log.Warn().Msg("Cancelling execution")
+            return err
+          case <-time.After(cfg.StartDelay + time.Second): // + one second for good measure
+            for {
+              if stat, err := mod.tsch.GetLastRunInfo(ctx, &itaskschedulerservice.GetLastRunInfoRequest{Path: taskPath}); err != nil {
+                log.Warn().Err(err).Msg("Failed to get last run info. Assuming task was executed")
+                break
+              } else if stat.LastRuntime.AsTime().IsZero() {
+                log.Warn().Msg("Task was not yet run. Waiting 10 additional seconds")
+                time.Sleep(10 * time.Second)
+              } else {
+                break
+              }
+            }
+            break
+          }
+        } else {
+          log.Info().Time("when", stopTime).Msg("Task is scheduled to delete")
+        }
+      }
+    }
+  } else if ecfg.ExecutionMethod == MethodDemand {
+    if cfg, ok := ecfg.ExecutionMethodConfig.(MethodDemandConfig); !ok {
+      return errors.New("invalid configuration")
+
+    } else {
+      taskPath := cfg.TaskPath
+      if taskPath == "" {
+        log.Debug().Msg("Task path not defined. Using random path")
+        taskPath = `\` + util.RandomString()
+      }
+      st := newSettings(true, true, false)
+      tk := newTask(st, nil, triggers{}, ecfg.ExecutableName, ecfg.ExecutableArgs)
+
+      // The taskPath is changed here to the *actual path returned by SchRpcRegisterTask
+      taskPath, err = mod.registerTask(ctx, *tk, taskPath)
+      if err != nil {
+        return fmt.Errorf("call registerTask: %w", err)
+      }
+      if !cfg.NoDelete {
+        defer mod.deleteTask(ctx, taskPath)
+      }
+      _, err := mod.tsch.Run(ctx, &itaskschedulerservice.RunRequest{
+        Path:  taskPath,
+        Flags: 0, // Maybe we want to use these?
+      })
+      if err != nil {
+        log.Error().Str("task", taskPath).Err(err).Msg("Failed to run task")
+        return fmt.Errorf("force run task: %w", err)
+      }
+      log.Info().Str("task", taskPath).Msg("Started task")
+    }
+  } else {
+    return fmt.Errorf("method '%s' not implemented", ecfg.ExecutionMethod)
+  }
+
+  return nil
 }
diff --git a/internal/exec/tsch/module.go b/internal/exec/tsch/module.go
index 8e31840..6e379a7 100644
--- a/internal/exec/tsch/module.go
+++ b/internal/exec/tsch/module.go
@@ -11,8 +11,6 @@ type Module struct {
   dce dcerpc.Conn
   // tsch holds the ITaskSchedulerService client
   tsch itaskschedulerservice.TaskSchedulerServiceClient
-  // createdTasks holds any tasks that are created - for cleanup
-  createdTasks []string
 }
 
 type MethodRegisterConfig struct {
diff --git a/internal/util/util.go b/internal/util/util.go
index 36d7ea2..f8f12ba 100644
--- a/internal/util/util.go
+++ b/internal/util/util.go
@@ -3,6 +3,7 @@ package util
 import (
   "math/rand" // not crypto secure
   "regexp"
+  "strings"
 )
 
 const randHostnameCharset = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-"
@@ -33,3 +34,17 @@ func RandomStringFromCharset(charset string, length int) string {
   }
   return string(b)
 }
+
+func RandomStringIfBlank(s string) string {
+  if s == "" {
+    return RandomString()
+  }
+  return s
+}
+
+func CheckNullString(s string) string {
+  if !strings.HasSuffix(s, "\x00") {
+    return s + "\x00"
+  }
+  return s
+}