From 0945da8dfefcbb723eecea0e5f4eafff63398253 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Paul=20B=C3=BCtow?= Date: Sun, 26 Jan 2020 11:26:53 +0000 Subject: Introduce drun command, refactor code to use context package --- cmd/dcat/main.go | 19 +++++++------ cmd/dexec/main.go | 80 --------------------------------------------------- cmd/dgrep/main.go | 18 ++++++------ cmd/dmap/main.go | 19 +++++++------ cmd/drun/main.go | 82 +++++++++++++++++++++++++++++++++++++++++++++++++++++ cmd/dserver/main.go | 22 ++++++++------ cmd/dtail/main.go | 19 +++++++------ 7 files changed, 137 insertions(+), 122 deletions(-) delete mode 100644 cmd/dexec/main.go create mode 100644 cmd/drun/main.go (limited to 'cmd') diff --git a/cmd/dcat/main.go b/cmd/dcat/main.go index b02d369..1ec945d 100644 --- a/cmd/dcat/main.go +++ b/cmd/dcat/main.go @@ -1,12 +1,14 @@ package main import ( + "context" "flag" + "os" "github.com/mimecast/dtail/internal/clients" "github.com/mimecast/dtail/internal/color" "github.com/mimecast/dtail/internal/config" - "github.com/mimecast/dtail/internal/logger" + "github.com/mimecast/dtail/internal/io/logger" "github.com/mimecast/dtail/internal/pprof" "github.com/mimecast/dtail/internal/user" "github.com/mimecast/dtail/internal/version" @@ -27,7 +29,6 @@ func main() { var sshPort int var trustAllHosts bool - pingTimeoutS := 60 userName := user.Name() flag.BoolVar(&debugEnable, "debug", false, "Activate debug messages") @@ -37,7 +38,6 @@ func main() { flag.BoolVar(&silentEnable, "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(&pingTimeoutS, "pingTimeout", 10, "The server ping timeout (0 means disable pings)") flag.IntVar(&sshPort, "port", 2222, "SSH server port") flag.StringVar(&cfgFile, "cfg", "", "Config file path") flag.StringVar(&discovery, "discovery", "", "Server discovery method") @@ -54,9 +54,10 @@ func main() { version.PrintAndExit() } + ctx := context.Background() serverEnable := false - logger.Start(serverEnable, debugEnable, silentEnable, silentEnable) - defer logger.Stop() + + logger.Start(ctx, serverEnable, debugEnable, silentEnable, silentEnable) if pprofEnable || config.Common.PProfEnable { pprof.Start() @@ -67,14 +68,16 @@ func main() { ServersStr: serversStr, Discovery: discovery, UserName: userName, - Files: files, + What: files, TrustAllHosts: trustAllHosts, - PingTimeout: pingTimeoutS, } client, err := clients.NewCatClient(args) if err != nil { panic(err) } - client.Start() + + status := client.Start(ctx) + logger.Flush() + os.Exit(status) } diff --git a/cmd/dexec/main.go b/cmd/dexec/main.go deleted file mode 100644 index 7a7ab1f..0000000 --- a/cmd/dexec/main.go +++ /dev/null @@ -1,80 +0,0 @@ -package main - -import ( - "flag" - - "github.com/mimecast/dtail/internal/clients" - "github.com/mimecast/dtail/internal/color" - "github.com/mimecast/dtail/internal/config" - "github.com/mimecast/dtail/internal/logger" - "github.com/mimecast/dtail/internal/pprof" - "github.com/mimecast/dtail/internal/user" - "github.com/mimecast/dtail/internal/version" -) - -// The evil begins here. -func main() { - var cfgFile string - var connectionsPerCPU int - var debugEnable bool - var discovery string - var displayVersion bool - var command string - var noColor bool - var pprofEnable bool - var serversStr string - var silentEnable bool - var sshPort int - var trustAllHosts bool - - pingTimeoutS := 60 - userName := user.Name() - - 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(&silentEnable, "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(&pingTimeoutS, "pingTimeout", 10, "The server ping timeout (0 means disable pings)") - flag.IntVar(&sshPort, "port", 2222, "SSH server port") - flag.StringVar(&cfgFile, "cfg", "", "Config file path") - flag.StringVar(&discovery, "discovery", "", "Server discovery method") - flag.StringVar(&command, "command", "", "Command to run") - flag.StringVar(&serversStr, "servers", "", "Remote servers to connect") - flag.StringVar(&userName, "user", userName, "Your system user name") - - flag.Parse() - - config.Read(cfgFile, sshPort) - color.Colored = !noColor - - if displayVersion { - version.PrintAndExit() - } - - serverEnable := false - logger.Start(serverEnable, debugEnable, silentEnable, silentEnable) - defer logger.Stop() - - if pprofEnable || config.Common.PProfEnable { - pprof.Start() - } - - args := clients.Args{ - ConnectionsPerCPU: connectionsPerCPU, - ServersStr: serversStr, - Discovery: discovery, - UserName: userName, - Files: files, - TrustAllHosts: trustAllHosts, - PingTimeout: pingTimeoutS, - } - - client, err := clients.NewExecClient(args) - if err != nil { - panic(err) - } - client.Start() -} diff --git a/cmd/dgrep/main.go b/cmd/dgrep/main.go index d1a7d52..74a501f 100644 --- a/cmd/dgrep/main.go +++ b/cmd/dgrep/main.go @@ -1,12 +1,14 @@ package main import ( + "context" "flag" + "os" "github.com/mimecast/dtail/internal/clients" "github.com/mimecast/dtail/internal/color" "github.com/mimecast/dtail/internal/config" - "github.com/mimecast/dtail/internal/logger" + "github.com/mimecast/dtail/internal/io/logger" "github.com/mimecast/dtail/internal/pprof" "github.com/mimecast/dtail/internal/user" "github.com/mimecast/dtail/internal/version" @@ -28,7 +30,6 @@ func main() { var sshPort int var trustAllHosts bool - pingTimeoutS := 60 userName := user.Name() flag.BoolVar(&debugEnable, "debug", false, "Activate debug messages") @@ -38,7 +39,6 @@ func main() { flag.BoolVar(&silentEnable, "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(&pingTimeoutS, "pingTimeout", 10, "The server ping timeout (0 means disable pings)") flag.IntVar(&sshPort, "port", 2222, "SSH server port") flag.StringVar(&cfgFile, "cfg", "", "Config file path") flag.StringVar(&discovery, "discovery", "", "Server discovery method") @@ -56,9 +56,9 @@ func main() { version.PrintAndExit() } + ctx := context.Background() serverEnable := false - logger.Start(serverEnable, debugEnable, silentEnable, silentEnable) - defer logger.Stop() + logger.Start(ctx, serverEnable, debugEnable, silentEnable, silentEnable) if pprofEnable || config.Common.PProfEnable { pprof.Start() @@ -69,9 +69,8 @@ func main() { ServersStr: serversStr, Discovery: discovery, UserName: userName, - Files: files, + What: files, TrustAllHosts: trustAllHosts, - PingTimeout: pingTimeoutS, Regex: regex, } @@ -79,5 +78,8 @@ func main() { if err != nil { panic(err) } - client.Start() + + status := client.Start(ctx) + logger.Flush() + os.Exit(status) } diff --git a/cmd/dmap/main.go b/cmd/dmap/main.go index 83dad50..f3f706a 100644 --- a/cmd/dmap/main.go +++ b/cmd/dmap/main.go @@ -1,13 +1,15 @@ package main import ( + "context" "flag" + "os" - "github.com/mimecast/dtail/internal/omode" "github.com/mimecast/dtail/internal/clients" "github.com/mimecast/dtail/internal/color" "github.com/mimecast/dtail/internal/config" - "github.com/mimecast/dtail/internal/logger" + "github.com/mimecast/dtail/internal/io/logger" + "github.com/mimecast/dtail/internal/omode" "github.com/mimecast/dtail/internal/pprof" "github.com/mimecast/dtail/internal/user" "github.com/mimecast/dtail/internal/version" @@ -29,7 +31,6 @@ func main() { var sshPort int var trustAllHosts bool - pingTimeoutS := 900 userName := user.Name() flag.BoolVar(&debugEnable, "debug", false, "Activate debug messages") @@ -39,7 +40,6 @@ func main() { flag.BoolVar(&silentEnable, "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(&pingTimeoutS, "pingTimeout", 10, "The server ping timeout (0 means disable pings)") flag.IntVar(&sshPort, "port", 2222, "SSH server port") flag.StringVar(&cfgFile, "cfg", "", "Config file path") flag.StringVar(&discovery, "discovery", "", "Server discovery method") @@ -57,10 +57,10 @@ func main() { version.PrintAndExit() } + ctx := context.Background() serverEnable := false - logger.Start(serverEnable, debugEnable, silentEnable, silentEnable) - defer logger.Stop() + logger.Start(ctx, serverEnable, debugEnable, silentEnable, silentEnable) if pprofEnable || config.Common.PProfEnable { pprof.Start() } @@ -70,9 +70,8 @@ func main() { ServersStr: serversStr, Discovery: discovery, UserName: userName, - Files: files, + What: files, TrustAllHosts: trustAllHosts, - PingTimeout: pingTimeoutS, Mode: omode.MapClient, } @@ -81,5 +80,7 @@ func main() { panic(err) } - client.Start() + status := client.Start(ctx) + logger.Flush() + os.Exit(status) } diff --git a/cmd/drun/main.go b/cmd/drun/main.go new file mode 100644 index 0000000..b1936d4 --- /dev/null +++ b/cmd/drun/main.go @@ -0,0 +1,82 @@ +package main + +import ( + "context" + "flag" + "os" + + "github.com/mimecast/dtail/internal/clients" + "github.com/mimecast/dtail/internal/color" + "github.com/mimecast/dtail/internal/config" + "github.com/mimecast/dtail/internal/io/logger" + "github.com/mimecast/dtail/internal/pprof" + "github.com/mimecast/dtail/internal/user" + "github.com/mimecast/dtail/internal/version" +) + +// The evil begins here. +func main() { + var cfgFile string + var command string + var connectionsPerCPU int + var debugEnable bool + var discovery string + var displayVersion bool + var noColor bool + var pprofEnable bool + var serversStr string + var silentEnable bool + var sshPort int + var trustAllHosts bool + + userName := user.Name() + + 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(&silentEnable, "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(&sshPort, "port", 2222, "SSH server port") + flag.StringVar(&cfgFile, "cfg", "", "Config file path") + flag.StringVar(&command, "command", "", "Command to run") + flag.StringVar(&discovery, "discovery", "", "Server discovery method") + flag.StringVar(&serversStr, "servers", "", "Remote servers to connect") + flag.StringVar(&userName, "user", userName, "Your system user name") + + flag.Parse() + + config.Read(cfgFile, sshPort) + color.Colored = !noColor + + if displayVersion { + version.PrintAndExit() + } + + ctx := context.Background() + serverEnable := false + + logger.Start(ctx, serverEnable, debugEnable, silentEnable, silentEnable) + if pprofEnable || config.Common.PProfEnable { + pprof.Start() + } + + args := clients.Args{ + ConnectionsPerCPU: connectionsPerCPU, + ServersStr: serversStr, + Discovery: discovery, + UserName: userName, + What: command, + TrustAllHosts: trustAllHosts, + } + + client, err := clients.NewRunClient(args) + if err != nil { + panic(err) + } + + status := client.Start(ctx) + logger.Flush() + os.Exit(status) +} diff --git a/cmd/dserver/main.go b/cmd/dserver/main.go index 489910b..aa209a8 100644 --- a/cmd/dserver/main.go +++ b/cmd/dserver/main.go @@ -1,13 +1,14 @@ package main import ( + "context" "flag" "os" "time" "github.com/mimecast/dtail/internal/color" "github.com/mimecast/dtail/internal/config" - "github.com/mimecast/dtail/internal/logger" + "github.com/mimecast/dtail/internal/io/logger" "github.com/mimecast/dtail/internal/pprof" "github.com/mimecast/dtail/internal/server" "github.com/mimecast/dtail/internal/user" @@ -24,7 +25,7 @@ func main() { var shutdownAfter int var sshPort int - userName := user.Name() + user.NoRootCheck() flag.BoolVar(&debugEnable, "debug", false, "Activate debug messages") flag.BoolVar(&displayVersion, "version", false, "Display version") @@ -43,19 +44,23 @@ func main() { version.PrintAndExit() } + ctx := context.Background() + serverEnable := true silentEnable := false nothingEnable := false - logger.Start(serverEnable, debugEnable, silentEnable, nothingEnable) - defer logger.Stop() + logger.Start(ctx, serverEnable, debugEnable, silentEnable, nothingEnable) 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") + select { + case <-time.After(time.Duration(shutdownAfter) * time.Second): + logger.Info("Auto shutdown timer reached, shutting down now") + case <-ctx.Done(): + } }() } @@ -63,7 +68,8 @@ func main() { pprof.Start() } - logger.Info("Launching server", version.String(), userName) sshServer := server.New() - sshServer.Start() + status := sshServer.Start(ctx) + logger.Flush() + os.Exit(status) } diff --git a/cmd/dtail/main.go b/cmd/dtail/main.go index 1bf77c7..76070ff 100644 --- a/cmd/dtail/main.go +++ b/cmd/dtail/main.go @@ -1,13 +1,14 @@ package main import ( + "context" "flag" "os" "github.com/mimecast/dtail/internal/clients" "github.com/mimecast/dtail/internal/color" "github.com/mimecast/dtail/internal/config" - "github.com/mimecast/dtail/internal/logger" + "github.com/mimecast/dtail/internal/io/logger" "github.com/mimecast/dtail/internal/omode" "github.com/mimecast/dtail/internal/pprof" "github.com/mimecast/dtail/internal/user" @@ -32,7 +33,6 @@ func main() { var sshPort int var trustAllHosts bool - pingTimeoutS := 5 userName := user.Name() flag.BoolVar(&checkHealth, "checkHealth", false, "Only check for server health") @@ -43,7 +43,6 @@ func main() { flag.BoolVar(&silentEnable, "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(&pingTimeoutS, "pingTimeout", 10, "The server ping timeout (0 means disable pings)") flag.IntVar(&sshPort, "port", 2222, "SSH server port") flag.StringVar(&cfgFile, "cfg", "", "Config file path") flag.StringVar(&discovery, "discovery", "", "Server discovery method") @@ -62,17 +61,18 @@ func main() { version.PrintAndExit() } + ctx := context.Background() + if checkHealth { healthClient, _ := clients.NewHealthClient(omode.HealthClient) - os.Exit(healthClient.Start()) + os.Exit(healthClient.Start(ctx)) } serverEnable := false if checkHealth { silentEnable = true } - logger.Start(serverEnable, debugEnable, silentEnable, silentEnable) - defer logger.Stop() + logger.Start(ctx, serverEnable, debugEnable, silentEnable, silentEnable) if pprofEnable || config.Common.PProfEnable { pprof.Start() @@ -83,9 +83,8 @@ func main() { ServersStr: serversStr, Discovery: discovery, UserName: userName, - Files: files, + What: files, TrustAllHosts: trustAllHosts, - PingTimeout: pingTimeoutS, Regex: regex, Mode: omode.TailClient, } @@ -104,5 +103,7 @@ func main() { } } - client.Start() + status := client.Start(ctx) + logger.Flush() + os.Exit(status) } -- cgit v1.2.3