diff --git a/pkg/scheduler/jobs.go b/pkg/scheduler/jobs.go new file mode 100644 index 0000000..6990da0 --- /dev/null +++ b/pkg/scheduler/jobs.go @@ -0,0 +1 @@ +package scheduler diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go new file mode 100644 index 0000000..6990da0 --- /dev/null +++ b/pkg/scheduler/scheduler.go @@ -0,0 +1 @@ +package scheduler 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() +}