Release of DTail v1.0.0v1.0.0

1 files changed, 250 insertions, 0 deletions

diff --git a/main.go b/main.go
new file mode 100644
index 0000000..e54d320
--- /dev/null
+++ b/main.go
@@ -0,0 +1,250 @@
+package main
+
+import (
+	"dtail/clients"
+	"dtail/color"
+	"dtail/config"
+	"dtail/logger"
+	"dtail/omode"
+	"dtail/server"
+	"dtail/version"
+	"flag"
+	"fmt"
+	"net/http"
+	_ "net/http"
+	_ "net/http/pprof"
+	"os"
+	"os/user"
+	"runtime"
+	"sync"
+	"time"
+)
+
+// The evil begins here.
+func main() {
+	var cfgFile, modeStr string
+	var checkHealth bool
+	var clientServerEnable bool
+	var connectionsPerCPU int
+	var debugEnable bool
+	var discovery string
+	var displayVersion bool
+	var files string
+	var grep, regex string
+	var maxInitConnections int
+	var noColor bool
+	var pprofEnable bool
+	var queryStr string
+	var serversStr string
+	var shutdownAfter int
+	var silent bool
+	var sshPort int
+	var trustAllHosts bool
+	var userName string
+
+	user, err := user.Current()
+	if err != nil {
+		panic(err)
+	}
+
+	if user.Uid == "0" {
+		panic("Not allowed to run as UID 0")
+	}
+
+	if user.Gid == "0" {
+		panic("Not allowed to run as GID 0")
+	}
+
+	defaultMode := omode.Default()
+	serverEnable := defaultMode == omode.Server
+	clientEnable := !serverEnable
+
+	// Based on the mode we have different default timeouts
+	var pingTimeoutS int
+	switch defaultMode {
+	case omode.CatClient:
+		fallthrough
+	case omode.GrepClient:
+		pingTimeoutS = 60
+	case omode.MapClient:
+		pingTimeoutS = 900
+	default:
+		pingTimeoutS = 5
+	}
+
+	flag.BoolVar(&checkHealth, "checkHealth", false, "Only check for server health")
+	flag.BoolVar(&clientServerEnable, "clientServer", false, "Enable client and server (dev purposes)")
+	flag.BoolVar(&debugEnable, "debug", false, "Activate debug messages")
+	flag.BoolVar(&displayVersion, "version", false, "Display version")
+	flag.BoolVar(&noColor, "noColor", false, "Disable ANSII terminal colors")
+	flag.BoolVar(&pprofEnable, "pprofEnable", false, "Enable pprof server")
+	flag.BoolVar(&serverEnable, "server", serverEnable, "Start as a DTail server")
+	flag.BoolVar(&silent, "silent", false, "Reduce output")
+	flag.BoolVar(&trustAllHosts, "trustAllHosts", false, "Auto trust all unknown host keys")
+	flag.IntVar(&connectionsPerCPU, "cpc", 10, "How many connections established per CPU core concurrently")
+	flag.IntVar(&maxInitConnections, "mic", 20, "Max cpc")
+	flag.IntVar(&shutdownAfter, "shutdownAfter", 0, "Automatically shutdown after so many seconds")
+	flag.IntVar(&sshPort, "port", 2222, "SSH server port")
+	flag.IntVar(&pingTimeoutS, "pingTimeout", 10, "The server ping timeout (0 means disable pings)")
+	flag.StringVar(&cfgFile, "cfg", "", "Config file path")
+	flag.StringVar(&discovery, "discovery", "", "Server discovery method")
+	flag.StringVar(&files, "files", "", "File(s) to read")
+	flag.StringVar(&grep, "grep", "", "Regular expression (deprecated)")
+	flag.StringVar(&modeStr, "mode", defaultMode.String(), "Operating mode (tail, grep, cat, map, server)")
+	flag.StringVar(&queryStr, "query", "", "Map reduce query")
+	flag.StringVar(&regex, "regex", "", "Regular expression")
+	flag.StringVar(&serversStr, "servers", "", "Remote servers to connect")
+	flag.StringVar(&userName, "user", user.Username, "Your system user name")
+
+	mode := omode.New(modeStr)
+
+	flag.Parse()
+
+	config.Init(cfgFile)
+	color.Init(!noColor)
+
+	if displayVersion {
+		fmt.Println(version.PaintedString())
+		os.Exit(0)
+	}
+
+	// Figure out how many SSH sessions can be established concurrently.
+	if connectionsPerCPU*runtime.NumCPU() < maxInitConnections {
+		maxInitConnections = connectionsPerCPU * runtime.NumCPU()
+	}
+
+	// Figure out in which mode I am? Server or client or both (the latter for dev purposes)?
+	if serverEnable {
+		clientEnable = false
+	}
+	if clientServerEnable {
+		clientEnable = true
+		serverEnable = true
+	}
+
+	// If non-standard port specified, overwrite config
+	if sshPort != 2222 {
+		config.Common.SSHPort = sshPort
+	}
+
+	// Figure out the log level.
+	var logMode logger.LogMode
+	switch {
+	case debugEnable:
+		logMode = logger.DebugMode
+	case checkHealth:
+		logMode = logger.NothingMode
+	case config.Common.TraceEnable:
+		logMode = logger.TraceMode
+	case config.Common.DebugEnable:
+		logMode = logger.DebugMode
+	case silent:
+		logMode = logger.SilentMode
+	default:
+		logMode = logger.NormalMode
+	}
+
+	// Figure out the log strategy.
+	var logStrategy logger.LogStrategy
+	switch config.Common.LogStrategy {
+	case "daily":
+		logStrategy = logger.DailyStrategy
+	case "stdout":
+		fallthrough
+	default:
+		logStrategy = logger.StdoutStrategy
+	}
+
+	logger.Init(serverEnable, logMode, logStrategy)
+
+	// Wait group for shutting down logger.
+	var wg sync.WaitGroup
+	if serverEnable {
+		wg.Add(1)
+	}
+	if clientEnable {
+		wg.Add(1)
+	}
+
+	logger.Debug("Common config", config.Common)
+	logger.Debug("Client config", config.Client)
+	logger.Debug("Server config", config.Server)
+
+	if grep != "" {
+		logger.Warn("Flag 'grep' is deprecated and may be removed in the future, please use 'regex' instead")
+		if regex == "" {
+			regex = grep
+		}
+	}
+
+	if checkHealth {
+		healthClient, _ := clients.NewHealthClient(omode.HealthClient)
+		os.Exit(healthClient.Start(&wg))
+	}
+
+	if shutdownAfter > 0 {
+		go func() {
+			defer os.Exit(1)
+
+			logger.Info("Enabling auto shutdown timer", shutdownAfter)
+			time.Sleep(time.Duration(shutdownAfter) * time.Second)
+			logger.Info("Auto shutdown timer reached, shutting down now")
+		}()
+	}
+
+	if pprofEnable || config.Common.PProfEnable {
+		bindAddr := fmt.Sprintf("%s:%d", config.Common.PProfBindAddress, config.Common.PProfPort)
+		logger.Info("Starting PProf server", bindAddr)
+		go http.ListenAndServe(bindAddr, nil)
+	}
+
+	if serverEnable {
+		logger.Info("Launching server", mode, version.String())
+		sshServer := server.New()
+		go sshServer.Start(&wg)
+	}
+
+	if clientEnable {
+		var client clients.Client
+		var err error
+
+		logger.Info("Launching client", mode, version.String())
+
+		args := clients.Args{
+			Mode:               mode,
+			ServersStr:         serversStr,
+			Discovery:          discovery,
+			UserName:           userName,
+			Files:              files,
+			Regex:              regex,
+			TrustAllHosts:      trustAllHosts,
+			MaxInitConnections: maxInitConnections,
+			PingTimeout:        pingTimeoutS,
+		}
+
+		switch mode {
+		case omode.TailClient:
+			switch queryStr {
+			case "":
+				client, err = clients.NewTailClient(args)
+			default:
+				client, err = clients.NewMaprClient(args, queryStr)
+			}
+		case omode.GrepClient:
+			client, err = clients.NewGrepClient(args)
+		case omode.CatClient:
+			client, err = clients.NewCatClient(args)
+		case omode.MapClient:
+			client, err = clients.NewMaprClient(args, queryStr)
+		}
+
+		if err != nil {
+			panic(err)
+		}
+
+		go client.Start(&wg)
+	}
+
+	wg.Wait()
+	logger.Stop()
+}