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)
|
||||
}
|
||||
197
cmd/serve.go
Normal file
197
cmd/serve.go
Normal file
@@ -0,0 +1,197 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"os/signal"
|
||||
"path/filepath"
|
||||
"syscall"
|
||||
|
||||
"github.com/mymmrac/telego"
|
||||
"github.com/spf13/cobra"
|
||||
"github.com/spf13/viper"
|
||||
|
||||
"github.com/salvacybersec/keyhunter/pkg/bot"
|
||||
"github.com/salvacybersec/keyhunter/pkg/config"
|
||||
"github.com/salvacybersec/keyhunter/pkg/engine"
|
||||
"github.com/salvacybersec/keyhunter/pkg/providers"
|
||||
"github.com/salvacybersec/keyhunter/pkg/recon"
|
||||
reconSources "github.com/salvacybersec/keyhunter/pkg/recon/sources"
|
||||
"github.com/salvacybersec/keyhunter/pkg/scheduler"
|
||||
"github.com/salvacybersec/keyhunter/pkg/storage"
|
||||
"github.com/salvacybersec/keyhunter/pkg/verify"
|
||||
)
|
||||
|
||||
var (
|
||||
serveTelegram bool
|
||||
servePort int
|
||||
)
|
||||
|
||||
var serveCmd = &cobra.Command{
|
||||
Use: "serve",
|
||||
Short: "Start the scheduler (and optionally Telegram bot)",
|
||||
Long: `Start KeyHunter in server mode. The scheduler runs all enabled recurring scan
|
||||
jobs defined via 'keyhunter schedule add'. If --telegram is specified, the
|
||||
Telegram bot is also started for remote control.`,
|
||||
RunE: func(cmd *cobra.Command, args []string) error {
|
||||
cfg := config.Load()
|
||||
|
||||
// Open database.
|
||||
dbPath := viper.GetString("database.path")
|
||||
if dbPath == "" {
|
||||
dbPath = cfg.DBPath
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(dbPath), 0700); err != nil {
|
||||
return fmt.Errorf("creating database directory: %w", err)
|
||||
}
|
||||
db, err := storage.Open(dbPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("opening database: %w", err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
// Derive encryption key.
|
||||
encKey, err := loadOrCreateEncKey(db, cfg.Passphrase)
|
||||
if err != nil {
|
||||
return fmt.Errorf("preparing encryption key: %w", err)
|
||||
}
|
||||
|
||||
// Initialize provider registry and engine.
|
||||
reg, err := providers.NewRegistry()
|
||||
if err != nil {
|
||||
return fmt.Errorf("loading providers: %w", err)
|
||||
}
|
||||
eng := engine.NewEngine(reg)
|
||||
|
||||
// Initialize scheduler.
|
||||
sched, err := scheduler.New(scheduler.Deps{
|
||||
Engine: eng,
|
||||
DB: db,
|
||||
EncKey: encKey,
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("creating scheduler: %w", err)
|
||||
}
|
||||
|
||||
// Optionally start Telegram bot.
|
||||
var telegramBot *bot.Bot
|
||||
if serveTelegram {
|
||||
token := viper.GetString("telegram.token")
|
||||
if token == "" {
|
||||
token = os.Getenv("KEYHUNTER_TELEGRAM_TOKEN")
|
||||
}
|
||||
if token == "" {
|
||||
return fmt.Errorf("telegram token required: set telegram.token in config or KEYHUNTER_TELEGRAM_TOKEN env var")
|
||||
}
|
||||
|
||||
api, err := telego.NewBot(token)
|
||||
if err != nil {
|
||||
return fmt.Errorf("creating telegram bot: %w", err)
|
||||
}
|
||||
|
||||
verifier := verify.NewHTTPVerifier(0) // default timeout
|
||||
reconEng := buildServeReconEngine()
|
||||
|
||||
telegramBot = bot.New(api, bot.Deps{
|
||||
Engine: eng,
|
||||
Verifier: verifier,
|
||||
Recon: reconEng,
|
||||
DB: db,
|
||||
Registry: reg,
|
||||
EncKey: encKey,
|
||||
})
|
||||
|
||||
// Wire scheduler notifications to Telegram.
|
||||
sched.OnFindings = func(jobName string, findings []engine.Finding) {
|
||||
msg := fmt.Sprintf("Scheduled scan %q found %d key(s):\n", jobName, len(findings))
|
||||
for i, f := range findings {
|
||||
if i >= 10 {
|
||||
msg += fmt.Sprintf("\n... and %d more", len(findings)-10)
|
||||
break
|
||||
}
|
||||
msg += fmt.Sprintf("[%s] %s %s:%d\n", f.ProviderName, f.KeyMasked, f.Source, f.LineNumber)
|
||||
}
|
||||
// Broadcast to all subscribed chats via bot status message.
|
||||
// For now, log. Subscribe/broadcast will be wired in a future plan.
|
||||
log.Printf("scheduler notification: %s", msg)
|
||||
}
|
||||
|
||||
updates, err := api.UpdatesViaLongPolling(context.Background(), nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("starting telegram long polling: %w", err)
|
||||
}
|
||||
|
||||
bh := telegramBot.RegisterHandlers(updates)
|
||||
go bh.Start()
|
||||
fmt.Println("Telegram bot started.")
|
||||
}
|
||||
|
||||
// Load and start scheduler.
|
||||
if err := sched.LoadAndStart(); err != nil {
|
||||
return fmt.Errorf("starting scheduler: %w", err)
|
||||
}
|
||||
fmt.Printf("Scheduler started (port %d placeholder for future web dashboard).\n", servePort)
|
||||
fmt.Println("Press Ctrl+C to stop.")
|
||||
|
||||
// Wait for signal.
|
||||
sigCh := make(chan os.Signal, 1)
|
||||
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
|
||||
<-sigCh
|
||||
|
||||
fmt.Println("\nShutting down...")
|
||||
if err := sched.Stop(); err != nil {
|
||||
log.Printf("scheduler shutdown error: %v", err)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
|
||||
// buildServeReconEngine creates a recon engine with all registered sources
|
||||
// for the serve command's Telegram bot integration. Reuses the same
|
||||
// credential lookup pattern as cmd/recon.go buildReconEngine.
|
||||
func buildServeReconEngine() *recon.Engine {
|
||||
eng := recon.NewEngine()
|
||||
reg, err := providers.NewRegistry()
|
||||
if err != nil {
|
||||
log.Printf("serve: failed to load providers for recon: %v", err)
|
||||
return eng
|
||||
}
|
||||
reconSources.RegisterAll(eng, reconSources.SourcesConfig{
|
||||
Registry: reg,
|
||||
Limiters: recon.NewLimiterRegistry(),
|
||||
GitHubToken: firstNonEmpty(os.Getenv("GITHUB_TOKEN"), viper.GetString("recon.github.token")),
|
||||
GitLabToken: firstNonEmpty(os.Getenv("GITLAB_TOKEN"), viper.GetString("recon.gitlab.token")),
|
||||
BitbucketToken: firstNonEmpty(os.Getenv("BITBUCKET_TOKEN"), viper.GetString("recon.bitbucket.token")),
|
||||
BitbucketWorkspace: firstNonEmpty(os.Getenv("BITBUCKET_WORKSPACE"), viper.GetString("recon.bitbucket.workspace")),
|
||||
CodebergToken: firstNonEmpty(os.Getenv("CODEBERG_TOKEN"), viper.GetString("recon.codeberg.token")),
|
||||
HuggingFaceToken: firstNonEmpty(os.Getenv("HUGGINGFACE_TOKEN"), viper.GetString("recon.huggingface.token")),
|
||||
KaggleUser: firstNonEmpty(os.Getenv("KAGGLE_USERNAME"), viper.GetString("recon.kaggle.username")),
|
||||
KaggleKey: firstNonEmpty(os.Getenv("KAGGLE_KEY"), viper.GetString("recon.kaggle.key")),
|
||||
GoogleAPIKey: firstNonEmpty(os.Getenv("GOOGLE_API_KEY"), viper.GetString("recon.google.api_key")),
|
||||
GoogleCX: firstNonEmpty(os.Getenv("GOOGLE_CX"), viper.GetString("recon.google.cx")),
|
||||
BingAPIKey: firstNonEmpty(os.Getenv("BING_API_KEY"), viper.GetString("recon.bing.api_key")),
|
||||
YandexUser: firstNonEmpty(os.Getenv("YANDEX_USER"), viper.GetString("recon.yandex.user")),
|
||||
YandexAPIKey: firstNonEmpty(os.Getenv("YANDEX_API_KEY"), viper.GetString("recon.yandex.api_key")),
|
||||
BraveAPIKey: firstNonEmpty(os.Getenv("BRAVE_API_KEY"), viper.GetString("recon.brave.api_key")),
|
||||
ShodanAPIKey: firstNonEmpty(os.Getenv("SHODAN_API_KEY"), viper.GetString("recon.shodan.api_key")),
|
||||
CensysAPIId: firstNonEmpty(os.Getenv("CENSYS_API_ID"), viper.GetString("recon.censys.api_id")),
|
||||
CensysAPISecret: firstNonEmpty(os.Getenv("CENSYS_API_SECRET"), viper.GetString("recon.censys.api_secret")),
|
||||
ZoomEyeAPIKey: firstNonEmpty(os.Getenv("ZOOMEYE_API_KEY"), viper.GetString("recon.zoomeye.api_key")),
|
||||
FOFAEmail: firstNonEmpty(os.Getenv("FOFA_EMAIL"), viper.GetString("recon.fofa.email")),
|
||||
FOFAAPIKey: firstNonEmpty(os.Getenv("FOFA_API_KEY"), viper.GetString("recon.fofa.api_key")),
|
||||
NetlasAPIKey: firstNonEmpty(os.Getenv("NETLAS_API_KEY"), viper.GetString("recon.netlas.api_key")),
|
||||
BinaryEdgeAPIKey: firstNonEmpty(os.Getenv("BINARYEDGE_API_KEY"), viper.GetString("recon.binaryedge.api_key")),
|
||||
CircleCIToken: firstNonEmpty(os.Getenv("CIRCLECI_TOKEN"), viper.GetString("recon.circleci.token")),
|
||||
VirusTotalAPIKey: firstNonEmpty(os.Getenv("VIRUSTOTAL_API_KEY"), viper.GetString("recon.virustotal.api_key")),
|
||||
IntelligenceXAPIKey: firstNonEmpty(os.Getenv("INTELX_API_KEY"), viper.GetString("recon.intelx.api_key")),
|
||||
SecurityTrailsAPIKey: firstNonEmpty(os.Getenv("SECURITYTRAILS_API_KEY"), viper.GetString("recon.securitytrails.api_key")),
|
||||
})
|
||||
return eng
|
||||
}
|
||||
|
||||
func init() {
|
||||
serveCmd.Flags().BoolVar(&serveTelegram, "telegram", false, "start Telegram bot alongside scheduler")
|
||||
serveCmd.Flags().IntVar(&servePort, "port", 8080, "port for future web dashboard (reserved)")
|
||||
}
|
||||
12
cmd/stubs.go
12
cmd/stubs.go
@@ -25,16 +25,8 @@ var verifyCmd = &cobra.Command{
|
||||
|
||||
// keysCmd is implemented in cmd/keys.go (Phase 6).
|
||||
|
||||
var serveCmd = &cobra.Command{
|
||||
Use: "serve",
|
||||
Short: "Start the web dashboard (Phase 18)",
|
||||
RunE: notImplemented("serve", "Phase 18"),
|
||||
}
|
||||
// serveCmd is implemented in cmd/serve.go (Phase 17).
|
||||
|
||||
// dorksCmd is implemented in cmd/dorks.go (Phase 8).
|
||||
|
||||
var scheduleCmd = &cobra.Command{
|
||||
Use: "schedule",
|
||||
Short: "Manage scheduled recurring scans (Phase 17)",
|
||||
RunE: notImplemented("schedule", "Phase 17"),
|
||||
}
|
||||
// scheduleCmd is implemented in cmd/schedule.go (Phase 17).
|
||||
|
||||
Reference in New Issue
Block a user