keyhunter/cmd/serve.go

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