keyhunter/pkg/bot/handlers.go

package bot

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/mymmrac/telego"
	th "github.com/mymmrac/telego/telegohandler"

	"github.com/salvacybersec/keyhunter/pkg/engine"
	"github.com/salvacybersec/keyhunter/pkg/recon"
	"github.com/salvacybersec/keyhunter/pkg/storage"
)

// commandHelp lists all available bot commands with descriptions.
var commandHelp = []struct {
	Cmd  string
	Desc string
}{
	{"/help", "Show this help message"},
	{"/scan <path>", "Scan a file or directory for leaked API keys"},
	{"/verify <id>", "Verify a stored finding by ID"},
	{"/recon [--sources=x,y]", "Run OSINT recon sweep across sources"},
	{"/status", "Show bot status and uptime"},
	{"/stats", "Show provider and finding statistics"},
	{"/providers", "List loaded provider definitions"},
	{"/key <id>", "Show full key detail (private chat only)"},
}

// handleHelp responds with a list of all available commands.
func (b *Bot) handleHelp(ctx *th.Context, msg telego.Message) error {
	var sb strings.Builder
	sb.WriteString("KeyHunter Bot Commands:\n\n")
	for _, c := range commandHelp {
		sb.WriteString(fmt.Sprintf("%s - %s\n", c.Cmd, c.Desc))
	}
	b.reply(ctx, &msg, sb.String())
	return nil
}

// handleScan triggers a scan on the given path and returns masked findings.
// Usage: /scan <path>
func (b *Bot) handleScan(ctx *th.Context, msg telego.Message) error {
	path := extractArg(msg.Text)
	if path == "" {
		b.reply(ctx, &msg, "Usage: /scan <path>\nExample: /scan /tmp/myrepo")
		return nil
	}

	b.reply(ctx, &msg, fmt.Sprintf("Scanning %s ...", path))

	scanCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
	defer cancel()

	findings, err := b.runScan(scanCtx, path)
	if err != nil {
		b.reply(ctx, &msg, fmt.Sprintf("Scan error: %s", err))
		return nil
	}

	if len(findings) == 0 {
		b.reply(ctx, &msg, "Scan complete. No API keys found.")
		return nil
	}

	var sb strings.Builder
	sb.WriteString(fmt.Sprintf("Scan complete. Found %d key(s):\n\n", len(findings)))
	for i, f := range findings {
		if i >= 20 {
			sb.WriteString(fmt.Sprintf("\n... and %d more", len(findings)-20))
			break
		}
		sb.WriteString(fmt.Sprintf("[%s] %s %s:%d\n", f.ProviderName, f.KeyMasked, f.Source, f.LineNumber))
	}
	b.reply(ctx, &msg, sb.String())
	return nil
}

// handleVerify verifies a stored finding by its database ID.
// Usage: /verify <id>
func (b *Bot) handleVerify(ctx *th.Context, msg telego.Message) error {
	arg := extractArg(msg.Text)
	if arg == "" {
		b.reply(ctx, &msg, "Usage: /verify <id>\nExample: /verify 42")
		return nil
	}

	id, err := strconv.ParseInt(arg, 10, 64)
	if err != nil || id <= 0 {
		b.reply(ctx, &msg, "Invalid ID. Must be a positive integer.")
		return nil
	}

	f, err := b.db.GetFinding(id, b.encKey)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			b.reply(ctx, &msg, fmt.Sprintf("No finding with ID %d.", id))
			return nil
		}
		b.reply(ctx, &msg, fmt.Sprintf("Error: %s", err))
		return nil
	}

	ef := storageToEngine(*f)
	results := b.verifier.VerifyAll(ctx, []engine.Finding{ef}, b.registry, 1)
	r := <-results
	for range results {
	}

	// Persist verification result.
	var metaJSON interface{}
	if r.Metadata != nil {
		byt, _ := json.Marshal(r.Metadata)
		metaJSON = string(byt)
	} else {
		metaJSON = sql.NullString{}
	}
	_, _ = b.db.SQL().Exec(
		`UPDATE findings SET verified=1, verify_status=?, verify_http_code=?, verify_metadata_json=? WHERE id=?`,
		r.Status, r.HTTPCode, metaJSON, id,
	)

	var sb strings.Builder
	sb.WriteString(fmt.Sprintf("Verification for finding #%d:\n", id))
	sb.WriteString(fmt.Sprintf("Provider: %s\n", f.ProviderName))
	sb.WriteString(fmt.Sprintf("Key: %s\n", f.KeyMasked))
	sb.WriteString(fmt.Sprintf("Status: %s\n", r.Status))
	if r.HTTPCode != 0 {
		sb.WriteString(fmt.Sprintf("HTTP Code: %d\n", r.HTTPCode))
	}
	if r.Error != "" {
		sb.WriteString(fmt.Sprintf("Error: %s\n", r.Error))
	}
	if len(r.Metadata) > 0 {
		sb.WriteString("Metadata:\n")
		for k, v := range r.Metadata {
			sb.WriteString(fmt.Sprintf("  %s: %s\n", k, v))
		}
	}
	b.reply(ctx, &msg, sb.String())
	return nil
}

// handleRecon runs a recon sweep and returns a summary.
// Usage: /recon [--sources=github,gitlab]
func (b *Bot) handleRecon(ctx *th.Context, msg telego.Message) error {
	arg := extractArg(msg.Text)

	b.reply(ctx, &msg, "Running recon sweep...")

	cfg := recon.Config{
		RespectRobots: true,
	}

	eng := b.recon
	// Parse optional --sources filter
	if strings.HasPrefix(arg, "--sources=") {
		filter := strings.TrimPrefix(arg, "--sources=")
		names := strings.Split(filter, ",")
		filtered := recon.NewEngine()
		for _, name := range names {
			name = strings.TrimSpace(name)
			if src, ok := eng.Get(name); ok {
				filtered.Register(src)
			}
		}
		eng = filtered
	}

	reconCtx, cancel := context.WithTimeout(ctx, 10*time.Minute)
	defer cancel()

	all, err := eng.SweepAll(reconCtx, cfg)
	if err != nil {
		b.reply(ctx, &msg, fmt.Sprintf("Recon error: %s", err))
		return nil
	}

	deduped := recon.Dedup(all)

	if len(deduped) == 0 {
		b.reply(ctx, &msg, fmt.Sprintf("Recon complete. Swept %d sources, no findings.", len(eng.List())))
		return nil
	}

	var sb strings.Builder
	sb.WriteString(fmt.Sprintf("Recon complete: %d sources, %d findings (%d after dedup)\n\n",
		len(eng.List()), len(all), len(deduped)))
	for i, f := range deduped {
		if i >= 20 {
			sb.WriteString(fmt.Sprintf("\n... and %d more", len(deduped)-20))
			break
		}
		sb.WriteString(fmt.Sprintf("[%s] %s %s %s\n", f.SourceType, f.ProviderName, f.KeyMasked, f.Source))
	}
	b.reply(ctx, &msg, sb.String())
	return nil
}

// handleStatus returns bot uptime and last scan time.
func (b *Bot) handleStatus(ctx *th.Context, msg telego.Message) error {
	b.mu.Lock()
	startedAt := b.startedAt
	lastScan := b.lastScan
	b.mu.Unlock()

	uptime := time.Since(startedAt).Truncate(time.Second)

	var sb strings.Builder
	sb.WriteString("KeyHunter Bot Status\n\n")
	sb.WriteString(fmt.Sprintf("Uptime: %s\n", uptime))
	sb.WriteString(fmt.Sprintf("Started: %s\n", startedAt.Format(time.RFC3339)))
	if !lastScan.IsZero() {
		sb.WriteString(fmt.Sprintf("Last scan: %s\n", lastScan.Format(time.RFC3339)))
	} else {
		sb.WriteString("Last scan: none\n")
	}

	// DB stats
	findings, err := b.db.ListFindingsFiltered(b.encKey, storage.Filters{Limit: 0})
	if err == nil {
		sb.WriteString(fmt.Sprintf("Total findings: %d\n", len(findings)))
	}

	sb.WriteString(fmt.Sprintf("Providers loaded: %d\n", len(b.registry.List())))
	sb.WriteString(fmt.Sprintf("Recon sources: %d\n", len(b.recon.List())))

	b.reply(ctx, &msg, sb.String())
	return nil
}

// handleStats returns provider and finding statistics.
func (b *Bot) handleStats(ctx *th.Context, msg telego.Message) error {
	stats := b.registry.Stats()

	var sb strings.Builder
	sb.WriteString("KeyHunter Statistics\n\n")
	sb.WriteString(fmt.Sprintf("Total providers: %d\n", stats.Total))
	sb.WriteString("By tier:\n")
	for tier := 1; tier <= 9; tier++ {
		if count := stats.ByTier[tier]; count > 0 {
			sb.WriteString(fmt.Sprintf("  Tier %d: %d\n", tier, count))
		}
	}
	sb.WriteString("By confidence:\n")
	for conf, count := range stats.ByConfidence {
		sb.WriteString(fmt.Sprintf("  %s: %d\n", conf, count))
	}

	// Finding counts
	findings, err := b.db.ListFindingsFiltered(b.encKey, storage.Filters{Limit: 0})
	if err == nil {
		verified := 0
		for _, f := range findings {
			if f.Verified {
				verified++
			}
		}
		sb.WriteString(fmt.Sprintf("\nTotal findings: %d\n", len(findings)))
		sb.WriteString(fmt.Sprintf("Verified: %d\n", verified))
		sb.WriteString(fmt.Sprintf("Unverified: %d\n", len(findings)-verified))
	}

	b.reply(ctx, &msg, sb.String())
	return nil
}

// handleProviders lists all loaded provider definitions.
func (b *Bot) handleProviders(ctx *th.Context, msg telego.Message) error {
	provs := b.registry.List()

	var sb strings.Builder
	sb.WriteString(fmt.Sprintf("Loaded providers (%d):\n\n", len(provs)))
	for i, p := range provs {
		if i >= 50 {
			sb.WriteString(fmt.Sprintf("\n... and %d more", len(provs)-50))
			break
		}
		sb.WriteString(fmt.Sprintf("%-20s tier=%d patterns=%d\n", p.Name, p.Tier, len(p.Patterns)))
	}
	b.reply(ctx, &msg, sb.String())
	return nil
}

// handleKey shows full key detail. Only works in private chats for security.
// Usage: /key <id>
func (b *Bot) handleKey(ctx *th.Context, msg telego.Message) error {
	if !isPrivateChat(&msg) {
		b.reply(ctx, &msg, "The /key command is only available in private chats for security reasons.")
		return nil
	}

	arg := extractArg(msg.Text)
	if arg == "" {
		b.reply(ctx, &msg, "Usage: /key <id>\nExample: /key 42")
		return nil
	}

	id, err := strconv.ParseInt(arg, 10, 64)
	if err != nil || id <= 0 {
		b.reply(ctx, &msg, "Invalid ID. Must be a positive integer.")
		return nil
	}

	f, err := b.db.GetFinding(id, b.encKey)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			b.reply(ctx, &msg, fmt.Sprintf("No finding with ID %d.", id))
			return nil
		}
		b.reply(ctx, &msg, fmt.Sprintf("Error: %s", err))
		return nil
	}

	var sb strings.Builder
	sb.WriteString(fmt.Sprintf("Finding #%d\n\n", f.ID))
	sb.WriteString(fmt.Sprintf("Provider:    %s\n", f.ProviderName))
	sb.WriteString(fmt.Sprintf("Confidence:  %s\n", f.Confidence))
	sb.WriteString(fmt.Sprintf("Key:         %s\n", f.KeyValue))
	sb.WriteString(fmt.Sprintf("Key Masked:  %s\n", f.KeyMasked))
	sb.WriteString(fmt.Sprintf("Source:      %s\n", f.SourcePath))
	sb.WriteString(fmt.Sprintf("Source Type: %s\n", f.SourceType))
	sb.WriteString(fmt.Sprintf("Line:        %d\n", f.LineNumber))
	if !f.CreatedAt.IsZero() {
		sb.WriteString(fmt.Sprintf("Created:     %s\n", f.CreatedAt.Format(time.RFC3339)))
	}
	sb.WriteString(fmt.Sprintf("Verified:    %t\n", f.Verified))
	if f.VerifyStatus != "" {
		sb.WriteString(fmt.Sprintf("Status:      %s\n", f.VerifyStatus))
	}
	if f.VerifyHTTPCode != 0 {
		sb.WriteString(fmt.Sprintf("HTTP Code:   %d\n", f.VerifyHTTPCode))
	}
	if len(f.VerifyMetadata) > 0 {
		sb.WriteString("Metadata:\n")
		for k, v := range f.VerifyMetadata {
			sb.WriteString(fmt.Sprintf("  %s: %s\n", k, v))
		}
	}
	b.reply(ctx, &msg, sb.String())
	return nil
}

// extractArg extracts the argument after the /command from message text.
// For "/scan /tmp/repo", returns "/tmp/repo".
// For "/help", returns "".
func extractArg(text string) string {
	parts := strings.SplitN(text, " ", 2)
	if len(parts) < 2 {
		return ""
	}
	return strings.TrimSpace(parts[1])
}

// storageToEngine converts a storage.Finding to an engine.Finding for verification.
func storageToEngine(f storage.Finding) engine.Finding {
	return engine.Finding{
		ProviderName:   f.ProviderName,
		KeyValue:       f.KeyValue,
		KeyMasked:      f.KeyMasked,
		Confidence:     f.Confidence,
		Source:         f.SourcePath,
		SourceType:     f.SourceType,
		LineNumber:     f.LineNumber,
		DetectedAt:     f.CreatedAt,
		Verified:       f.Verified,
		VerifyStatus:   f.VerifyStatus,
		VerifyHTTPCode: f.VerifyHTTPCode,
		VerifyMetadata: f.VerifyMetadata,
	}
}