keyhunter/pkg/bot/bot.go

// Package bot implements the Telegram bot interface for KeyHunter.
// It wraps telego v1.8.0 with long-polling updates, per-chat authorization,
// per-user rate limiting, and command dispatch to handler stubs.
package bot

import (
	"context"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/mymmrac/telego"
	"github.com/mymmrac/telego/telegoutil"
	"github.com/salvacybersec/keyhunter/pkg/engine"
	"github.com/salvacybersec/keyhunter/pkg/providers"
	"github.com/salvacybersec/keyhunter/pkg/recon"
	"github.com/salvacybersec/keyhunter/pkg/storage"
)

// Config holds all dependencies and settings for the Telegram bot.
type Config struct {
	// Token is the Telegram bot token from BotFather.
	Token string

	// AllowedChats restricts bot access to these chat IDs.
	// Empty slice means allow all chats.
	AllowedChats []int64

	// DB is the SQLite database for subscriber queries and finding lookups.
	DB *storage.DB

	// ScanEngine is the scanning engine for /scan commands.
	ScanEngine *engine.Engine

	// ReconEngine is the recon engine for /recon commands.
	ReconEngine *recon.Engine

	// ProviderRegistry is the provider registry for /providers and /verify.
	ProviderRegistry *providers.Registry

	// EncKey is the encryption key for finding decryption.
	EncKey []byte
}

// Bot wraps a telego.Bot with KeyHunter command handling and authorization.
type Bot struct {
	cfg       Config
	bot       *telego.Bot
	cancel    context.CancelFunc
	startTime time.Time

	rateMu     sync.Mutex
	rateLimits map[int64]time.Time
}

// commands is the list of bot commands registered with Telegram.
var commands = []telego.BotCommand{
	{Command: "scan", Description: "Scan a target for API keys"},
	{Command: "verify", Description: "Verify a found API key"},
	{Command: "recon", Description: "Run OSINT recon for a keyword"},
	{Command: "status", Description: "Show bot and scan status"},
	{Command: "stats", Description: "Show finding statistics"},
	{Command: "providers", Description: "List supported providers"},
	{Command: "help", Description: "Show available commands"},
	{Command: "key", Description: "Show full details for a finding"},
	{Command: "subscribe", Description: "Subscribe to scan notifications"},
	{Command: "unsubscribe", Description: "Unsubscribe from notifications"},
}

// New creates a new Bot from the given config. Returns an error if the token
// is invalid or telego cannot initialize.
func New(cfg Config) (*Bot, error) {
	tb, err := telego.NewBot(cfg.Token)
	if err != nil {
		return nil, fmt.Errorf("creating telego bot: %w", err)
	}
	return &Bot{
		cfg:        cfg,
		bot:        tb,
		rateLimits: make(map[int64]time.Time),
	}, nil
}

// Start begins long-polling for updates and dispatching commands. It blocks
// until the provided context is cancelled or an error occurs.
func (b *Bot) Start(ctx context.Context) error {
	ctx, b.cancel = context.WithCancel(ctx)

	// Register command list with Telegram.
	err := b.bot.SetMyCommands(ctx, &telego.SetMyCommandsParams{
		Commands: commands,
	})
	if err != nil {
		return fmt.Errorf("setting bot commands: %w", err)
	}

	updates, err := b.bot.UpdatesViaLongPolling(ctx, nil)
	if err != nil {
		return fmt.Errorf("starting long polling: %w", err)
	}

	for update := range updates {
		if update.Message == nil {
			continue
		}
		b.dispatch(ctx, update.Message)
	}

	return nil
}

// Stop cancels the bot context, which stops long polling and the update loop.
func (b *Bot) Stop() {
	if b.cancel != nil {
		b.cancel()
	}
}

// isAllowed returns true if the given chat ID is authorized to use the bot.
// If AllowedChats is empty, all chats are allowed.
func (b *Bot) isAllowed(chatID int64) bool {
	if len(b.cfg.AllowedChats) == 0 {
		return true
	}
	for _, id := range b.cfg.AllowedChats {
		if id == chatID {
			return true
		}
	}
	return false
}

// checkRateLimit returns true if the user is allowed to execute a command,
// false if they are still within the cooldown window.
func (b *Bot) checkRateLimit(userID int64, cooldown time.Duration) bool {
	b.rateMu.Lock()
	defer b.rateMu.Unlock()

	last, ok := b.rateLimits[userID]
	if ok && time.Since(last) < cooldown {
		return false
	}
	b.rateLimits[userID] = time.Now()
	return true
}

// dispatch routes an incoming message to the appropriate handler.
func (b *Bot) dispatch(ctx context.Context, msg *telego.Message) {
	chatID := msg.Chat.ID
	if !b.isAllowed(chatID) {
		_ = b.replyPlain(ctx, chatID, "Unauthorized: your chat ID is not in the allowed list.")
		return
	}

	text := strings.TrimSpace(msg.Text)
	if text == "" {
		return
	}

	// Extract command (first word, with optional @mention suffix removed).
	cmd := strings.SplitN(text, " ", 2)[0]
	if at := strings.Index(cmd, "@"); at > 0 {
		cmd = cmd[:at]
	}

	// Determine cooldown based on command type.
	var cooldown time.Duration
	switch cmd {
	case "/scan", "/verify", "/recon":
		cooldown = 60 * time.Second
	default:
		cooldown = 5 * time.Second
	}

	if msg.From != nil && !b.checkRateLimit(msg.From.ID, cooldown) {
		_ = b.replyPlain(ctx, chatID, "Rate limited. Please wait before sending another command.")
		return
	}

	switch cmd {
	case "/scan":
		b.handleScan(ctx, msg)
	case "/verify":
		b.handleVerify(ctx, msg)
	case "/recon":
		b.handleRecon(ctx, msg)
	case "/status":
		b.handleStatus(ctx, msg)
	case "/stats":
		b.handleStats(ctx, msg)
	case "/providers":
		b.handleProviders(ctx, msg)
	case "/help", "/start":
		b.handleHelp(ctx, msg)
	case "/key":
		b.handleKey(ctx, msg)
	case "/subscribe":
		b.handleSubscribe(ctx, msg)
	case "/unsubscribe":
		b.handleUnsubscribe(ctx, msg)
	}
}

// reply sends a MarkdownV2-formatted message to the given chat.
func (b *Bot) reply(ctx context.Context, chatID int64, text string) error {
	params := telegoutil.Message(telego.ChatID{ID: chatID}, text).
		WithParseMode("MarkdownV2")
	_, err := b.bot.SendMessage(ctx, params)
	return err
}

// replyPlain sends a plain text message to the given chat.
func (b *Bot) replyPlain(ctx context.Context, chatID int64, text string) error {
	params := telegoutil.Message(telego.ChatID{ID: chatID}, text)
	_, err := b.bot.SendMessage(ctx, params)
	return err
}

// Command handlers are in handlers.go (17-03).
// Subscribe/unsubscribe handlers are in subscribe.go (17-04).
// Notification dispatcher is in notify.go (17-04).