keyhunter/pkg/scheduler/scheduler.go

package scheduler

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/go-co-op/gocron/v2"
	"github.com/salvacybersec/keyhunter/pkg/storage"
)

// Config holds the dependencies for a Scheduler.
type Config struct {
	// DB is the storage backend for persisting jobs and subscribers.
	DB *storage.DB
	// ScanFunc executes a scan command and returns the finding count.
	ScanFunc func(ctx context.Context, scanCommand string) (int, error)
	// OnComplete is called after a job finishes. May be nil.
	OnComplete func(result JobResult)
}

// Scheduler wraps gocron with SQLite persistence for scheduled scan jobs.
type Scheduler struct {
	cfg   Config
	sched gocron.Scheduler
	jobs  map[string]gocron.Job
	mu    sync.Mutex
}

// New creates a new Scheduler with the given configuration.
func New(cfg Config) (*Scheduler, error) {
	s, err := gocron.NewScheduler()
	if err != nil {
		return nil, fmt.Errorf("creating gocron scheduler: %w", err)
	}
	return &Scheduler{
		cfg:   cfg,
		sched: s,
		jobs:  make(map[string]gocron.Job),
	}, nil
}

// Start loads all enabled jobs from the database and begins scheduling.
func (s *Scheduler) Start(ctx context.Context) error {
	jobs, err := s.cfg.DB.ListScheduledJobs()
	if err != nil {
		return fmt.Errorf("loading scheduled jobs: %w", err)
	}
	for _, sj := range jobs {
		if !sj.Enabled {
			continue
		}
		if err := s.registerJob(ctx, sj); err != nil {
			return fmt.Errorf("registering job %q: %w", sj.Name, err)
		}
	}
	s.sched.Start()
	return nil
}

// Stop shuts down the gocron scheduler.
func (s *Scheduler) Stop() error {
	return s.sched.Shutdown()
}

// AddJob creates a new scheduled job, persists it, and registers it with gocron.
func (s *Scheduler) AddJob(name, cronExpr, scanCommand string, notifyTelegram bool) error {
	sj := storage.ScheduledJob{
		Name:           name,
		CronExpr:       cronExpr,
		ScanCommand:    scanCommand,
		NotifyTelegram: notifyTelegram,
		Enabled:        true,
	}
	if _, err := s.cfg.DB.SaveScheduledJob(sj); err != nil {
		return fmt.Errorf("saving job %q: %w", name, err)
	}
	return s.registerJob(context.Background(), sj)
}

// RemoveJob removes a job from gocron and deletes it from the database.
func (s *Scheduler) RemoveJob(name string) error {
	s.mu.Lock()
	j, ok := s.jobs[name]
	if ok {
		delete(s.jobs, name)
	}
	s.mu.Unlock()

	if ok {
		if err := s.sched.RemoveJob(j.ID()); err != nil {
			return fmt.Errorf("removing gocron job %q: %w", name, err)
		}
	}

	if _, err := s.cfg.DB.DeleteScheduledJob(name); err != nil {
		return fmt.Errorf("deleting job %q from DB: %w", name, err)
	}
	return nil
}

// ListJobs returns all scheduled jobs from the database.
func (s *Scheduler) ListJobs() ([]storage.ScheduledJob, error) {
	return s.cfg.DB.ListScheduledJobs()
}

// RunJob manually triggers a job by name. Looks up the job in the DB,
// runs ScanFunc, updates last_run, and calls OnComplete.
func (s *Scheduler) RunJob(ctx context.Context, name string) (JobResult, error) {
	sj, err := s.cfg.DB.GetScheduledJob(name)
	if err != nil {
		return JobResult{}, fmt.Errorf("getting job %q: %w", name, err)
	}

	start := time.Now()
	findings, scanErr := s.cfg.ScanFunc(ctx, sj.ScanCommand)
	dur := time.Since(start)

	result := JobResult{
		JobName:      name,
		FindingCount: findings,
		Duration:     dur,
		Error:        scanErr,
	}

	// Update last_run regardless of error
	now := time.Now().UTC()
	_ = s.cfg.DB.UpdateJobLastRun(name, now, nil)

	if s.cfg.OnComplete != nil {
		s.cfg.OnComplete(result)
	}

	return result, nil
}

// JobCount returns the number of registered gocron jobs.
func (s *Scheduler) JobCount() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.jobs)
}

// registerJob creates a gocron cron job and stores the handle.
func (s *Scheduler) registerJob(ctx context.Context, sj storage.ScheduledJob) error {
	jobName := sj.Name
	scanCmd := sj.ScanCommand
	notify := sj.NotifyTelegram

	j, err := s.sched.NewJob(
		gocron.CronJob(sj.CronExpr, false),
		gocron.NewTask(func() {
			start := time.Now()
			findings, scanErr := s.cfg.ScanFunc(ctx, scanCmd)
			dur := time.Since(start)

			now := time.Now().UTC()
			_ = s.cfg.DB.UpdateJobLastRun(jobName, now, nil)

			if notify && s.cfg.OnComplete != nil {
				s.cfg.OnComplete(JobResult{
					JobName:      jobName,
					FindingCount: findings,
					Duration:     dur,
					Error:        scanErr,
				})
			}
		}),
	)
	if err != nil {
		return fmt.Errorf("creating gocron job: %w", err)
	}

	s.mu.Lock()
	s.jobs[sj.Name] = j
	s.mu.Unlock()
	return nil
}