diff --git a/.planning/REQUIREMENTS.md b/.planning/REQUIREMENTS.md index 1043627..dd22aa3 100644 --- a/.planning/REQUIREMENTS.md +++ b/.planning/REQUIREMENTS.md @@ -242,7 +242,7 @@ Requirements for initial release. Each maps to roadmap phases. ### Scheduled Scanning -- [ ] **SCHED-01**: Cron-based recurring scan scheduling +- [x] **SCHED-01**: Cron-based recurring scan scheduling - [ ] **SCHED-02**: keyhunter schedule add/list/remove commands - [ ] **SCHED-03**: Auto-notify on scheduled scan completion diff --git a/.planning/phases/17-telegram-scheduler/17-02-SUMMARY.md b/.planning/phases/17-telegram-scheduler/17-02-SUMMARY.md new file mode 100644 index 0000000..da13dfb --- /dev/null +++ b/.planning/phases/17-telegram-scheduler/17-02-SUMMARY.md @@ -0,0 +1,105 @@ +--- +phase: 17-telegram-scheduler +plan: 02 +subsystem: scheduler +tags: [gocron, sqlite, cron, scheduler, telegram] + +requires: + - phase: 01-foundation + provides: pkg/storage DB wrapper with schema.sql embed pattern +provides: + - pkg/scheduler/ package with gocron wrapper, start/stop lifecycle + - Storage CRUD for subscribers table (Add/Remove/List/IsSubscribed) + - Storage CRUD for scheduled_jobs table (Save/List/Get/Delete/UpdateLastRun/SetEnabled) + - subscribers and scheduled_jobs SQLite tables in schema.sql +affects: [17-telegram-scheduler, 17-03, 17-04, 17-05] + +tech-stack: + added: [gocron/v2 v2.19.1] + patterns: [scheduler wraps gocron with DB persistence, ScanFunc abstraction decouples from engine] + +key-files: + created: + - pkg/scheduler/scheduler.go + - pkg/scheduler/jobs.go + - pkg/scheduler/scheduler_test.go + - pkg/storage/subscribers.go + - pkg/storage/scheduled_jobs.go + modified: + - pkg/storage/schema.sql + - go.mod + - go.sum + +key-decisions: + - "Scheduler.ScanFunc callback decouples from engine -- Plan 17-04 wires the real scan logic" + - "OnComplete callback bridges scheduler to notification system without direct bot dependency" + - "Disabled jobs skipped during Start() but remain in DB for re-enabling" + +patterns-established: + - "Scheduler pattern: gocron wrapper with DB persistence and callback-based extensibility" + +requirements-completed: [SCHED-01] + +duration: 2min +completed: 2026-04-06 +--- + +# Phase 17 Plan 02: Scheduler + Storage Summary + +**gocron v2.19.1 wrapper with SQLite persistence for subscribers and scheduled scan jobs, callback-based scan/notify extensibility** + +## Performance + +- **Duration:** 2 min +- **Started:** 2026-04-06T14:25:04Z +- **Completed:** 2026-04-06T14:27:08Z +- **Tasks:** 2 +- **Files modified:** 8 + +## Accomplishments +- Created pkg/scheduler/ package wrapping gocron with Start/Stop lifecycle and DB-backed job persistence +- Implemented full CRUD for subscribers (Add/Remove/List/IsSubscribed) and scheduled_jobs (Save/List/Get/Delete/UpdateLastRun/SetEnabled) +- Added subscribers and scheduled_jobs tables to schema.sql +- All 5 tests pass: storage round-trip, subscriber round-trip, scheduler start/add/remove/run + +## Task Commits + +Each task was committed atomically: + +1. **Task 1: Add gocron dependency, create storage tables, and subscriber/job CRUD** - `c8f7592` (feat) +2. **Task 2 RED: Failing tests for scheduler package** - `89cc133` (test) +3. **Task 2 GREEN: Implement scheduler package** - `c71faa9` (feat) + +## Files Created/Modified +- `pkg/scheduler/scheduler.go` - Scheduler struct wrapping gocron with Start/Stop/AddJob/RemoveJob/RunJob/ListJobs +- `pkg/scheduler/jobs.go` - Job and JobResult types +- `pkg/scheduler/scheduler_test.go` - 5 tests covering storage, subscriber, and scheduler lifecycle +- `pkg/storage/subscribers.go` - Subscriber struct and CRUD methods on DB +- `pkg/storage/scheduled_jobs.go` - ScheduledJob struct and CRUD methods on DB +- `pkg/storage/schema.sql` - subscribers and scheduled_jobs CREATE TABLE statements +- `go.mod` - gocron/v2 v2.19.1 promoted to direct dependency +- `go.sum` - Updated checksums + +## Decisions Made +- ScanFunc callback decouples scheduler from engine -- Plan 17-04 wires real scan logic +- OnComplete callback bridges scheduler to notification system without direct bot dependency +- Disabled jobs skipped during Start() but remain in DB for re-enabling via SetJobEnabled + +## Deviations from Plan + +None - plan executed exactly as written. + +## Issues Encountered +None + +## User Setup Required +None - no external service configuration required. + +## Next Phase Readiness +- pkg/scheduler/ ready for CLI wiring in Plan 17-03 (schedule add/list/remove commands) +- Subscriber storage ready for bot /subscribe handler in Plan 17-01 +- OnComplete callback ready for notification bridge in Plan 17-04 + +--- +*Phase: 17-telegram-scheduler* +*Completed: 2026-04-06* diff --git a/go.mod b/go.mod index 798b5fe..0b672e5 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/mymmrac/telego v1.8.0 @@ -51,6 +52,7 @@ require ( github.com/grbit/go-json v0.11.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 // indirect + github.com/jonboulle/clockwork v0.5.0 // indirect github.com/kevinburke/ssh_config v1.2.0 // indirect github.com/klauspost/compress v1.18.2 // indirect github.com/klauspost/cpuid/v2 v2.2.9 // indirect @@ -63,6 +65,7 @@ require ( github.com/pmezard/go-difflib v1.0.0 // indirect github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect github.com/rivo/uniseg v0.4.7 // indirect + github.com/robfig/cron/v3 v3.0.1 // indirect github.com/sagikazarmark/locafero v0.11.0 // indirect github.com/sergi/go-diff v1.3.2-0.20230802210424-5b0b94c5c0d3 // indirect github.com/skeema/knownhosts v1.3.1 // indirect @@ -80,7 +83,10 @@ require ( 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 +<<<<<<< HEAD golang.org/x/arch v0.0.0-20210923205945-b76863e36670 // indirect +======= +>>>>>>> worktree-agent-a282d1fe 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 ca15ca4..606adfb 100644 --- a/go.sum +++ b/go.sum @@ -53,6 +53,8 @@ github.com/fsnotify/fsnotify v1.9.0 h1:2Ml+OJNzbYCTzsxtv8vKSFD9PbJjmhYF14k/jKC7S github.com/fsnotify/fsnotify v1.9.0/go.mod h1:8jBTzvmWwFyi3Pb8djgCCO5IBqzKJ/Jwo8TRcHyHii0= github.com/gliderlabs/ssh v0.3.8 h1:a4YXD1V7xMF9g5nTkdfnja3Sxy1PVDCj1Zg4Wb8vY6c= github.com/gliderlabs/ssh v0.3.8/go.mod h1:xYoytBv1sV0aL3CavoDuJIQNURXkkfPA/wxQ1pL1fAU= +github.com/go-co-op/gocron/v2 v2.19.1 h1:B4iLeA0NB/2iO3EKQ7NfKn5KsQgZfjb2fkvoZJU3yBI= +github.com/go-co-op/gocron/v2 v2.19.1/go.mod h1:5lEiCKk1oVJV39Zg7/YG10OnaVrDAV5GGR6O0663k6U= github.com/go-git/gcfg v1.5.1-0.20230307220236-3a3c6141e376 h1:+zs/tPmkDkHx3U66DAb0lQFJrpS6731Oaa12ikc+DiI= github.com/go-git/gcfg v1.5.1-0.20230307220236-3a3c6141e376/go.mod h1:an3vInlBmSxCcxctByoQdvwPiA7DTK7jaaFDBTtu0ic= github.com/go-git/go-billy/v5 v5.8.0 h1:I8hjc3LbBlXTtVuFNJuwYuMiHvQJDq1AT6u4DwDzZG0= @@ -79,6 +81,8 @@ github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 h1:BQSFePA1RWJOlocH6Fxy8MmwDt+yVQYULKfN0RoTN8A= github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99/go.mod h1:1lJo3i6rXxKeerYnT8Nvf0QmHCRC1n8sfWVwXF2Frvo= +github.com/jonboulle/clockwork v0.5.0 h1:Hyh9A8u51kptdkR+cqRpT1EebBwTn1oK9YfGYbdFz6I= +github.com/jonboulle/clockwork v0.5.0/go.mod h1:3mZlmanh0g2NDKO5TWZVJAfofYk64M7XN3SzBPjZF60= github.com/kevinburke/ssh_config v1.2.0 h1:x584FjTGwHzMwvHx18PXxbBVzfnxogHaAReU4gf13a4= github.com/kevinburke/ssh_config v1.2.0/go.mod h1:CT57kijsi8u/K/BOFA39wgDQJ9CxiF4nAY/ojJ6r6mM= github.com/klauspost/compress v1.18.2 h1:iiPHWW0YrcFgpBYhsA6D1+fqHssJscY/Tm/y2Uqnapk= @@ -123,6 +127,8 @@ github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qq github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= +github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs= +github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro= github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= @@ -181,10 +187,15 @@ 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= +<<<<<<< HEAD github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU= github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= go.uber.org/mock v0.6.0 h1:hyF9dfmbgIX5EfOdasqLsWD6xqpNZlXblLB/Dbnwv3Y= go.uber.org/mock v0.6.0/go.mod h1:KiVJ4BqZJaMj4svdfmHM0AUx4NJYO8ZNpPnZn1Z+BBU= +======= +go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= +go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= +>>>>>>> worktree-agent-a282d1fe 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/arch v0.0.0-20210923205945-b76863e36670 h1:18EFjUmQOcUvxNYSkA6jO9VAiXCnxFY6NyDX0bHDmkU= diff --git a/pkg/scheduler/jobs.go b/pkg/scheduler/jobs.go new file mode 100644 index 0000000..539abf0 --- /dev/null +++ b/pkg/scheduler/jobs.go @@ -0,0 +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 new file mode 100644 index 0000000..f1325e7 --- /dev/null +++ b/pkg/scheduler/scheduler.go @@ -0,0 +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 +} diff --git a/pkg/scheduler/scheduler_test.go b/pkg/scheduler/scheduler_test.go new file mode 100644 index 0000000..9ce8879 --- /dev/null +++ b/pkg/scheduler/scheduler_test.go @@ -0,0 +1,204 @@ +package scheduler_test + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/salvacybersec/keyhunter/pkg/scheduler" + "github.com/salvacybersec/keyhunter/pkg/storage" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func openTestDB(t *testing.T) *storage.DB { + t.Helper() + db, err := storage.Open(":memory:") + require.NoError(t, err) + t.Cleanup(func() { db.Close() }) + return db +} + +func TestStorageRoundTrip(t *testing.T) { + db := openTestDB(t) + + id, err := db.SaveScheduledJob(storage.ScheduledJob{ + Name: "nightly-scan", + CronExpr: "0 2 * * *", + ScanCommand: "/tmp/repos", + NotifyTelegram: true, + Enabled: true, + }) + require.NoError(t, err) + assert.Greater(t, id, int64(0)) + + jobs, err := db.ListScheduledJobs() + require.NoError(t, err) + require.Len(t, jobs, 1) + assert.Equal(t, "nightly-scan", jobs[0].Name) + assert.Equal(t, "0 2 * * *", jobs[0].CronExpr) + assert.Equal(t, "/tmp/repos", jobs[0].ScanCommand) + assert.True(t, jobs[0].NotifyTelegram) + assert.True(t, jobs[0].Enabled) + + got, err := db.GetScheduledJob("nightly-scan") + require.NoError(t, err) + assert.Equal(t, "nightly-scan", got.Name) + + now := time.Now().UTC() + next := now.Add(24 * time.Hour) + err = db.UpdateJobLastRun("nightly-scan", now, &next) + require.NoError(t, err) + + got2, err := db.GetScheduledJob("nightly-scan") + require.NoError(t, err) + require.NotNil(t, got2.LastRun) + require.NotNil(t, got2.NextRun) + + n, err := db.DeleteScheduledJob("nightly-scan") + require.NoError(t, err) + assert.Equal(t, int64(1), n) + + jobs, err = db.ListScheduledJobs() + require.NoError(t, err) + assert.Empty(t, jobs) +} + +func TestSubscriberRoundTrip(t *testing.T) { + db := openTestDB(t) + + err := db.AddSubscriber(12345, "alice") + require.NoError(t, err) + + subs, err := db.ListSubscribers() + require.NoError(t, err) + require.Len(t, subs, 1) + assert.Equal(t, int64(12345), subs[0].ChatID) + assert.Equal(t, "alice", subs[0].Username) + + ok, err := db.IsSubscribed(12345) + require.NoError(t, err) + assert.True(t, ok) + + ok, err = db.IsSubscribed(99999) + require.NoError(t, err) + assert.False(t, ok) + + n, err := db.RemoveSubscriber(12345) + require.NoError(t, err) + assert.Equal(t, int64(1), n) + + subs, err = db.ListSubscribers() + require.NoError(t, err) + assert.Empty(t, subs) +} + +func TestSchedulerStartLoadsJobs(t *testing.T) { + db := openTestDB(t) + + _, err := db.SaveScheduledJob(storage.ScheduledJob{ + Name: "job-a", CronExpr: "0 * * * *", ScanCommand: "/a", Enabled: true, + }) + require.NoError(t, err) + _, err = db.SaveScheduledJob(storage.ScheduledJob{ + Name: "job-b", CronExpr: "0 * * * *", ScanCommand: "/b", Enabled: true, + }) + require.NoError(t, err) + // Disabled job should not be registered + _, err = db.SaveScheduledJob(storage.ScheduledJob{ + Name: "job-c", CronExpr: "0 * * * *", ScanCommand: "/c", Enabled: false, + }) + require.NoError(t, err) + + s, err := scheduler.New(scheduler.Config{ + DB: db, + ScanFunc: func(ctx context.Context, cmd string) (int, error) { + return 0, nil + }, + }) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + err = s.Start(ctx) + require.NoError(t, err) + defer s.Stop() + + assert.Equal(t, 2, s.JobCount()) +} + +func TestSchedulerAddRemoveJob(t *testing.T) { + db := openTestDB(t) + + s, err := scheduler.New(scheduler.Config{ + DB: db, + ScanFunc: func(ctx context.Context, cmd string) (int, error) { + return 0, nil + }, + }) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + err = s.Start(ctx) + require.NoError(t, err) + defer s.Stop() + + err = s.AddJob("test-job", "0 * * * *", "/test", true) + require.NoError(t, err) + assert.Equal(t, 1, s.JobCount()) + + jobs, err := db.ListScheduledJobs() + require.NoError(t, err) + require.Len(t, jobs, 1) + assert.Equal(t, "test-job", jobs[0].Name) + + err = s.RemoveJob("test-job") + require.NoError(t, err) + assert.Equal(t, 0, s.JobCount()) + + jobs, err = db.ListScheduledJobs() + require.NoError(t, err) + assert.Empty(t, jobs) +} + +func TestSchedulerRunJob(t *testing.T) { + db := openTestDB(t) + + var mu sync.Mutex + var scanCalled string + var completeCalled bool + + s, err := scheduler.New(scheduler.Config{ + DB: db, + ScanFunc: func(ctx context.Context, cmd string) (int, error) { + mu.Lock() + scanCalled = cmd + mu.Unlock() + return 5, nil + }, + OnComplete: func(result scheduler.JobResult) { + mu.Lock() + completeCalled = true + mu.Unlock() + }, + }) + require.NoError(t, err) + + _, err = db.SaveScheduledJob(storage.ScheduledJob{ + Name: "manual-run", CronExpr: "0 * * * *", ScanCommand: "/manual", NotifyTelegram: true, Enabled: true, + }) + require.NoError(t, err) + + result, err := s.RunJob(context.Background(), "manual-run") + require.NoError(t, err) + assert.Equal(t, "manual-run", result.JobName) + assert.Equal(t, 5, result.FindingCount) + + mu.Lock() + assert.Equal(t, "/manual", scanCalled) + assert.True(t, completeCalled) + mu.Unlock() +} diff --git a/pkg/storage/scheduled_jobs.go b/pkg/storage/scheduled_jobs.go new file mode 100644 index 0000000..7315e3f --- /dev/null +++ b/pkg/storage/scheduled_jobs.go @@ -0,0 +1,111 @@ +package storage + +import ( + "database/sql" + "time" +) + +// ScheduledJob represents a cron-based scheduled scan job. +type ScheduledJob struct { + ID int64 + Name string + CronExpr string + ScanCommand string + NotifyTelegram bool + Enabled bool + LastRun *time.Time + NextRun *time.Time + CreatedAt time.Time +} + +// SaveScheduledJob inserts a new scheduled job. Returns the new row ID. +func (db *DB) SaveScheduledJob(j ScheduledJob) (int64, error) { + res, err := db.sql.Exec( + `INSERT INTO scheduled_jobs (name, cron_expr, scan_command, notify_telegram, enabled) + VALUES (?, ?, ?, ?, ?)`, + j.Name, j.CronExpr, j.ScanCommand, j.NotifyTelegram, j.Enabled, + ) + if err != nil { + return 0, err + } + return res.LastInsertId() +} + +// ListScheduledJobs returns all scheduled jobs. +func (db *DB) ListScheduledJobs() ([]ScheduledJob, error) { + rows, err := db.sql.Query( + `SELECT id, name, cron_expr, scan_command, notify_telegram, enabled, last_run, next_run, created_at + FROM scheduled_jobs ORDER BY id`, + ) + if err != nil { + return nil, err + } + defer rows.Close() + + var jobs []ScheduledJob + for rows.Next() { + var j ScheduledJob + var lastRun, nextRun sql.NullTime + if err := rows.Scan(&j.ID, &j.Name, &j.CronExpr, &j.ScanCommand, + &j.NotifyTelegram, &j.Enabled, &lastRun, &nextRun, &j.CreatedAt); err != nil { + return nil, err + } + if lastRun.Valid { + j.LastRun = &lastRun.Time + } + if nextRun.Valid { + j.NextRun = &nextRun.Time + } + jobs = append(jobs, j) + } + return jobs, rows.Err() +} + +// GetScheduledJob returns a single scheduled job by name. +func (db *DB) GetScheduledJob(name string) (*ScheduledJob, error) { + var j ScheduledJob + var lastRun, nextRun sql.NullTime + err := db.sql.QueryRow( + `SELECT id, name, cron_expr, scan_command, notify_telegram, enabled, last_run, next_run, created_at + FROM scheduled_jobs WHERE name = ?`, name, + ).Scan(&j.ID, &j.Name, &j.CronExpr, &j.ScanCommand, + &j.NotifyTelegram, &j.Enabled, &lastRun, &nextRun, &j.CreatedAt) + if err != nil { + return nil, err + } + if lastRun.Valid { + j.LastRun = &lastRun.Time + } + if nextRun.Valid { + j.NextRun = &nextRun.Time + } + return &j, nil +} + +// DeleteScheduledJob removes a scheduled job by name. Returns rows affected. +func (db *DB) DeleteScheduledJob(name string) (int64, error) { + res, err := db.sql.Exec(`DELETE FROM scheduled_jobs WHERE name = ?`, name) + if err != nil { + return 0, err + } + return res.RowsAffected() +} + +// UpdateJobLastRun updates the last_run and next_run timestamps for a job. +func (db *DB) UpdateJobLastRun(name string, lastRun time.Time, nextRun *time.Time) error { + var nr sql.NullTime + if nextRun != nil { + nr = sql.NullTime{Time: *nextRun, Valid: true} + } + _, err := db.sql.Exec( + `UPDATE scheduled_jobs SET last_run = ?, next_run = ? WHERE name = ?`, + lastRun, nr, name, + ) + return err +} + +// SetJobEnabled updates the enabled flag for a scheduled job. +func (db *DB) SetJobEnabled(name string, enabled bool) error { + _, err := db.sql.Exec(`UPDATE scheduled_jobs SET enabled = ? WHERE name = ?`, enabled, name) + return err +} diff --git a/pkg/storage/schema.sql b/pkg/storage/schema.sql index 964dad9..84e91d9 100644 --- a/pkg/storage/schema.sql +++ b/pkg/storage/schema.sql @@ -55,3 +55,23 @@ CREATE TABLE IF NOT EXISTS custom_dorks ( CREATE INDEX IF NOT EXISTS idx_custom_dorks_source ON custom_dorks(source); CREATE INDEX IF NOT EXISTS idx_custom_dorks_category ON custom_dorks(category); + +-- Phase 17: Telegram bot subscribers for auto-notifications. +CREATE TABLE IF NOT EXISTS subscribers ( + chat_id INTEGER PRIMARY KEY, + username TEXT, + subscribed_at DATETIME DEFAULT CURRENT_TIMESTAMP +); + +-- Phase 17: Cron-based scheduled scan jobs. +CREATE TABLE IF NOT EXISTS scheduled_jobs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name TEXT UNIQUE NOT NULL, + cron_expr TEXT NOT NULL, + scan_command TEXT NOT NULL, + notify_telegram BOOLEAN DEFAULT FALSE, + enabled BOOLEAN DEFAULT TRUE, + last_run DATETIME, + next_run DATETIME, + created_at DATETIME DEFAULT CURRENT_TIMESTAMP +); diff --git a/pkg/storage/subscribers.go b/pkg/storage/subscribers.go new file mode 100644 index 0000000..fa6fd3b --- /dev/null +++ b/pkg/storage/subscribers.go @@ -0,0 +1,54 @@ +package storage + +import "time" + +// Subscriber represents a Telegram chat subscribed to scan notifications. +type Subscriber struct { + ChatID int64 + Username string + SubscribedAt time.Time +} + +// AddSubscriber inserts or replaces a subscriber in the database. +func (db *DB) AddSubscriber(chatID int64, username string) error { + _, err := db.sql.Exec( + `INSERT OR REPLACE INTO subscribers (chat_id, username, subscribed_at) VALUES (?, ?, CURRENT_TIMESTAMP)`, + chatID, username, + ) + return err +} + +// RemoveSubscriber deletes a subscriber by chat ID. Returns rows affected. +func (db *DB) RemoveSubscriber(chatID int64) (int64, error) { + res, err := db.sql.Exec(`DELETE FROM subscribers WHERE chat_id = ?`, chatID) + if err != nil { + return 0, err + } + return res.RowsAffected() +} + +// ListSubscribers returns all subscribers ordered by subscription time. +func (db *DB) ListSubscribers() ([]Subscriber, error) { + rows, err := db.sql.Query(`SELECT chat_id, username, subscribed_at FROM subscribers ORDER BY subscribed_at`) + if err != nil { + return nil, err + } + defer rows.Close() + + var subs []Subscriber + for rows.Next() { + var s Subscriber + if err := rows.Scan(&s.ChatID, &s.Username, &s.SubscribedAt); err != nil { + return nil, err + } + subs = append(subs, s) + } + return subs, rows.Err() +} + +// IsSubscribed returns true if the given chat ID is subscribed. +func (db *DB) IsSubscribed(chatID int64) (bool, error) { + var count int + err := db.sql.QueryRow(`SELECT COUNT(*) FROM subscribers WHERE chat_id = ?`, chatID).Scan(&count) + return count > 0, err +}