Move commands to cmd/ and move internal dependencies to internal/

67 files changed, 6793 insertions, 0 deletions

diff --git a/internal/clients/args.go b/internal/clients/args.go
new file mode 100644
index 0000000..5fe0a72
--- /dev/null
+++ b/internal/clients/args.go
@@ -0,0 +1,18 @@
+package clients
+
+import (
+	"github.com/mimecast/dtail/internal/omode"
+)
+
+// Args is a helper struct to summarize common client arguments.
+type Args struct {
+	Mode              omode.Mode
+	ServersStr        string
+	UserName          string
+	Files             string
+	Regex             string
+	TrustAllHosts     bool
+	Discovery         string
+	ConnectionsPerCPU int
+	PingTimeout       int
+}
diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go
new file mode 100644
index 0000000..574ae94
--- /dev/null
+++ b/internal/clients/baseclient.go
@@ -0,0 +1,137 @@
+package clients
+
+import (
+	"regexp"
+	"sync"
+	"time"
+
+	"github.com/mimecast/dtail/internal/clients/remote"
+	"github.com/mimecast/dtail/internal/discovery"
+	"github.com/mimecast/dtail/internal/logger"
+	"github.com/mimecast/dtail/internal/omode"
+	"github.com/mimecast/dtail/internal/ssh/client"
+
+	gossh "golang.org/x/crypto/ssh"
+)
+
+// This is the main client data structure.
+type baseClient struct {
+	Args
+	// To display client side stats
+	stats *stats
+	// List of remote servers to connect to.
+	servers []string
+	// We have one connection per remote server.
+	connections []*remote.Connection
+	// SSH auth methods to use to connect to the remote servers.
+	sshAuthMethods []gossh.AuthMethod
+	// To deal with SSH host keys
+	hostKeyCallback *client.HostKeyCallback
+	// To stop the client.
+	stop chan struct{}
+	// To indicate that the client has stopped.
+	stopped chan struct{}
+	// Throttle how fast we initiate SSH connections concurrently
+	throttleCh chan struct{}
+	// Retry connection upon failure?
+	retry bool
+	// Connection helper.
+	maker connectionMaker
+}
+
+func (c *baseClient) init(maker connectionMaker) {
+	logger.Info("Initiating base client")
+
+	c.maker = maker
+	//c.connections = make(map[string]*remote.Connection)
+	c.sshAuthMethods, c.hostKeyCallback = client.InitSSHAuthMethods(c.TrustAllHosts, c.throttleCh)
+
+	// Retrieve a shuffled list of remote dtail servers.
+	shuffleServers := true
+	discoveryService := discovery.New(c.Discovery, c.ServersStr, shuffleServers)
+	for _, server := range discoveryService.ServerList() {
+		c.connections = append(c.connections, c.maker.makeConnection(server, c.sshAuthMethods, c.hostKeyCallback))
+	}
+
+	if _, err := regexp.Compile(c.Regex); err != nil {
+		logger.FatalExit(c.Regex, "Can't test compile regex", err)
+	}
+
+	// Periodically check for unknown hosts, and ask the user whether to trust them or not.
+	go c.hostKeyCallback.PromptAddHosts(c.stop)
+
+	// Periodically print out connection stats to the client.
+	c.stats = newTailStats(len(c.connections))
+	go c.stats.periodicLogStats(c.throttleCh, c.stop)
+}
+
+func (c *baseClient) Start() (status int) {
+	active := make(chan struct{}, len(c.connections))
+
+	var wg sync.WaitGroup
+	wg.Add(len(c.connections))
+
+	for i, conn := range c.connections {
+		go func(i int, conn *remote.Connection) {
+			active <- struct{}{}
+			defer func() {
+				logger.Debug(conn.Server, "Disconnected completely...")
+				<-active
+			}()
+			wg.Done()
+
+			for {
+				conn.Start(c.throttleCh, c.stats.connectionsEstCh)
+				if !c.retry {
+					return
+				}
+				time.Sleep(time.Second * 2)
+				logger.Debug(conn.Server, "Reconencting")
+				conn = c.maker.makeConnection(conn.Server, c.sshAuthMethods, c.hostKeyCallback)
+				c.connections[i] = conn
+			}
+		}(i, conn)
+	}
+
+	wg.Wait()
+	c.waitUntilDone(active)
+
+	return
+}
+
+func (c *baseClient) waitUntilDone(active chan struct{}) {
+	defer close(c.stopped)
+
+	if c.Mode != omode.TailClient {
+		c.waitUntilZero(active)
+		logger.Info("All connections stopped")
+		return
+	}
+
+	<-c.stop
+	logger.Info("Stopping client")
+	for _, conn := range c.connections {
+		conn.Stop()
+	}
+
+	c.waitUntilZero(active)
+}
+
+func (c *baseClient) waitUntilZero(active chan struct{}) {
+	for {
+		logger.Debug("Active connections", len(active))
+		if len(active) == 0 {
+			return
+		}
+		time.Sleep(time.Second)
+	}
+}
+
+func (c *baseClient) Stop() {
+	close(c.stop)
+	<-c.WaitC()
+}
+
+func (c *baseClient) WaitC() <-chan struct{} {
+	return c.stopped
+}
diff --git a/internal/clients/catclient.go b/internal/clients/catclient.go
new file mode 100644
index 0000000..5ea701d
--- /dev/null
+++ b/internal/clients/catclient.go
@@ -0,0 +1,53 @@
+package clients
+
+import (
+	"errors"
+	"fmt"
+	"runtime"
+	"strings"
+
+	"github.com/mimecast/dtail/internal/clients/handlers"
+	"github.com/mimecast/dtail/internal/clients/remote"
+	"github.com/mimecast/dtail/internal/omode"
+	"github.com/mimecast/dtail/internal/ssh/client"
+
+	gossh "golang.org/x/crypto/ssh"
+)
+
+// CatClient is a client for returning a whole file from the beginning to the end.
+type CatClient struct {
+	baseClient
+}
+
+// NewCatClient returns a new cat client.
+func NewCatClient(args Args) (*CatClient, error) {
+	if args.Regex != "" {
+		return nil, errors.New("Can't use regex with 'cat' operating mode")
+	}
+
+	args.Regex = "."
+	args.Mode = omode.CatClient
+
+	c := CatClient{
+		baseClient: baseClient{
+			Args:       args,
+			stop:       make(chan struct{}),
+			stopped:    make(chan struct{}),
+			throttleCh: make(chan struct{}, args.ConnectionsPerCPU*runtime.NumCPU()),
+			retry:      false,
+		},
+	}
+
+	c.init(c)
+
+	return &c, nil
+}
+
+func (c CatClient) makeConnection(server string, sshAuthMethods []gossh.AuthMethod, hostKeyCallback *client.HostKeyCallback) *remote.Connection {
+	conn := remote.NewConnection(server, c.UserName, sshAuthMethods, hostKeyCallback)
+	conn.Handler = handlers.NewClientHandler(server, c.PingTimeout)
+	for _, file := range strings.Split(c.Files, ",") {
+		conn.Commands = append(conn.Commands, fmt.Sprintf("%s %s regex %s", c.Mode.String(), file, c.Regex))
+	}
+	return conn
+}
diff --git a/internal/clients/client.go b/internal/clients/client.go
new file mode 100644
index 0000000..85d1aae
--- /dev/null
+++ b/internal/clients/client.go
@@ -0,0 +1,7 @@
+package clients
+
+// Client is the interface for the end user command line client.
+type Client interface {
+	Start() int
+	Stop()
+}
diff --git a/internal/clients/connectionmaker.go b/internal/clients/connectionmaker.go
new file mode 100644
index 0000000..0617992
--- /dev/null
+++ b/internal/clients/connectionmaker.go
@@ -0,0 +1,12 @@
+package clients
+
+import (
+	"github.com/mimecast/dtail/internal/clients/remote"
+	"github.com/mimecast/dtail/internal/ssh/client"
+
+	gossh "golang.org/x/crypto/ssh"
+)
+
+type connectionMaker interface {
+	makeConnection(server string, sshAuthMethods []gossh.AuthMethod, hostKeyCallback *client.HostKeyCallback) *remote.Connection
+}
diff --git a/internal/clients/grepclient.go b/internal/clients/grepclient.go
new file mode 100644
index 0000000..c568f63
--- /dev/null
+++ b/internal/clients/grepclient.go
@@ -0,0 +1,53 @@
+package clients
+
+import (
+	"errors"
+	"fmt"
+	"runtime"
+	"strings"
+
+	"github.com/mimecast/dtail/internal/clients/handlers"
+	"github.com/mimecast/dtail/internal/clients/remote"
+	"github.com/mimecast/dtail/internal/omode"
+	"github.com/mimecast/dtail/internal/ssh/client"
+
+	gossh "golang.org/x/crypto/ssh"
+)
+
+// GrepClient searches a remote file for all lines matching a regular expression. Only the matching lines are displayed.
+type GrepClient struct {
+	baseClient
+}
+
+// NewGrepClient creates a new grep client.
+func NewGrepClient(args Args) (*GrepClient, error) {
+	if args.Regex == "" {
+		return nil, errors.New("No regex specified, use '-regex' flag")
+	}
+	args.Mode = omode.GrepClient
+
+	c := GrepClient{
+		baseClient: baseClient{
+			Args:       args,
+			stop:       make(chan struct{}),
+			stopped:    make(chan struct{}),
+			throttleCh: make(chan struct{}, args.ConnectionsPerCPU*runtime.NumCPU()),
+			retry:      false,
+		},
+	}
+
+	c.init(c)
+
+	return &c, nil
+}
+
+func (c GrepClient) makeConnection(server string, sshAuthMethods []gossh.AuthMethod, hostKeyCallback *client.HostKeyCallback) *remote.Connection {
+	conn := remote.NewConnection(server, c.UserName, sshAuthMethods, hostKeyCallback)
+	conn.Handler = handlers.NewClientHandler(server, c.PingTimeout)
+
+	for _, file := range strings.Split(c.Files, ",") {
+		conn.Commands = append(conn.Commands, fmt.Sprintf("%s %s regex %s", c.Mode.String(), file, c.Regex))
+	}
+
+	return conn
+}
diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go
new file mode 100644
index 0000000..19246f9
--- /dev/null
+++ b/internal/clients/handlers/basehandler.go
@@ -0,0 +1,134 @@
+package handlers
+
+import (
+	"github.com/mimecast/dtail/internal/logger"
+	"errors"
+	"fmt"
+	"io"
+	"strings"
+	"time"
+)
+
+type baseHandler struct {
+	server       string
+	shellStarted bool
+	commands     chan string
+	pong         chan struct{}
+	receiveBuf   []byte
+	stop         chan struct{}
+	pingTimeout  int
+}
+
+func (h *baseHandler) Server() string {
+	return h.server
+}
+
+// Used to determine whether server is still responding to requests or not.
+func (h *baseHandler) Ping() error {
+	if h.pingTimeout == 0 {
+		// Server ping disabled
+		return nil
+	}
+
+	if err := h.SendCommand("ping"); err != nil {
+		return err
+	}
+
+	select {
+	case <-h.pong:
+		return nil
+	case <-time.After(time.Duration(h.pingTimeout) * time.Second):
+	}
+
+	return errors.New("Didn't receive any server pongs (ping replies)")
+}
+
+func (h *baseHandler) SendCommand(command string) error {
+	if command == "ping" {
+		logger.Trace("Sending command", h.server, command)
+	} else {
+		logger.Debug("Sending command", h.server, command)
+	}
+
+	select {
+	case h.commands <- fmt.Sprintf("%s;", command):
+	case <-time.After(time.Second * 5):
+		return errors.New("Timed out sending command " + command)
+	case <-h.stop:
+	}
+
+	return nil
+}
+
+// Read data from the dtail server via Writer interface.
+func (h *baseHandler) Write(p []byte) (n int, err error) {
+	for _, b := range p {
+		h.receiveBuf = append(h.receiveBuf, b)
+		if b == '\n' {
+			if len(h.receiveBuf) == 0 {
+				continue
+			}
+			message := string(h.receiveBuf)
+			h.handleMessageType(message)
+		}
+	}
+
+	return len(p), nil
+}
+
+// Send data to the dtail server via Reader interface.
+func (h *baseHandler) Read(p []byte) (n int, err error) {
+	select {
+	case command := <-h.commands:
+		n = copy(p, []byte(command))
+	case <-h.stop:
+		return 0, io.EOF
+	}
+	return
+}
+
+// Handle various message types.
+func (h *baseHandler) handleMessageType(message string) {
+	if len(h.receiveBuf) == 0 {
+		return
+	}
+	// Hidden server commands starti with a dot "."
+	if h.receiveBuf[0] == '.' {
+		h.handleHiddenMessage(message)
+		h.receiveBuf = h.receiveBuf[:0]
+		return
+	}
+
+	// Silent mode will only print out remote logs but not remote server
+	// commands. But remote server commands will be still logged to ./log/.
+	if logger.Mode == logger.SilentMode {
+		if h.receiveBuf[0] == 'R' {
+			logger.Raw(message)
+		}
+		h.receiveBuf = h.receiveBuf[:0]
+		return
+	}
+	logger.Raw(message)
+	h.receiveBuf = h.receiveBuf[:0]
+}
+
+// Handle messages received from server which are not meant to be displayed
+// to the end user.
+func (h *baseHandler) handleHiddenMessage(message string) {
+	switch {
+	case strings.HasPrefix(message, ".pong"):
+		h.pong <- struct{}{}
+	case strings.HasPrefix(message, ".syn close connection"):
+		h.SendCommand("ack close connection")
+	}
+}
+
+// Stop the handler.
+func (h *baseHandler) Stop() {
+	select {
+	case <-h.stop:
+	default:
+		logger.Debug("Stopping base handler", h.server)
+		close(h.stop)
+	}
+}
diff --git a/internal/clients/handlers/clienthandler.go b/internal/clients/handlers/clienthandler.go
new file mode 100644
index 0000000..4738cd3
--- /dev/null
+++ b/internal/clients/handlers/clienthandler.go
@@ -0,0 +1,26 @@
+package handlers
+
+import (
+	"github.com/mimecast/dtail/internal/logger"
+)
+
+// ClientHandler is the basic client handler interface.
+type ClientHandler struct {
+	baseHandler
+}
+
+// NewClientHandler creates a new client handler.
+func NewClientHandler(server string, pingTimeout int) *ClientHandler {
+	logger.Debug(server, "Creating new client handler")
+
+	return &ClientHandler{
+		baseHandler{
+			server:       server,
+			shellStarted: false,
+			commands:     make(chan string),
+			pong:         make(chan struct{}, 1),
+			stop:         make(chan struct{}),
+			pingTimeout:  pingTimeout,
+		},
+	}
+}
diff --git a/internal/clients/handlers/handler.go b/internal/clients/handlers/handler.go
new file mode 100644
index 0000000..2013be0
--- /dev/null
+++ b/internal/clients/handlers/handler.go
@@ -0,0 +1,12 @@
+package handlers
+
+import "io"
+
+// Handler provides all methods which can be run on any client handler.
+type Handler interface {
+	io.ReadWriter
+	Ping() error
+	Stop()
+	SendCommand(command string) error
+	Server() string
+}
diff --git a/internal/clients/handlers/healthhandler.go b/internal/clients/handlers/healthhandler.go
new file mode 100644
index 0000000..4051e2c
--- /dev/null
+++ b/internal/clients/handlers/healthhandler.go
@@ -0,0 +1,75 @@
+package handlers
+
+import (
+	"errors"
+	"fmt"
+	"time"
+)
+
+// HealthHandler implements the handler required for health checks.
+type HealthHandler struct {
+	// Buffer of incoming data from server.
+	receiveBuf []byte
+	// To send commands to the server.
+	commands chan string
+	// To receive messages from the server.
+	receive chan<- string
+	// The remote server address
+	server string
+}
+
+// NewHealthHandler returns a new health check handler.
+func NewHealthHandler(server string, receive chan<- string) *HealthHandler {
+	h := HealthHandler{
+		server:   server,
+		receive:  receive,
+		commands: make(chan string),
+	}
+
+	return &h
+}
+
+// Server returns the remote server name.
+func (h *HealthHandler) Server() string {
+	return h.server
+}
+
+// Stop is not of use for health check handler.
+func (h *HealthHandler) Stop() {
+	// Nothing done here.
+}
+
+// Ping is not of use for health check handler.
+func (h *HealthHandler) Ping() error {
+	return nil
+}
+
+// SendCommand send a DTail command to the server.
+func (h *HealthHandler) SendCommand(command string) error {
+	select {
+	case h.commands <- fmt.Sprintf("%s;", command):
+	case <-time.NewTimer(time.Second * 10).C:
+		return errors.New("Timed out sending command " + command)
+	}
+
+	return nil
+}
+
+// Server writes byte stream to client.
+func (h *HealthHandler) Write(p []byte) (n int, err error) {
+	for _, b := range p {
+		h.receiveBuf = append(h.receiveBuf, b)
+		if b == '\n' {
+			h.receive <- string(h.receiveBuf)
+			h.receiveBuf = h.receiveBuf[:0]
+		}
+	}
+
+	return len(p), nil
+}
+
+// Server reads byte stream from client.
+func (h *HealthHandler) Read(p []byte) (n int, err error) {
+	n = copy(p, []byte(<-h.commands))
+	return
+}
diff --git a/internal/clients/handlers/maprhandler.go b/internal/clients/handlers/maprhandler.go
new file mode 100644
index 0000000..d76cdfd
--- /dev/null
+++ b/internal/clients/handlers/maprhandler.go
@@ -0,0 +1,74 @@
+package handlers
+
+import (
+	"github.com/mimecast/dtail/internal/logger"
+	"github.com/mimecast/dtail/internal/mapr"
+	"github.com/mimecast/dtail/internal/mapr/client"
+	"strings"
+)
+
+// MaprHandler is the handler used on the client side for running mapreduce aggregations.
+type MaprHandler struct {
+	baseHandler
+	aggregate *client.Aggregate
+	query     *mapr.Query
+	count     uint64
+}
+
+// NewMaprHandler returns a new mapreduce client handler.
+func NewMaprHandler(server string, query *mapr.Query, globalGroup *mapr.GlobalGroupSet, pingTimeout int) *MaprHandler {
+	return &MaprHandler{
+		baseHandler: baseHandler{
+			server:       server,
+			shellStarted: false,
+			commands:     make(chan string),
+			pong:         make(chan struct{}, 1),
+			stop:         make(chan struct{}),
+			pingTimeout:  pingTimeout,
+		},
+		query:     query,
+		aggregate: client.NewAggregate(server, query, globalGroup),
+	}
+}
+
+// Read data from the dtail server via Writer interface.
+func (h *MaprHandler) Write(p []byte) (n int, err error) {
+	for _, b := range p {
+		h.baseHandler.receiveBuf = append(h.baseHandler.receiveBuf, b)
+		if b == '\n' {
+			if len(h.baseHandler.receiveBuf) == 0 {
+				continue
+			}
+			message := string(h.baseHandler.receiveBuf)
+
+			if h.baseHandler.receiveBuf[0] == 'A' {
+				h.handleAggregateMessage(strings.TrimSpace(message))
+				h.baseHandler.receiveBuf = h.baseHandler.receiveBuf[:0]
+				continue
+			}
+			h.baseHandler.handleMessageType(message)
+		}
+	}
+
+	return len(p), nil
+}
+
+// Handle a message received from server including mapr aggregation
+// related data.
+func (h *MaprHandler) handleAggregateMessage(message string) {
+	h.count++
+	parts := strings.Split(message, "|")
+
+	// Index 0 contains 'AGGREGATE', 1 c