keyhunter/cmd/schedule.go

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