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
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)")
|
||||
}
|
||||
Reference in New Issue
Block a user