4 files changed, 45 insertions, 143 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
-}