initial scheduled jobs

8 files changed, 56 insertions, 19 deletions

diff --git a/internal/clients/args.go b/internal/clients/args.go
index dea5a9e..f55ce90 100644
--- a/internal/clients/args.go
+++ b/internal/clients/args.go
@@ -2,16 +2,20 @@ package clients
 
 import (
 	"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
-	Regex             string
-	TrustAllHosts     bool
-	Discovery         string
-	ConnectionsPerCPU int
+	Mode               omode.Mode
+	ServersStr         string
+	UserName           string
+	What               string
+	Regex              string
+	TrustAllHosts      bool
+	Discovery          string
+	ConnectionsPerCPU  int
+	SSHAuthMethods     []gossh.AuthMethod
+	SSHHostKeyCallback gossh.HostKeyCallback
 }
diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go
index b1540ea..725910a 100644
--- a/internal/clients/baseclient.go
+++ b/internal/clients/baseclient.go
@@ -40,7 +40,8 @@ func (c *baseClient) init(maker maker) {
 	logger.Info("Initiating base client")
 
 	c.maker = maker
-	c.sshAuthMethods, c.hostKeyCallback = client.InitSSHAuthMethods(c.TrustAllHosts, c.throttleCh)
+	c.sshAuthMethods, c.hostKeyCallback = client.InitSSHAuthMethods(c.Args, c.throttleCh)
+
 	discoveryService := discovery.New(c.Discovery, c.ServersStr, discovery.Shuffle)
 
 	for _, server := range discoveryService.ServerList() {
diff --git a/internal/config/config.go b/internal/config/config.go
index 0f26635..d20a480 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -9,6 +9,9 @@ import (
 // ControlUser is used for various DTail specific operations.
 const ControlUser string = "DTAIL-CONTROL-USER"
 
+// ScheduledUser is used for scheduled queries.
+const ScheduledUser string = "DTAIL-SCHEDULED-USER"
+
 // Client holds a DTail client configuration.
 var Client *ClientConfig
 
diff --git a/internal/config/server.go b/internal/config/server.go
index 7883b33..be6d45e 100644
--- a/internal/config/server.go
+++ b/internal/config/server.go
@@ -13,6 +13,16 @@ type Permissions struct {
 	Users map[string][]string
 }
 
+// Scheduled allows to configure scheduled mapreduce jobs.
+type Scheduled struct {
+	Name              string
+	ConnectionsPerCPU int
+	Discovery         string
+	Files             string
+	Query             string
+	Servers           string
+}
+
 // ServerConfig represents the server configuration.
 type ServerConfig struct {
 	// The SSH server bind port.
@@ -31,6 +41,8 @@ type ServerConfig struct {
 	HostKeyFile string
 	// The host key size in bits
 	HostKeyBits int
+	// Scheduled mapreduce jobs.
+	Schedule []Scheduled `json:",omitempty"`
 }
 
 // Create a new default server configuration.
diff --git a/internal/server/server.go b/internal/server/server.go
index 42eb74c..5c46147 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -27,6 +27,8 @@ type Server struct {
 	catLimiterCh chan struct{}
 	// To control the max amount of concurrent tails
 	tailLimiterCh chan struct{}
+	// To run scheduled tasks (if configured)
+	sched *scheduler
 }
 
 // New returns a new server.
@@ -37,9 +39,10 @@ func New() *Server {
 		sshServerConfig: &gossh.ServerConfig{},
 		catLimiterCh:    make(chan struct{}, config.Server.MaxConcurrentCats),
 		tailLimiterCh:   make(chan struct{}, config.Server.MaxConcurrentTails),
+		sched:           newScheduler(),
 	}
 
-	s.sshServerConfig.PasswordCallback = s.controlUserCallback
+	s.sshServerConfig.PasswordCallback = s.backgroundUserCallback
 	s.sshServerConfig.PublicKeyCallback = server.PublicKeyCallback
 
 	private, err := gossh.ParsePrivateKey(server.PrivateHostKey())
@@ -62,7 +65,8 @@ func (s *Server) Start(ctx context.Context) int {
 		logger.FatalExit("Failed to open listening TCP socket", err)
 	}
 
-	go s.stats.periodicLogServerStats(ctx)
+	go s.stats.start(ctx)
+	go s.sched.start(ctx)
 
 	for {
 		conn, err := listener.Accept() // Blocking
@@ -192,13 +196,18 @@ func (s *Server) handleRequests(ctx context.Context, sshConn gossh.Conn, in <-ch
 	return nil
 }
 
-func (*Server) controlUserCallback(c gossh.ConnMetadata, authPayload []byte) (*gossh.Permissions, error) {
+func (s *Server) backgroundUserCallback(c gossh.ConnMetadata, authPayload []byte) (*gossh.Permissions, error) {
 	user := user.New(c.User(), c.RemoteAddr().String())
 
 	if user.Name == config.ControlUser && string(authPayload) == config.ControlUser {
-		logger.Debug(user, "Initiating master control program")
+		logger.Debug(user, "Granting permissions to control user")
 		return nil, nil
 	}
 
-	return nil, fmt.Errorf("Not authorized")
+	if user.Name == config.ScheduledUser && string(authPayload) == s.sched.authPayload {
+		logger.Debug(user, "Granting permissions to schedule user")
+		return nil, nil
+	}
+
+	return nil, fmt.Errorf("user %s not authorized", user)
 }
diff --git a/internal/server/stats.go b/internal/server/stats.go
index 4d661f7..ac579ad 100644
--- a/internal/server/stats.go
+++ b/internal/server/stats.go
@@ -67,7 +67,7 @@ func (s *stats) serverLimitExceeded() error {
 	return nil
 }
 
-func (s *stats) periodicLogServerStats(ctx context.Context) {
+func (s *stats) start(ctx context.Context) {
 	for {
 		select {
 		case <-time.NewTimer(time.Second * 10).C:
diff --git a/internal/ssh/client/authmethods.go b/internal/ssh/client/authmethods.go
index 967866f..072fad0 100644
--- a/internal/ssh/client/authmethods.go
+++ b/internal/ssh/client/authmethods.go
@@ -1,18 +1,26 @@
 package client
 
 import (
+	"os"
+
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/logger"
 	"github.com/mimecast/dtail/internal/ssh"
-	"os"
 
 	gossh "golang.org/x/crypto/ssh"
 )
 
 // InitSSHAuthMethods initialises all known SSH auth methods on othe client side.
-func InitSSHAuthMethods(trustAllHosts bool, throttleCh chan struct{}) ([]gossh.AuthMethod, *HostKeyCallback) {
-	var sshAuthMethods []gossh.AuthMethod
+func InitSSHAuthMethods(args clients.Args, trustAllHosts bool, throttleCh chan struct{}) ([]gossh.AuthMethod, *HostKeyCallback) {
+	if len(args.SSHAuthMethods) > 0 {
+		hostKeyCallback, err := NewSimpleCallback(trustAllHosts)
+		if err != nil {
+			logger.FatalExit(err)
+		}
+		return args.SSHAuthMethods, hostKeyCallback
+	}
 
+	var sshAuthMethods []gossh.AuthMethod
 	if config.Common.ExperimentalFeaturesEnable {
 		sshAuthMethods = append(sshAuthMethods, gossh.Password("experimental feature test"))
 		logger.Info("Added experimental method to list of auth methods")
diff --git a/internal/version/version.go b/internal/version/version.go
index 3fd729b..82fd86f 100644
--- a/internal/version/version.go
+++ b/internal/version/version.go
@@ -11,7 +11,7 @@ const (
 	// Name of DTail.
 	Name string = "DTail"
 	// Version of DTail.
-	Version string = "2.0.3"
+	Version string = "2.1.0"
 	// Additional information for DTail
 	Additional string = "develop"
 	// ProtocolCompat -ibility version.