feat(17-05): implement serve and schedule commands replacing stubs
- cmd/serve.go: starts scheduler, optionally starts Telegram bot with --telegram flag - cmd/schedule.go: add/list/remove/run subcommands for scheduled scan job CRUD - pkg/scheduler/: gocron v2 based scheduler with DB-backed jobs and scan execution - pkg/storage/scheduled_jobs.go: scheduled_jobs table CRUD with tests - Remove serve and schedule stubs from cmd/stubs.go Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
salvacybersec
279
cmd/schedule.go
Normal file
279
cmd/schedule.go
Normal file
@@ -0,0 +1,279 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/spf13/cobra"
|
||||
"github.com/spf13/viper"
|
||||
|
||||
"github.com/salvacybersec/keyhunter/pkg/config"
|
||||
"github.com/salvacybersec/keyhunter/pkg/engine"
|
||||
"github.com/salvacybersec/keyhunter/pkg/providers"
|
||||
"github.com/salvacybersec/keyhunter/pkg/storage"
|
||||
)
|
||||
|
||||
var (
|
||||
schedCron string
|
||||
schedScan string
|
||||
schedName string
|
||||
schedNotify bool
|
||||
)
|
||||
|
||||
var scheduleCmd = &cobra.Command{
|
||||
Use: "schedule",
|
||||
Short: "Manage scheduled recurring scans",
|
||||
Long: `Add, list, remove, or manually run scheduled scan jobs.
|
||||
Jobs are stored in the database and executed by 'keyhunter serve'.`,
|
||||
}
|
||||
|
||||
var scheduleAddCmd = &cobra.Command{
|
||||
Use: "add",
|
||||
Short: "Add a new scheduled scan job",
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
if schedCron == "" {
|
||||
return fmt.Errorf("--cron is required (e.g. --cron=\"0 */6 * * *\")")
|
||||
}
|
||||
if schedScan == "" {
|
||||
return fmt.Errorf("--scan is required (path to scan)")
|
||||
}
|
||||
|
||||
db, cleanup, err := openScheduleDB()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer cleanup()
|
||||
|
||||
name := schedName
|
||||
if name == "" {
|
||||
name = fmt.Sprintf("scan-%s", filepath.Base(schedScan))
|
||||
}
|
||||
|
||||
job := storage.ScheduledJob{
|
||||
Name: name,
|
||||
CronExpr: schedCron,
|
||||
ScanPath: schedScan,
|
||||
Enabled: true,
|
||||
Notify: schedNotify,
|
||||
}
|
||||
|
||||
id, err := db.SaveScheduledJob(job)
|
||||
if err != nil {
|
||||
return fmt.Errorf("saving scheduled job: %w", err)
|
||||
}
|
||||
|
||||
fmt.Printf("Scheduled job #%d added:\n", id)
|
||||
fmt.Printf(" Name: %s\n", name)
|
||||
fmt.Printf(" Cron: %s\n", schedCron)
|
||||
fmt.Printf(" Path: %s\n", schedScan)
|
||||
fmt.Printf(" Notify: %t\n", schedNotify)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
var scheduleListCmd = &cobra.Command{
|
||||
Use: "list",
|
||||
Short: "List all scheduled scan jobs",
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
db, cleanup, err := openScheduleDB()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer cleanup()
|
||||
|
||||
jobs, err := db.ListScheduledJobs()
|
||||
if err != nil {
|
||||
return fmt.Errorf("listing scheduled jobs: %w", err)
|
||||
}
|
||||
|
||||
if len(jobs) == 0 {
|
||||
fmt.Println("No scheduled jobs. Use 'keyhunter schedule add' to create one.")
|
||||
return nil
|
||||
}
|
||||
|
||||
fmt.Printf("%-4s %-20s %-20s %-30s %-8s %-7s %s\n",
|
||||
"ID", "Name", "Cron", "Path", "Enabled", "Notify", "Last Run")
|
||||
fmt.Println(strings.Repeat("-", 120))
|
||||
|
||||
for _, j := range jobs {
|
||||
lastRun := "never"
|
||||
if j.LastRunAt != nil {
|
||||
lastRun = j.LastRunAt.Format(time.RFC3339)
|
||||
}
|
||||
enabled := "yes"
|
||||
if !j.Enabled {
|
||||
enabled = "no"
|
||||
}
|
||||
notify := "yes"
|
||||
if !j.Notify {
|
||||
notify = "no"
|
||||
}
|
||||
fmt.Printf("%-4d %-20s %-20s %-30s %-8s %-7s %s\n",
|
||||
j.ID, truncateStr(j.Name, 20), j.CronExpr, truncateStr(j.ScanPath, 30),
|
||||
enabled, notify, lastRun)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
var scheduleRemoveCmd = &cobra.Command{
|
||||
Use: "remove <id>",
|
||||
Short: "Remove a scheduled scan job by ID",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
id, err := strconv.ParseInt(args[0], 10, 64)
|
||||
if err != nil || id <= 0 {
|
||||
return fmt.Errorf("invalid job ID: %s", args[0])
|
||||
}
|
||||
|
||||
db, cleanup, err := openScheduleDB()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer cleanup()
|
||||
|
||||
affected, err := db.DeleteScheduledJob(id)
|
||||
if err != nil {
|
||||
return fmt.Errorf("removing job: %w", err)
|
||||
}
|
||||
if affected == 0 {
|
||||
return fmt.Errorf("no job with ID %d", id)
|
||||
}
|
||||
fmt.Printf("Scheduled job #%d removed.\n", id)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
var scheduleRunCmd = &cobra.Command{
|
||||
Use: "run <id>",
|
||||
Short: "Manually run a scheduled scan job now",
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
id, err := strconv.ParseInt(args[0], 10, 64)
|
||||
if err != nil || id <= 0 {
|
||||
return fmt.Errorf("invalid job ID: %s", args[0])
|
||||
}
|
||||
|
||||
db, cleanup, err := openScheduleDB()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer cleanup()
|
||||
|
||||
job, err := db.GetScheduledJob(id)
|
||||
if err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return fmt.Errorf("no job with ID %d", id)
|
||||
}
|
||||
return fmt.Errorf("fetching job: %w", err)
|
||||
}
|
||||
|
||||
cfg := config.Load()
|
||||
|
||||
// Derive encryption key.
|
||||
encKey, err := loadOrCreateEncKey(db, cfg.Passphrase)
|
||||
if err != nil {
|
||||
return fmt.Errorf("preparing encryption key: %w", err)
|
||||
}
|
||||
|
||||
// Initialize engine.
|
||||
reg, err := providers.NewRegistry()
|
||||
if err != nil {
|
||||
return fmt.Errorf("loading providers: %w", err)
|
||||
}
|
||||
eng := engine.NewEngine(reg)
|
||||
|
||||
fmt.Printf("Running job #%d (%s) scanning %s...\n", job.ID, job.Name, job.ScanPath)
|
||||
|
||||
// Select source and scan.
|
||||
src, err := selectSource([]string{job.ScanPath}, sourceFlags{})
|
||||
if err != nil {
|
||||
return fmt.Errorf("selecting source: %w", err)
|
||||
}
|
||||
|
||||
scanCfg := engine.ScanConfig{
|
||||
Workers: 0, // auto
|
||||
Verify: false,
|
||||
Unmask: false,
|
||||
}
|
||||
|
||||
ch, scanErr := eng.Scan(context.Background(), src, scanCfg)
|
||||
if scanErr != nil {
|
||||
return fmt.Errorf("starting scan: %w", scanErr)
|
||||
}
|
||||
|
||||
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 := db.SaveFinding(sf, encKey); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "warning: failed to save finding: %v\n", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Update last run time.
|
||||
if err := db.UpdateJobLastRun(job.ID, time.Now()); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "warning: failed to update last_run: %v\n", err)
|
||||
}
|
||||
|
||||
fmt.Printf("Scan complete. Found %d key(s).\n", len(findings))
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
// openScheduleDB opens the database for schedule commands.
|
||||
func openScheduleDB() (*storage.DB, func(), error) {
|
||||
cfg := config.Load()
|
||||
dbPath := viper.GetString("database.path")
|
||||
if dbPath == "" {
|
||||
dbPath = cfg.DBPath
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(dbPath), 0700); err != nil {
|
||||
return nil, nil, fmt.Errorf("creating database directory: %w", err)
|
||||
}
|
||||
db, err := storage.Open(dbPath)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("opening database: %w", err)
|
||||
}
|
||||
return db, func() { db.Close() }, nil
|
||||
}
|
||||
|
||||
// truncateStr shortens a string to max length with ellipsis.
|
||||
// Named differently from dorks.go truncate to avoid redeclaration.
|
||||
func truncateStr(s string, max int) string {
|
||||
if len(s) <= max {
|
||||
return s
|
||||
}
|
||||
return s[:max-3] + "..."
|
||||
}
|
||||
|
||||
func init() {
|
||||
scheduleAddCmd.Flags().StringVar(&schedCron, "cron", "", "cron expression (e.g. \"0 */6 * * *\")")
|
||||
scheduleAddCmd.Flags().StringVar(&schedScan, "scan", "", "path to scan")
|
||||
scheduleAddCmd.Flags().StringVar(&schedName, "name", "", "job name (default: auto-generated)")
|
||||
scheduleAddCmd.Flags().BoolVar(&schedNotify, "notify", true, "send Telegram notification on findings")
|
||||
|
||||
scheduleCmd.AddCommand(scheduleAddCmd)
|
||||
scheduleCmd.AddCommand(scheduleListCmd)
|
||||
scheduleCmd.AddCommand(scheduleRemoveCmd)
|
||||
scheduleCmd.AddCommand(scheduleRunCmd)
|
||||
}
|
||||
Reference in New Issue
Block a user