move args to config package

logger package rewrite as dlog

65 files changed, 1236 insertions, 425 deletions

diff --git a/internal/clients/args.go b/internal/clients/args.go
deleted file mode 100644
index dc71e83..0000000
--- a/internal/clients/args.go
+++ /dev/null
@@ -1,63 +0,0 @@
-package clients
-
-import (
-	"flag"
-	"fmt"
-	"strings"
-
-	"github.com/mimecast/dtail/internal/omode"
-
-	gossh "golang.org/x/crypto/ssh"
-)
-
-// Args is a helper struct to summarize common client arguments.
-type Args struct {
-	Mode               omode.Mode
-	ServersStr         string
-	UserName           string
-	What               string
-	Arguments          []string
-	RegexStr           string
-	RegexInvert        bool
-	TrustAllHosts      bool
-	Discovery          string
-	ConnectionsPerCPU  int
-	Timeout            int
-	SSHAuthMethods     []gossh.AuthMethod
-	SSHHostKeyCallback gossh.HostKeyCallback
-	PrivateKeyPathFile string
-	Quiet              bool
-	Spartan            bool
-	NoColor            bool
-	Serverless         bool
-}
-
-// Transform the arguments based on certain conditions.
-func (a *Args) Transform(args []string) {
-	// Interpret additional args as file list.
-	if a.What == "" {
-		var files []string
-		for _, file := range flag.Args() {
-			files = append(files, file)
-		}
-		a.What = strings.Join(files, ",")
-	}
-
-	if a.Spartan {
-		a.Quiet = true
-		a.NoColor = true
-	}
-
-	if a.Discovery == "" && a.ServersStr == "" {
-		a.Serverless = true
-	}
-}
-
-// TransformAfterConfigFile same as Transform, but after the config file has been read.
-func (a *Args) TransformAfterConfigFile() {
-	// TODO: Remove this method. It's not used.
-}
-
-func (a *Args) SerializeOptions() string {
-	return fmt.Sprintf("quiet=%v:spartan=%v", a.Quiet, a.Spartan)
-}
diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go
index 5523052..fc01955 100644
--- a/internal/clients/baseclient.go
+++ b/internal/clients/baseclient.go
@@ -6,8 +6,9 @@ import (
 	"time"
 
 	"github.com/mimecast/dtail/internal/clients/connectors"
+	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/discovery"
-	"github.com/mimecast/dtail/internal/io/logger"
+	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/omode"
 	"github.com/mimecast/dtail/internal/regex"
 	"github.com/mimecast/dtail/internal/ssh/client"
@@ -17,7 +18,7 @@ import (
 
 // This is the main client data structure.
 type baseClient struct {
-	Args
+	config.Args
 	// To display client side stats
 	stats *stats
 	// List of remote servers to connect to.
@@ -39,7 +40,8 @@ type baseClient struct {
 }
 
 func (c *baseClient) init() {
-	logger.Debug("Initiating base client")
+	dlog.Client.Debug("Initiating base client")
+	dlog.Client.Debug(c.Args.String())
 
 	flag := regex.Default
 	if c.Args.RegexInvert {
@@ -47,11 +49,14 @@ func (c *baseClient) init() {
 	}
 	regex, err := regex.New(c.Args.RegexStr, flag)
 	if err != nil {
-		logger.FatalExit(c.Regex, "invalid regex!", err, regex)
+		dlog.Client.FatalPanic(c.Regex, "invalid regex!", err, regex)
 	}
 	c.Regex = regex
-	logger.Debug("Regex", c.Regex)
+	dlog.Client.Debug("Regex", c.Regex)
 
+	if c.Args.Serverless {
+		return
+	}
 	c.sshAuthMethods, c.hostKeyCallback = client.InitSSHAuthMethods(c.Args.SSHAuthMethods, c.Args.SSHHostKeyCallback, c.Args.TrustAllHosts, c.throttleCh, c.Args.PrivateKeyPathFile)
 }
 
@@ -67,8 +72,11 @@ func (c *baseClient) makeConnections(maker maker) {
 }
 
 func (c *baseClient) Start(ctx context.Context, statsCh <-chan string) (status int) {
-	// Periodically check for unknown hosts, and ask the user whether to trust them or not.
-	go c.hostKeyCallback.PromptAddHosts(ctx)
+	// Can be nil when serverless.
+	if c.hostKeyCallback != nil {
+		// Periodically check for unknown hosts, and ask the user whether to trust them or not.
+		go c.hostKeyCallback.PromptAddHosts(ctx)
+	}
 	// Print client stats every time something on statsCh is recieved.
 	go c.stats.Start(ctx, c.throttleCh, statsCh, c.Args.Quiet)
 
@@ -112,7 +120,7 @@ func (c *baseClient) start(ctx context.Context, active chan struct{}, i int, con
 		}
 
 		time.Sleep(time.Second * 2)
-		logger.Debug(conn.Server(), "Reconnecting")
+		dlog.Client.Debug(conn.Server(), "Reconnecting")
 		conn = c.makeConnection(conn.Server(), c.sshAuthMethods, c.hostKeyCallback)
 		c.connections[i] = conn
 	}
@@ -127,7 +135,7 @@ func (c *baseClient) makeConnection(server string, sshAuthMethods []gossh.AuthMe
 }
 
 func (c *baseClient) waitUntilDone(ctx context.Context, active chan struct{}) {
-	defer logger.Debug("Terminated connection")
+	defer dlog.Client.Debug("Terminated connection")
 
 	// We want to have at least one active connection
 	<-active
@@ -143,7 +151,7 @@ func (c *baseClient) waitUntilDone(ctx context.Context, active chan struct{}) {
 		if numActive == 0 {
 			return
 		}
-		logger.Debug("Active connections", numActive)
+		dlog.Client.Debug("Active connections", numActive)
 		time.Sleep(time.Second)
 	}
 }
diff --git a/internal/clients/catclient.go b/internal/clients/catclient.go
index d14cdcc..2726e7e 100644
--- a/internal/clients/catclient.go
+++ b/internal/clients/catclient.go
@@ -7,6 +7,8 @@ import (
 	"strings"
 
 	"github.com/mimecast/dtail/internal/clients/handlers"
+	"github.com/mimecast/dtail/internal/config"
+	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/omode"
 )
 
@@ -16,7 +18,7 @@ type CatClient struct {
 }
 
 // NewCatClient returns a new cat client.
-func NewCatClient(args Args) (*CatClient, error) {
+func NewCatClient(args config.Args) (*CatClient, error) {
 	if args.RegexStr != "" {
 		return nil, errors.New("Can't use regex with 'cat' operating mode")
 	}
@@ -42,11 +44,13 @@ func (c CatClient) makeHandler(server string) handlers.Handler {
 }
 
 func (c CatClient) makeCommands() (commands []string) {
+	regex, err := c.Regex.Serialize()
+	if err != nil {
+		dlog.Client.FatalPanic(err)
+	}
 	for _, file := range strings.Split(c.What, ",") {
 		commands = append(commands, fmt.Sprintf("%s:%s %s %s",
-			c.Mode.String(),
-			c.Args.SerializeOptions(),
-			file, c.Regex.Serialize()))
+			c.Mode.String(), c.Args.SerializeOptions(), file, regex))
 	}
 	return
 }
diff --git a/internal/clients/connectors/serverconnection.go b/internal/clients/connectors/serverconnection.go
index 0904ba1..5bc63ee 100644
--- a/internal/clients/connectors/serverconnection.go
+++ b/internal/clients/connectors/serverconnection.go
@@ -10,7 +10,7 @@ import (
 
 	"github.com/mimecast/dtail/internal/clients/handlers"
 	"github.com/mimecast/dtail/internal/config"
-	"github.com/mimecast/dtail/internal/io/logger"
+	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/ssh/client"
 
 	"golang.org/x/crypto/ssh"
@@ -30,7 +30,7 @@ type ServerConnection struct {
 
 // NewServerConnection returns a new DTail SSH server connection.
 func NewServerConnection(server string, userName string, authMethods []ssh.AuthMethod, hostKeyCallback client.HostKeyCallback, handler handlers.Handler, commands []string) *ServerConnection {
-	logger.Debug(server, "Creating new connection", server, handler, commands)
+	dlog.Client.Debug(server, "Creating new connection", server, handler, commands)
 
 	c := ServerConnection{
 		hostKeyCallback: hostKeyCallback,
@@ -81,10 +81,10 @@ func (c *ServerConnection) initServerPort() {
 	parts := strings.Split(c.server, ":")
 
 	if len(parts) == 2 {
-		logger.Debug("Parsing port from hostname", parts)
+		dlog.Client.Debug("Parsing port from hostname", parts)
 		port, err := strconv.Atoi(parts[1])
 		if err != nil {
-			logger.FatalExit("Unable to parse client port", c.server, parts, err)
+			dlog.Client.FatalPanic("Unable to parse client port", c.server, parts, err)
 		}
 		c.server = parts[0]
 		c.port = port
@@ -93,21 +93,21 @@ func (c *ServerConnection) initServerPort() {
 
 func (c *ServerConnection) Start(ctx context.Context, cancel context.CancelFunc, throttleCh, statsCh chan struct{}) {
 	// Throttle how many connections can be established concurrently (based on ch length)
-	logger.Debug(c.server, "Throttling connection", len(throttleCh), cap(throttleCh))
+	dlog.Client.Debug(c.server, "Throttling connection", len(throttleCh), cap(throttleCh))
 
 	select {
 	case throttleCh <- struct{}{}:
 	case <-ctx.Done():
-		logger.Debug(c.server, "Not establishing connection as context is done", len(throttleCh), cap(throttleCh))
+		dlog.Client.Debug(c.server, "Not establishing connection as context is done", len(throttleCh), cap(throttleCh))
 		return
 	}
 
-	logger.Debug(c.server, "Throttling says that the connection can be established", len(throttleCh), cap(throttleCh))
+	dlog.Client.Debug(c.server, "Throttling says that the connection can be established", len(throttleCh), cap(throttleCh))
 
 	go func() {
 		defer func() {
 			if !c.throttlingDone {
-				logger.Debug(c.server, "Unthrottling connection (1)", len(throttleCh), cap(throttleCh))
+				dlog.Client.Debug(c.server, "Unthrottling connection (1)", len(throttleCh), cap(throttleCh))
 				c.throttlingDone = true
 				<-throttleCh
 			}
@@ -115,9 +115,9 @@ func (c *ServerConnection) Start(ctx context.Context, cancel context.CancelFunc,
 		}()
 
 		if err := c.dial(ctx, cancel, throttleCh, statsCh); err != nil {
-			logger.Warn(c.server, c.port, err)
+			dlog.Client.Warn(c.server, c.port, err)
 			if c.hostKeyCallback.Untrusted(fmt.Sprintf("%s:%d", c.server, c.port)) {
-				logger.Debug(c.server, "Not trusting host")
+				dlog.Client.Debug(c.server, "Not trusting host")
 			}
 		}
 	}()
@@ -127,14 +127,14 @@ func (c *ServerConnection) Start(ctx context.Context, cancel context.CancelFunc,
 
 // Dail into a new SSH connection. Close connection in case of an error.
 func (c *ServerConnection) dial(ctx context.Context, cancel context.CancelFunc, throttleCh, statsCh chan struct{}) error {
-	logger.Debug(c.server, "Incrementing connection stats")
+	dlog.Client.Debug(c.server, "Incrementing connection stats")
 	statsCh <- struct{}{}
 	defer func() {
-		logger.Debug(c.server, "Decrementing connection stats")
+		dlog.Client.Debug(c.server, "Decrementing connection stats")
 		<-statsCh
 	}()
 
-	logger.Debug(c.server, "Dialing into the connection")
+	dlog.Client.Debug(c.server, "Dialing into the connection")
 	address := fmt.Sprintf("%s:%d", c.server, c.port)
 
 	client, err := ssh.Dial("tcp", address, c.config)
@@ -148,7 +148,7 @@ func (c *ServerConnection) dial(ctx context.Context, cancel context.CancelFunc,
 
 // Create the SSH session. Close the session in case of an error.
 func (c *ServerConnection) session(ctx context.Context, cancel context.CancelFunc, client *ssh.Client, throttleCh chan struct{}) error {
-	logger.Debug(c.server, "Creating SSH session")
+	dlog.Client.Debug(c.server, "Creating SSH session")
 
 	session, err := client.NewSession()
 	if err != nil {
@@ -160,7 +160,7 @@ func (c *ServerConnection) session(ctx context.Context, cancel context.CancelFun
 }
 
 func (c *ServerConnection) handle(ctx context.Context, cancel context.CancelFunc, session *ssh.Session, throttleCh chan struct{}) error {
-	logger.Debug(c.server, "Creating handler for SSH session")
+	dlog.Client.Debug(c.server, "Creating handler for SSH session")
 
 	stdinPipe, err := session.StdinPipe()
 	if err != nil {
@@ -196,12 +196,12 @@ func (c *ServerConnection) handle(ctx context.Context, cancel context.CancelFunc
 
 	// Send all commands to client.
 	for _, command := range c.commands {
-		logger.Debug(command)
+		dlog.Client.Debug(command)
 		c.handler.SendMessage(command)
 	}
 
 	if !c.throttlingDone {
-		logger.Debug(c.server, "Unthrottling connection (2)", len(throttleCh), cap(throttleCh))
+		dlog.Client.Debug(c.server, "Unthrottling connection (2)", len(throttleCh), cap(throttleCh))
 		c.throttlingDone = true
 		<-throttleCh
 	}
diff --git a/internal/clients/connectors/serverless.go b/internal/clients/connectors/serverless.go
index 0500645..c7b5f62 100644
--- a/internal/clients/connectors/serverless.go
+++ b/internal/clients/connectors/serverless.go
@@ -6,7 +6,7 @@ import (
 
 	"github.com/mimecast/dtail/internal/clients/handlers"
 	"github.com/mimecast/dtail/internal/config"
-	"github.com/mimecast/dtail/internal/io/logger"
+	"github.com/mimecast/dtail/internal/io/dlog"
 	serverHandlers "github.com/mimecast/dtail/internal/server/handlers"
 	user "github.com/mimecast/dtail/internal/user/server"
 )
@@ -26,7 +26,7 @@ func NewServerless(userName string, handler handlers.Handler, commands []string)
 		commands: commands,
 	}
 
-	logger.Debug("Creating new serverless connector", handler, commands)
+	dlog.Client.Debug("Creating new serverless connector", handler, commands)
 	return &s
 }
 
@@ -43,7 +43,7 @@ func (s *Serverless) Start(ctx context.Context, cancel context.CancelFunc, throt
 		defer cancel()
 
 		if err := s.handle(ctx, cancel); err != nil {
-			logger.Warn(err)
+			dlog.Client.Warn(err)
 		}
 	}()
 
@@ -51,7 +51,7 @@ func (s *Serverless) Start(ctx context.Context, cancel context.CancelFunc, throt
 }
 
 func (s *Serverless) handle(ctx context.Context, cancel context.CancelFunc) error {
-	logger.Debug("Creating server handler for a serverless session")
+	dlog.Client.Debug("Creating server handler for a serverless session")
 
 	serverHandler := serverHandlers.NewServerHandler(
 		user.New(s.userName, s.Server()),
@@ -59,14 +59,19 @@ func (s *Serverless) handle(ctx context.Context, cancel context.CancelFunc) erro
 		make(chan struct{}, config.Server.MaxConcurrentTails),
 	)
 
+	terminate := func() {
+		serverHandler.Shutdown()
+		cancel()
+	}
+
 	go func() {
 		io.Copy(serverHandler, s.handler)
-		cancel()
+		terminate()
 	}()
 
 	go func() {
 		io.Copy(s.handler, serverHandler)
-		cancel()
+		terminate()
 	}()
 
 	go func() {
@@ -74,12 +79,12 @@ func (s *Serverless) handle(ctx context.Context, cancel context.CancelFunc) erro
 		case <-s.handler.Done():
 		case <-ctx.Done():
 		}
-		cancel()
+		terminate()
 	}()
 
 	// Send all commands to client.
 	for _, command := range s.commands {
-		logger.Debug(command)
+		dlog.Client.Debug(command)
 		s.handler.SendMessage(command)
 	}
 
diff --git a/internal/clients/grepclient.go b/internal/clients/grepclient.go
index ea5022b..ae21ff2 100644
--- a/internal/clients/grepclient.go
+++ b/internal/clients/grepclient.go
@@ -7,6 +7,8 @@ import (
 	"strings"
 
 	"github.com/mimecast/dtail/internal/clients/handlers"
+	"github.com/mimecast/dtail/internal/config"
+	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/omode"
 )
 
@@ -16,7 +18,7 @@ type GrepClient struct {
 }
 
 // NewGrepClient creates a new grep client.
-func NewGrepClient(args Args) (*GrepClient, error) {
+func NewGrepClient(args config.Args) (*GrepClient, error) {
 	if args.RegexStr == "" {
 		return nil, errors.New("No regex specified, use '-regex' flag")
 	}
@@ -41,12 +43,13 @@ func (c GrepClient) makeHandler(server string) handlers.Handler {
 }
 
 func (c GrepClient) makeCommands() (commands []string) {
+	regex, err := c.Regex.Serialize()
+	if err != nil {
+		dlog.Client.FatalPanic(err)
+	}
 	for _, file := range strings.Split(c.What, ",") {
 		commands = append(commands, fmt.Sprintf("%s:%s %s %s",
-			c.Mode.String(),
-			c.Args.SerializeOptions(),
-			file,
-			c.Regex.Serialize()))
+			c.Mode.String(), c.Args.SerializeOptions(), file, regex))
 	}
 
 	return
diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go
index af1ad62..3291b43 100644
--- a/internal/clients/handlers/basehandler.go
+++ b/internal/clients/handlers/basehandler.go
@@ -9,7 +9,7 @@ import (
 	"time"
 
 	"github.com/mimecast/dtail/internal"
-	"github.com/mimecast/dtail/internal/io/logger"
+	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/protocol"
 )
 
@@ -51,7 +51,7 @@ func (h *baseHandler) Shutdown() {
 // SendMessage to the server.
 func (h *baseHandler) SendMessage(command string) error {
 	encoded := base64.StdEncoding.EncodeToString([]byte(command))
-	logger.Debug("Sending command", h.server, command, encoded)
+	dlog.Client.Debug("Sending command", h.server, command, encoded)
 
 	select {
 	case h.commands <- fmt.Sprintf("protocol %s base64 %v;", protocol.ProtocolCompat, encoded):
@@ -112,7 +112,7 @@ func (h *baseHandler) handleMessageType(message string) {
 		return
 	}
 
-	logger.Raw(message)
+	dlog.Client.Raw(message)
 }
 
 // Handle messages received from server which are not meant to be displayed
diff --git a/internal/clients/handlers/clienthandler.go b/internal/clients/handlers/clienthandler.go
index 2bcb038..27ac85e 100644
--- a/internal/clients/handlers/clienthandler.go
+++ b/internal/clients/handlers/clienthandler.go
@@ -2,7 +2,7 @@ package handlers
 
 import (
 	"github.com/mimecast/dtail/internal"
-	"github.com/mimecast/dtail/internal/io/logger"
+	"github.com/mimecast/dtail/internal/io/dlog"
 )
 
 // ClientHandler is the basic client handler interface.
@@ -12,7 +12,7 @@ type ClientHandler struct {
 
 // NewClientHandler creates a new client handler.
 func NewClientHandler(server string) *ClientHandler {
-	logger.Debug(server, "Creating new client handler")
+	dlog.Client.Debug(server, "Creating new client handler")
 
 	return &ClientHandler{
 		baseHandler{
diff --git a/internal/clients/handlers/maprhandler.go b/internal/clients/handlers/maprhandler.go
index 65b1454..848e7f0 100644
--- a/internal/clients/handlers/maprhandler.go
+++ b/internal/clients/handlers/maprhandler.go
@@ -4,7 +4,7 @@ import (
 	"strings"
 
 	"github.com/mimecast/dtail/internal"
-	"github.com/mimecast/dtail/internal/io/logger"
+	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/mapr"
 	"github.com/mimecast/dtail/internal/mapr/client"
 	"github.com/mimecast/dtail/internal/protocol"
@@ -40,7 +40,7 @@ func (h *MaprHandler) Write(p []byte) (n int, err error) {
 			continue
 		case protocol.MessageDelimiter:
 			message := h.baseHandler.receiveBuf.String()
-			logger.Debug(message)
+			dlog.Client.Debug(message)
 			if message[0] == 'A' {
 				h.handleAggregateMessage(message)
 			} else {
@@ -60,10 +60,10 @@ func (h *MaprHandler) Write(p []byte) (n int, err error) {
 func (h *MaprHandler) handleAggregateMessage(message string) {
 	parts := strings.SplitN(message, protocol.FieldDelimiter, 3)
 	if len(parts) != 3 {
-		logger.Error("Unable to aggregate data", h.server, message, parts, len(parts), "expected 3 parts")
+		dlog.Client.Error("Unable to aggregate data", h.server, message, parts, len(parts), "expected 3 parts")
 		return
 	}
 	if err := h.aggregate.Aggregate(parts[2]); err != nil {
-		logger.Error("Unable to aggregate data", h.server, message, err)
+		dlog.Client.Error("Unable to aggregate data", h.server, message, err)
 	}
 }
diff --git a/internal/clients/maprclient.go b/internal/clients/maprclient.go
index e6ab96b..f23aa08 100644
--- a/internal/clients/maprclient.go
+++ b/internal/clients/maprclient.go
@@ -11,7 +11,7 @@ import (
 	"github.com/mimecast/dtail/internal/clients/handlers"
 	"github.com/mimecast/dtail/internal/color"
 	"github.com/mimecast/dtail/internal/config"
-	"github.com/mimecast/dtail/internal/io/logger"
+	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/mapr"
 	"github.com/mimecast/dtail/internal/omode"
 )