keyhunter/pkg/scheduler/scheduler.go

// Package scheduler implements cron-based recurring scan scheduling for KeyHunter.
// It uses gocron v2 for job management and delegates scan execution to the engine.
package scheduler

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

	"github.com/go-co-op/gocron/v2"

	"github.com/salvacybersec/keyhunter/pkg/engine"
	"github.com/salvacybersec/keyhunter/pkg/storage"
)

// Scheduler manages recurring scan jobs backed by the database.
type Scheduler struct {
	cron   gocron.Scheduler
	engine *engine.Engine
	db     *storage.DB
	encKey []byte

	mu   sync.Mutex
	jobs map[int64]gocron.Job // DB job ID -> gocron job

	// OnFindings is called when a scheduled scan produces findings.
	// The caller can wire this to Telegram notifications.
	OnFindings func(jobName string, findings []engine.Finding)
}

// Deps bundles the dependencies for creating a Scheduler.
type Deps struct {
	Engine *engine.Engine
	DB     *storage.DB
	EncKey []byte
}

// New creates a new Scheduler. Call Start() to begin processing jobs.
func New(deps Deps) (*Scheduler, error) {
	cron, err := gocron.NewScheduler()
	if err != nil {
		return nil, fmt.Errorf("creating gocron scheduler: %w", err)
	}
	return &Scheduler{
		cron:   cron,
		engine: deps.Engine,
		db:     deps.DB,
		encKey: deps.EncKey,
		jobs:   make(map[int64]gocron.Job),
	}, nil
}

// LoadAndStart loads all enabled jobs from the database, registers them
// with gocron, and starts the scheduler.
func (s *Scheduler) LoadAndStart() error {
	jobs, err := s.db.ListEnabledScheduledJobs()
	if err != nil {
		return fmt.Errorf("loading scheduled jobs: %w", err)
	}
	for _, j := range jobs {
		if err := s.registerJob(j); err != nil {
			log.Printf("scheduler: failed to register job %d (%s): %v", j.ID, j.Name, err)
		}
	}
	s.cron.Start()
	return nil
}

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

// AddJob registers a new job from a storage.ScheduledJob and adds it to the
// running scheduler.
func (s *Scheduler) AddJob(job storage.ScheduledJob) error {
	return s.registerJob(job)
}

// RemoveJob removes a job from the running scheduler by its DB ID.
func (s *Scheduler) RemoveJob(id int64) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if j, ok := s.jobs[id]; ok {
		s.cron.RemoveJob(j.ID())
		delete(s.jobs, id)
	}
}

// RunNow executes a job immediately (outside of its cron schedule).
func (s *Scheduler) RunNow(job storage.ScheduledJob) ([]engine.Finding, error) {
	return s.executeScan(job)
}

// registerJob adds a single scheduled job to gocron.
func (s *Scheduler) registerJob(job storage.ScheduledJob) error {
	jobCopy := job // capture for closure
	cronJob, err := s.cron.NewJob(
		gocron.CronJob(job.CronExpr, false),
		gocron.NewTask(func() {
			findings, err := s.executeScan(jobCopy)
			if err != nil {
				log.Printf("scheduler: job %d (%s) failed: %v", jobCopy.ID, jobCopy.Name, err)
				return
			}
			// Update last run time in DB.
			if err := s.db.UpdateJobLastRun(jobCopy.ID, time.Now()); err != nil {
				log.Printf("scheduler: failed to update last_run for job %d: %v", jobCopy.ID, err)
			}
			if len(findings) > 0 && jobCopy.Notify && s.OnFindings != nil {
				s.OnFindings(jobCopy.Name, findings)
			}
		}),
	)
	if err != nil {
		return fmt.Errorf("registering cron job %q (%s): %w", job.Name, job.CronExpr, err)
	}

	s.mu.Lock()
	s.jobs[job.ID] = cronJob
	s.mu.Unlock()
	return nil
}

// executeScan runs a scan against the job's configured path and persists findings.
func (s *Scheduler) executeScan(job storage.ScheduledJob) ([]engine.Finding, error) {
	src, err := selectSchedulerSource(job.ScanPath)
	if err != nil {
		return nil, fmt.Errorf("selecting source for %q: %w", job.ScanPath, err)
	}

	cfg := engine.ScanConfig{
		Workers: 0, // auto
		Verify:  false,
		Unmask:  false,
	}

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Minute)
	defer cancel()

	ch, err := s.engine.Scan(ctx, src, cfg)
	if err != nil {
		return nil, fmt.Errorf("starting scan: %w", err)
	}

	var findings []engine.Finding
	for f := range ch {
		findings = append(findings, f)
	}

	// Persist findings.
	for _, f := range findings {
		sf := storage.Finding{
			ProviderName: f.ProviderName,
			KeyValue:     f.KeyValue,
			KeyMasked:    f.KeyMasked,
			Confidence:   f.Confidence,
			SourcePath:   f.Source,
			SourceType:   f.SourceType,
			LineNumber:   f.LineNumber,
		}
		if _, err := s.db.SaveFinding(sf, s.encKey); err != nil {
			log.Printf("scheduler: failed to save finding: %v", err)
		}
	}

	return findings, nil
}