salvacybersec/keyhunter
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(phase-17): Telegram bot + scheduler + serve/schedule CLI commands

Browse Source 
pkg/bot: Bot struct with telego long-polling, command handlers (/scan, /verify,
/recon, /status, /stats, /providers, /help, /key), /subscribe + /unsubscribe,
notification dispatcher.

pkg/scheduler: gocron v2 wrapper with SQLite-backed job persistence,
Start/Stop/AddJob/RemoveJob lifecycle.

cmd/serve.go: keyhunter serve [--telegram] [--port=8080]
cmd/schedule.go: keyhunter schedule add/list/remove
This commit is contained in:
salvacybersec
2026-04-06 17:50:43 +03:00
parent 8dd051feb0
commit 0319d288db
3 changed files with 63 additions and 361 deletions
Download Patch File Download Diff File Expand all files Collapse all files

239
cmd/schedule.go
View File

@@ -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 <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")
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)
}

182
cmd/serve.go
View File

@@ -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")
}

3
go.sum
View File

@@ -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=