can run scheduled queries across a storage volume

2 files changed, 41 insertions, 17 deletions

diff --git a/internal/server/scheduler.go b/internal/server/scheduler.go
index e0cecfd..d2657d0 100644
--- a/internal/server/scheduler.go
+++ b/internal/server/scheduler.go
@@ -3,7 +3,6 @@ package server
 import (
 	"context"
 	"fmt"
-	"math/rand"
 	"os"
 	"strconv"
 	"strings"
@@ -21,23 +20,17 @@ const authLength = 64
 const authCharset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789!@$%^&*()_+[]"
 
 type scheduler struct {
-	authPayload string
 }
 
 func newScheduler() *scheduler {
-	seededRand := rand.New(rand.NewSource(time.Now().UnixNano()))
-
-	b := make([]byte, authLength)
-	for i := range b {
-		b[i] = authCharset[seededRand.Intn(len(authCharset))]
-	}
-
-	return &scheduler{
-		authPayload: string(b),
-	}
+	return &scheduler{}
 }
 
 func (s *scheduler) start(ctx context.Context) {
+	// First run after just 10s!
+	time.Sleep(time.Second * 10)
+	s.runJobs(ctx)
+
 	for {
 		select {
 		case <-time.After(time.Minute):
@@ -75,7 +68,7 @@ func (s *scheduler) runJobs(ctx context.Context) {
 			continue
 		}
 
-		servers := scheduled.Servers
+		servers := strings.Join(scheduled.Servers, ",")
 		if servers == "" {
 			servers = config.Server.SSHBindAddress
 		}
@@ -88,7 +81,8 @@ func (s *scheduler) runJobs(ctx context.Context) {
 			Mode:              omode.MapClient,
 			UserName:          config.ScheduledUser,
 		}
-		args.SSHAuthMethods = append(args.SSHAuthMethods, gossh.Password(s.authPayload))
+
+		args.SSHAuthMethods = append(args.SSHAuthMethods, gossh.Password(scheduled.Name))
 
 		tmpOutfile := fmt.Sprintf("%s.tmp", outfile)
 		query := fmt.Sprintf("%s outfile %s", scheduled.Query, tmpOutfile)
diff --git a/internal/server/server.go b/internal/server/server.go
index 5c46147..34d0d88 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"net"
+	"strings"
 
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/logger"
@@ -198,16 +199,45 @@ func (s *Server) handleRequests(ctx context.Context, sshConn gossh.Conn, in <-ch
 
 func (s *Server) backgroundUserCallback(c gossh.ConnMetadata, authPayload []byte) (*gossh.Permissions, error) {
 	user := user.New(c.User(), c.RemoteAddr().String())
+	authInfo := string(authPayload)
 
-	if user.Name == config.ControlUser && string(authPayload) == config.ControlUser {
+	if user.Name == config.ControlUser && authInfo == config.ControlUser {
 		logger.Debug(user, "Granting permissions to control user")
 		return nil, nil
 	}
 
-	if user.Name == config.ScheduledUser && string(authPayload) == s.sched.authPayload {
-		logger.Debug(user, "Granting permissions to schedule user")
+	if user.Name == config.ScheduledUser && s.canRunScheduledJob(c.RemoteAddr().String(), user, authInfo) {
+		logger.Debug(user, "Schedule user can run scheduled job remotely")
 		return nil, nil
 	}
 
 	return nil, fmt.Errorf("user %s not authorized", user)
 }
+
+func (s *Server) canRunScheduledJob(addr string, user *user.User, jobName string) bool {
+	logger.Debug("canRunScheduledJob", user, jobName)
+	splitted := strings.Split(addr, ":")
+	ip := splitted[0]
+
+	for _, job := range config.Server.Schedule {
+		if job.Name != jobName {
+			continue
+		}
+		for _, myAddr := range job.AllowFrom {
+			myIps, err := net.LookupIP(myAddr)
+			if err != nil {
+				logger.Error(user, myAddr, err)
+				continue
+			}
+
+			for _, myIp := range myIps {
+				logger.Debug("canRunScheduledJob", "Comparing IP addresses", ip, myIp.String())
+				if ip == myIp.String() {
+					return true
+				}
+			}
+		}
+	}
+
+	return false
+}