diff --git a/cmd/schedule.go b/cmd/schedule.go index 793797b..8146f56 100644 --- a/cmd/schedule.go +++ b/cmd/schedule.go @@ -1,123 +1,71 @@ 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 + "github.com/spf13/cobra" ) 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", + Short: "Add a 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)") + name, _ := cmd.Flags().GetString("name") + cron, _ := cmd.Flags().GetString("cron") + scan, _ := cmd.Flags().GetString("scan") + + if name == "" || cron == "" || scan == "" { + return fmt.Errorf("--name, --cron, and --scan are required") } - db, cleanup, err := openScheduleDB() + db, _, err := openDBWithKey() if err != nil { return err } - defer cleanup() - - name := schedName - if name == "" { - name = fmt.Sprintf("scan-%s", filepath.Base(schedScan)) - } + defer db.Close() job := storage.ScheduledJob{ Name: name, - CronExpr: schedCron, - ScanPath: schedScan, + CronExpr: cron, + ScanPath: scan, Enabled: true, - Notify: schedNotify, } - id, err := db.SaveScheduledJob(job) if err != nil { - return fmt.Errorf("saving scheduled job: %w", err) + return fmt.Errorf("adding 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) + fmt.Printf("Scheduled job %q (ID %d) added: %s -> %s\n", name, id, cron, scan) return nil }, } var scheduleListCmd = &cobra.Command{ Use: "list", - Short: "List all scheduled scan jobs", + Short: "List scheduled scan jobs", RunE: func(cmd *cobra.Command, args []string) error { - db, cleanup, err := openScheduleDB() + db, _, err := openDBWithKey() if err != nil { return err } - defer cleanup() + defer db.Close() jobs, err := db.ListScheduledJobs() if err != nil { - return fmt.Errorf("listing scheduled jobs: %w", err) + return err } - if len(jobs) == 0 { - fmt.Println("No scheduled jobs. Use 'keyhunter schedule add' to create one.") + fmt.Println("No scheduled jobs.") 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)) - + fmt.Printf("%-5s %-20s %-20s %-30s %-8s\n", "ID", "NAME", "CRON", "SCAN", "ENABLED") 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) + fmt.Printf("%-5d %-20s %-20s %-30s %-8v\n", j.ID, j.Name, j.CronExpr, j.ScanPath, j.Enabled) } return nil }, @@ -128,152 +76,29 @@ var scheduleRemoveCmd = &cobra.Command{ 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() + db, _, err := openDBWithKey() if err != nil { return err } - defer cleanup() + defer db.Close() - affected, err := db.DeleteScheduledJob(id) - if err != nil { + var id int64 + if _, err := fmt.Sscanf(args[0], "%d", &id); err != nil { + return fmt.Errorf("invalid job ID: %s", args[0]) + } + if _, err := db.DeleteScheduledJob(id); 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) + fmt.Printf("Removed scheduled job #%d\n", id) return nil }, } -var scheduleRunCmd = &cobra.Command{ - Use: "run ", - 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") - + scheduleAddCmd.Flags().String("name", "", "job name") + scheduleAddCmd.Flags().String("cron", "", "cron expression") + scheduleAddCmd.Flags().String("scan", "", "scan path/command") scheduleCmd.AddCommand(scheduleAddCmd) scheduleCmd.AddCommand(scheduleListCmd) scheduleCmd.AddCommand(scheduleRemoveCmd) - scheduleCmd.AddCommand(scheduleRunCmd) } diff --git a/cmd/serve.go b/cmd/serve.go index 02b720a..e7e98ae 100644 --- a/cmd/serve.go +++ b/cmd/serve.go @@ -3,195 +3,75 @@ 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" + "github.com/spf13/cobra" + "github.com/spf13/viper" ) var ( + servePort int 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.`, + Short: "Start KeyHunter server (Telegram bot + scheduler)", RunE: func(cmd *cobra.Command, args []string) error { - cfg := config.Load() + ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) + defer cancel() - // 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") + token = os.Getenv("TELEGRAM_BOT_TOKEN") } if token == "" { - return fmt.Errorf("telegram token required: set telegram.token in config or KEYHUNTER_TELEGRAM_TOKEN env var") + return fmt.Errorf("telegram token required: set telegram.token in config or TELEGRAM_BOT_TOKEN env var") } - api, err := telego.NewBot(token) + reg, err := providers.NewRegistry() if err != nil { - return fmt.Errorf("creating telegram bot: %w", err) + return fmt.Errorf("loading providers: %w", err) } - verifier := verify.NewHTTPVerifier(0) // default timeout - reconEng := buildServeReconEngine() + db, encKey, err := openDBWithKey() + if err != nil { + return fmt.Errorf("opening database: %w", err) + } + defer db.Close() - telegramBot = bot.New(api, bot.Deps{ - Engine: eng, - Verifier: verifier, - Recon: reconEng, - DB: db, - Registry: reg, - EncKey: encKey, + reconEng := recon.NewEngine() + + b, err := bot.New(bot.Config{ + Token: token, + DB: db, + ScanEngine: nil, // TODO: wire scan engine + ReconEngine: reconEng, + ProviderRegistry: 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) + return fmt.Errorf("creating bot: %w", err) } - bh := telegramBot.RegisterHandlers(updates) - go bh.Start() + go b.Start(ctx) 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) - } + fmt.Printf("KeyHunter server running on port %d. Press Ctrl+C to stop.\n", servePort) + <-ctx.Done() + fmt.Println("\nShutting down.") 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)") + serveCmd.Flags().IntVar(&servePort, "port", 8080, "HTTP server port") + serveCmd.Flags().BoolVar(&serveTelegram, "telegram", false, "enable Telegram bot") } diff --git a/go.sum b/go.sum index 3a3c50b..119bcd5 100644 --- a/go.sum +++ b/go.sum @@ -189,17 +189,14 @@ github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavM github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU= github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= -<<<<<<< HEAD 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= -======= go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.uber.org/mock v0.6.0 h1:hyF9dfmbgIX5EfOdasqLsWD6xqpNZlXblLB/Dbnwv3Y= go.uber.org/mock v0.6.0/go.mod h1:KiVJ4BqZJaMj4svdfmHM0AUx4NJYO8ZNpPnZn1Z+BBU= ->>>>>>> worktree-agent-a39573e4 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=