From c71faa97f59a4dda0045dd2478db10374d886a9d Mon Sep 17 00:00:00 2001 From: salvacybersec Date: Mon, 6 Apr 2026 17:27:00 +0300 Subject: [PATCH] feat(17-02): implement scheduler package with gocron wrapper and job lifecycle - Scheduler wraps gocron with Start/Stop lifecycle - Start loads enabled jobs from DB and registers cron schedules - AddJob/RemoveJob persist to DB and sync with gocron - RunJob for manual trigger with OnComplete callback - JobResult struct for notification bridge - Promote gocron/v2 v2.19.1 to direct dependency --- go.mod | 6 +- go.sum | 2 + pkg/scheduler/jobs.go | 20 +++++ pkg/scheduler/scheduler.go | 178 +++++++++++++++++++++++++++++++++++++ 4 files changed, 203 insertions(+), 3 deletions(-) diff --git a/go.mod b/go.mod index f1a9d9c..001864f 100644 --- a/go.mod +++ b/go.mod @@ -5,6 +5,7 @@ go 1.26.1 require ( github.com/atotto/clipboard v0.1.4 github.com/charmbracelet/lipgloss v1.1.0 + github.com/go-co-op/gocron/v2 v2.19.1 github.com/go-git/go-git/v5 v5.17.2 github.com/mattn/go-isatty v0.0.20 github.com/panjf2000/ants/v2 v2.12.0 @@ -12,9 +13,11 @@ require ( github.com/spf13/cobra v1.10.2 github.com/spf13/viper v1.21.0 github.com/stretchr/testify v1.11.1 + github.com/temoto/robotstxt v1.1.2 github.com/tidwall/gjson v1.18.0 golang.org/x/crypto v0.49.0 golang.org/x/exp v0.0.0-20260312153236-7ab1446f8b90 + golang.org/x/net v0.52.0 golang.org/x/time v0.15.0 gopkg.in/yaml.v3 v3.0.1 modernc.org/sqlite v1.48.1 @@ -35,7 +38,6 @@ require ( github.com/dustin/go-humanize v1.0.1 // indirect github.com/emirpasic/gods v1.18.1 // indirect github.com/fsnotify/fsnotify v1.9.0 // indirect - github.com/go-co-op/gocron/v2 v2.19.1 // indirect github.com/go-git/gcfg v1.5.1-0.20230307220236-3a3c6141e376 // indirect github.com/go-git/go-billy/v5 v5.8.0 // indirect github.com/go-viper/mapstructure/v2 v2.4.0 // indirect @@ -63,13 +65,11 @@ require ( github.com/spf13/cast v1.10.0 // indirect github.com/spf13/pflag v1.0.10 // indirect github.com/subosito/gotenv v1.6.0 // indirect - github.com/temoto/robotstxt v1.1.2 // indirect github.com/tidwall/match v1.1.1 // indirect github.com/tidwall/pretty v1.2.0 // indirect github.com/xanzy/ssh-agent v0.3.3 // indirect github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect go.yaml.in/yaml/v3 v3.0.4 // indirect - golang.org/x/net v0.52.0 // indirect golang.org/x/sync v0.20.0 // indirect golang.org/x/sys v0.42.0 // indirect golang.org/x/text v0.35.0 // indirect diff --git a/go.sum b/go.sum index 8211d46..a0c74ea 100644 --- a/go.sum +++ b/go.sum @@ -154,6 +154,8 @@ github.com/xanzy/ssh-agent v0.3.3 h1:+/15pJfg/RsTxqYcX6fHqOXZwwMP+2VyYWJeWM2qQFM github.com/xanzy/ssh-agent v0.3.3/go.mod h1:6dzNDKs0J9rVPHPhaGCukekBHKqfl+L3KghI1Bc68Uw= github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= +go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= +go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc= go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= golang.org/x/crypto v0.0.0-20220622213112-05595931fe9d/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4= diff --git a/pkg/scheduler/jobs.go b/pkg/scheduler/jobs.go index 6990da0..539abf0 100644 --- a/pkg/scheduler/jobs.go +++ b/pkg/scheduler/jobs.go @@ -1 +1,21 @@ package scheduler + +import "time" + +// Job represents a scheduled scan job with its runtime function. +type Job struct { + Name string + CronExpr string + ScanCommand string + NotifyTelegram bool + Enabled bool + RunFunc func(ctx interface{}) (int, error) +} + +// JobResult contains the outcome of a scheduled or manually triggered scan job. +type JobResult struct { + JobName string + FindingCount int + Duration time.Duration + Error error +} diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go index 6990da0..f1325e7 100644 --- a/pkg/scheduler/scheduler.go +++ b/pkg/scheduler/scheduler.go @@ -1 +1,179 @@ 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 +}