More logging options

3 files changed, 142 insertions, 127 deletions

diff --git a/cmd/root.go b/cmd/root.go
index 6b5a416..c185be8 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -11,18 +11,26 @@ import (
 	"github.com/oiweiwei/go-msrpc/ssp/gssapi"
 	"github.com/rs/zerolog"
 	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+	"io"
 	"os"
 )
 
 var (
-	debug        bool
-	logJson      bool
 	returnCode   int
 	outputMethod string
 	outputPath   string
 	proxy        string
 
-	log zerolog.Logger
+	// === Logging ===
+	logJson   bool          // Log output in JSON lines
+	logDebug  bool          // Output debug log messages
+	logQuiet  bool          // Suppress logging output
+	logOutput string        // Log output file
+	logLevel  zerolog.Level = zerolog.InfoLevel
+	logFile   io.Writer     = os.Stderr
+	log       zerolog.Logger
+	// ===============
 
 	rpcClient dce.Client
 	smbClient smb.Client
@@ -39,19 +47,28 @@ var (
 	rootCmd = &cobra.Command{
 		Use:   "goexec",
 		Short: `Windows remote execution multitool`,
-		Long:  `TODO`,
+		Long:  ``,
 
-		PersistentPreRun: func(cmd *cobra.Command, args []string) {
+		PersistentPreRunE: func(cmd *cobra.Command, args []string) (err error) {
 
-			if logJson {
-				log = zerolog.New(os.Stderr)
-			} else {
-				log = zerolog.New(zerolog.ConsoleWriter{Out: os.Stderr})
-			}
-
-			log = log.Level(zerolog.InfoLevel).With().Timestamp().Logger()
-			if debug {
-				log = log.Level(zerolog.DebugLevel)
+			// Parse logging options
+			{
+				if logOutput != "" {
+					logFile, err = os.OpenFile(logOutput, os.O_WRONLY|os.O_CREATE, 0644)
+					if err != nil {
+						return
+					}
+				}
+				if logQuiet {
+					logLevel = zerolog.ErrorLevel
+				} else if logDebug {
+					logLevel = zerolog.DebugLevel
+				}
+				if logJson {
+					log = zerolog.New(logFile)
+				} else {
+					log = zerolog.New(zerolog.ConsoleWriter{Out: logFile}).With().Timestamp().Logger()
+				}
 			}
 
 			if proxy != "" {
@@ -72,6 +89,7 @@ var (
 					}
 				}
 			}
+			return
 		},
 	}
 )
@@ -83,9 +101,23 @@ func init() {
 
 		rootCmd.InitDefaultVersionFlag()
 		rootCmd.InitDefaultHelpCmd()
-		rootCmd.PersistentFlags().BoolVar(&debug, "debug", false, "Enable debug logging")
-		rootCmd.PersistentFlags().BoolVar(&logJson, "log-json", false, "Log in JSON format")
-		rootCmd.PersistentFlags().StringVarP(&proxy, "proxy", "x", "", "Proxy URL")
+
+		// Logging flags
+		{
+			logOpts := pflag.NewFlagSet("Logging", pflag.ExitOnError)
+			logOpts.BoolVar(&logDebug, "debug", false, "Enable debug logging")
+			logOpts.BoolVar(&logJson, "json", false, "Write logging output in JSON lines")
+			logOpts.BoolVar(&logQuiet, "quiet", false, "Disable info logging")
+			logOpts.StringVarP(&logOutput, "log-out", "O", "", "Write logging output to file")
+			rootCmd.PersistentFlags().AddFlagSet(logOpts)
+		}
+
+		// Global networking flags
+		{
+			netOpts := pflag.NewFlagSet("Network", pflag.ExitOnError)
+			netOpts.StringVarP(&proxy, "proxy", "x", "", "Proxy URL")
+			rootCmd.PersistentFlags().AddFlagSet(netOpts)
+		}
 
 		dcomCmdInit()
 		rootCmd.AddCommand(dcomCmd)
diff --git a/internal/util/util.go b/internal/util/util.go
index 555297a..d88ec28 100644
--- a/internal/util/util.go
+++ b/internal/util/util.go
@@ -1,55 +1,41 @@
 package util
 
 import (
-  "github.com/google/uuid"
-  "math/rand" // not crypto secure
-  "regexp"
-  "strings"
+	"github.com/google/uuid"
+	"math/rand" // not crypto secure
+	"regexp"
+	"strings"
 )
 
 const randHostnameCharset = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-"
 const randStringCharset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
 
 var (
-  // Up to 15 characters; only letters, digits, and hyphens (with hyphens not at the start or end).
-  randHostnameRegex = regexp.MustCompile(`^[a-zA-Z][a-zA-Z0-9-]{0,14}[a-zA-Z0-9]$`)
+	// Up to 15 characters; only letters, digits, and hyphens (with hyphens not at the start or end).
+	randHostnameRegex = regexp.MustCompile(`^[a-zA-Z][a-zA-Z0-9-]{0,14}[a-zA-Z0-9]$`)
 )
 
 func RandomHostname() (hostname string) {
-  for {
-    // between 2 and 10 characters
-    if hostname = RandomStringFromCharset(randHostnameCharset, rand.Intn(8)+2); randHostnameRegex.MatchString(hostname) {
-      return
-    }
-  }
+	for {
+		// between 2 and 10 characters
+		if hostname = RandomStringFromCharset(randHostnameCharset, rand.Intn(8)+2); randHostnameRegex.MatchString(hostname) {
+			return
+		}
+	}
 }
 
 func RandomWindowsTempFile() string {
-  return `\Windows\Temp\` + strings.ToUpper(uuid.New().String())
+	return `\Windows\Temp\` + strings.ToUpper(uuid.New().String())
 }
 
 func RandomString() string {
-  return RandomStringFromCharset(randStringCharset, rand.Intn(10)+6)
+	return RandomStringFromCharset(randStringCharset, rand.Intn(10)+6)
 }
 
 func RandomStringFromCharset(charset string, length int) string {
-  b := make([]byte, length)
-  for i := range length {
-    b[i] = charset[rand.Intn(len(charset))]
-  }
-  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
+	b := make([]byte, length)
+	for i := range length {
+		b[i] = charset[rand.Intn(len(charset))]
+	}
+	return string(b)
 }
diff --git a/pkg/goexec/smb/client.go b/pkg/goexec/smb/client.go
index 35fe5c8..3b41e39 100644
--- a/pkg/goexec/smb/client.go
+++ b/pkg/goexec/smb/client.go
@@ -1,115 +1,112 @@
 package smb
 
 import (
-  "context"
-  "errors"
-  "fmt"
-  "github.com/oiweiwei/go-smb2.fork"
-  "github.com/rs/zerolog"
-  "net"
+	"context"
+	"errors"
+	"fmt"
+	"github.com/oiweiwei/go-smb2.fork"
+	"github.com/rs/zerolog"
+	"net"
 )
 
 type Client struct {
-  ClientOptions
+	ClientOptions
 
-  conn  net.Conn
-  sess  *smb2.Session
-  mount *smb2.Share
+	conn  net.Conn
+	sess  *smb2.Session
+	mount *smb2.Share
 }
 
 func (c *Client) Session() (sess *smb2.Session) {
-  return c.sess
+	return c.sess
 }
 
 func (c *Client) String() string {
-  return ClientName
+	return ClientName
 }
 
 func (c *Client) Logger(ctx context.Context) zerolog.Logger {
-  return zerolog.Ctx(ctx).With().Str("client", c.String()).Logger()
+	return zerolog.Ctx(ctx).With().Str("client", c.String()).Logger()
 }
 
 func (c *Client) Mount(ctx context.Context, share string) (err error) {
 
-  if c.sess == nil {
-    return errors.New("SMB session not initialized")
-  }
+	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")
+	c.mount, err = c.sess.Mount(share)
+	zerolog.Ctx(ctx).Debug().Str("share", share).Msg("Mounted SMB share")
 
-  return
+	return
 }
 
 func (c *Client) Connect(ctx context.Context) (err error) {
 
-  log := c.Logger(ctx)
-  {
-    if c.netDialer == nil {
-      panic(fmt.Errorf("TCP dialer not initialized"))
-    }
-    if c.dialer == nil {
-      panic(fmt.Errorf("%s dialer not initialized", c.String()))
-    }
-  }
+	log := c.Logger(ctx)
+	{
+		if c.netDialer == nil {
+			panic(fmt.Errorf("TCP dialer not initialized"))
+		}
+		if c.dialer == nil {
+			panic(fmt.Errorf("%s dialer not initialized", c.String()))
+		}
+	}
 
-  // Establish TCP connection
-  c.conn, err = c.netDialer.Dial("tcp", net.JoinHostPort(c.Host, fmt.Sprintf("%d", c.Port)))
+	// Establish TCP connection
+	c.conn, err = c.netDialer.Dial("tcp", net.JoinHostPort(c.Host, fmt.Sprintf("%d", c.Port)))
 
-  if err != nil {
-    return err
-  }
+	if err != nil {
+		return err
+	}
 
-  log = log.With().Str("address", c.conn.RemoteAddr().String()).Logger()
-  log.Debug().Msgf("Connected to %s server", c.String())
+	log = log.With().Str("address", c.conn.RemoteAddr().String()).Logger()
+	log.Debug().Msgf("Connected to %s server", c.String())
 
-  // Open SMB session
-  c.sess, err = c.dialer.DialContext(ctx, c.conn)
+	// Open SMB session
+	c.sess, err = c.dialer.DialContext(ctx, c.conn)
 
-  if err != nil {
-    log.Error().Err(err).Msgf("Failed to open %s session", c.String())
-    return fmt.Errorf("dial %s: %w", c.String(), err)
-  }
+	if err != nil {
+		log.Error().Err(err).Msgf("Failed to open %s session", c.String())
+		return fmt.Errorf("dial %s: %w", c.String(), err)
+	}
 
-  log.Debug().Msgf("Opened %s session", c.String())
+	log.Debug().Msgf("Opened %s session", c.String())
 
-  return
+	return
 }
 
 func (c *Client) Close(ctx context.Context) (err error) {
 
-  log := c.Logger(ctx)
-
-  // Close TCP connection - Not needed?
-  /*
-     if c.conn != nil {
-       defer func() {
-         if err = c.conn.Close(); err != nil {
-           log.Debug().Err(err).Msgf("Failed to close %s connection", c.String())
-         }
-         log.Debug().Msgf("Closed %s connection", c.String())
-       }()
-     }
-  */
-
-  // Close SMB session
-  if c.sess != nil {
-    defer func() {
-      if err = c.sess.Logoff(); err != nil {
-        log.Debug().Err(err).Msgf("Failed to discard %s session", c.String())
-      }
-      log.Debug().Msgf("Discarded %s session", c.String())
-    }()
-  }
-
-  // Unmount SMB share
-  if c.mount != nil {
-    defer func() {
-      if err = c.mount.Umount(); err != nil {
-        log.Debug().Err(err).Msg("Failed to unmount share")
-      }
-      log.Debug().Msg("Unmounted file share")
-    }()
-  }
-  return
+	log := c.Logger(ctx)
+
+	// Close SMB session
+	if c.sess != nil {
+		defer func() {
+			if err = c.sess.Logoff(); err != nil {
+				log.Debug().Err(err).Msgf("Failed to discard SMB session")
+			}
+			log.Debug().Msgf("Discarded SMB session")
+		}()
+
+	} else if c.conn != nil {
+
+		defer func() {
+			if err = c.conn.Close(); err != nil {
+				log.Debug().Err(err).Msgf("Failed to disconnect SMB client")
+			}
+			log.Debug().Msgf("Disconnected SMB session")
+		}()
+	}
+
+	// Unmount SMB share
+	if c.mount != nil {
+		defer func() {
+			if err = c.mount.Umount(); err != nil {
+				log.Debug().Err(err).Msg("Failed to unmount share")
+			}
+			log.Debug().Msg("Unmounted file share")
+		}()
+	}
+	return
 }