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 }