merge: phase 17 wave 2

pkg/bot/bot.go

@@ -1,140 +1,253 @@
 // Package bot implements the Telegram bot interface for KeyHunter.
-// It wraps existing scan, verify, recon, and storage functionality,
-// exposing them through Telegram command handlers via the telego library.
+// 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"
-	th "github.com/mymmrac/telego/telegohandler"
-
+	"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"
-	"github.com/salvacybersec/keyhunter/pkg/verify"
 )
 
-// Bot holds the Telegram bot instance and all dependencies needed
-// to process commands. It delegates to the existing KeyHunter engine,
-// verifier, recon engine, and storage layer.
+// 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 {
-	api      *telego.Bot
-	handler  *th.BotHandler
-	engine   *engine.Engine
-	verifier *verify.HTTPVerifier
-	recon    *recon.Engine
-	db       *storage.DB
-	registry *providers.Registry
-	encKey   []byte
+	cfg    Config
+	bot    *telego.Bot
+	cancel context.CancelFunc
 
-	mu        sync.Mutex
-	startedAt time.Time
-	lastScan  time.Time
+	rateMu     sync.Mutex
+	rateLimits map[int64]time.Time
 }
 
-// Deps bundles the dependencies required to construct a Bot.
-type Deps struct {
-	Engine   *engine.Engine
-	Verifier *verify.HTTPVerifier
-	Recon    *recon.Engine
-	DB       *storage.DB
-	Registry *providers.Registry
-	EncKey   []byte
+// 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 Bot backed by the given telego API client and dependencies.
-// Call RegisterHandlers to wire up command handlers before starting the update loop.
-func New(api *telego.Bot, deps Deps) *Bot {
+// 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{
-		api:       api,
-		engine:    deps.Engine,
-		verifier:  deps.Verifier,
-		recon:     deps.Recon,
-		db:        deps.DB,
-		registry:  deps.Registry,
-		encKey:    deps.EncKey,
-		startedAt: time.Now(),
-	}
+		cfg:        cfg,
+		bot:        tb,
+		rateLimits: make(map[int64]time.Time),
+	}, nil
 }
 
-// RegisterHandlers wires all command handlers into a BotHandler that processes
-// updates from the Telegram API. The caller must call Start() on the returned
-// BotHandler to begin processing.
-func (b *Bot) RegisterHandlers(updates <-chan telego.Update) *th.BotHandler {
-	bh, _ := th.NewBotHandler(b.api, updates)
+// 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)
 
-	bh.HandleMessage(b.handleHelp, th.CommandEqual("help"))
-	bh.HandleMessage(b.handleScan, th.CommandEqual("scan"))
-	bh.HandleMessage(b.handleVerify, th.CommandEqual("verify"))
-	bh.HandleMessage(b.handleRecon, th.CommandEqual("recon"))
-	bh.HandleMessage(b.handleStatus, th.CommandEqual("status"))
-	bh.HandleMessage(b.handleStats, th.CommandEqual("stats"))
-	bh.HandleMessage(b.handleProviders, th.CommandEqual("providers"))
-	bh.HandleMessage(b.handleKey, th.CommandEqual("key"))
-
-	b.handler = bh
-	return bh
-}
-
-// reply sends a text message back to the chat that originated msg.
-func (b *Bot) reply(ctx context.Context, msg *telego.Message, text string) {
-	_, _ = b.api.SendMessage(ctx, &telego.SendMessageParams{
-		ChatID: telego.ChatID{ID: msg.Chat.ID},
-		Text:   text,
+	// Register command list with Telegram.
+	err := b.bot.SetMyCommands(ctx, &telego.SetMyCommandsParams{
+		Commands: commands,
 	})
-}
-
-// isPrivateChat returns true if the message was sent in a private (1:1) chat.
-func isPrivateChat(msg *telego.Message) bool {
-	return msg.Chat.Type == "private"
-}
-
-// runScan executes a scan against the given path and returns findings.
-// Findings are collected synchronously; the caller formats the output.
-func (b *Bot) runScan(ctx context.Context, path string) ([]engine.Finding, error) {
-	src, err := selectBotSource(path)
 	if err != nil {
-		return nil, err
+		return fmt.Errorf("setting bot commands: %w", err)
 	}
 
-	cfg := engine.ScanConfig{
-		Workers: 0, // auto
-		Verify:  false,
-		Unmask:  false,
-	}
-
-	ch, err := b.engine.Scan(ctx, src, cfg)
+	updates, err := b.bot.UpdatesViaLongPolling(ctx, nil)
 	if err != nil {
-		return nil, fmt.Errorf("starting scan: %w", err)
+		return fmt.Errorf("starting long polling: %w", err)
 	}
 
-	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,
+	for update := range updates {
+		if update.Message == nil {
+			continue
 		}
-		_, _ = b.db.SaveFinding(sf, b.encKey)
+		b.dispatch(ctx, update.Message)
 	}
 
-	b.mu.Lock()
-	b.lastScan = time.Now()
-	b.mu.Unlock()
-
-	return findings, nil
+	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
+}
+
+// --- Handler stubs (implemented in Plan 17-03/17-04) ---
+
+func (b *Bot) handleScan(ctx context.Context, msg *telego.Message) {
+	_ = b.replyPlain(ctx, msg.Chat.ID, "Not yet implemented: /scan")
+}
+
+func (b *Bot) handleVerify(ctx context.Context, msg *telego.Message) {
+	_ = b.replyPlain(ctx, msg.Chat.ID, "Not yet implemented: /verify")
+}
+
+func (b *Bot) handleRecon(ctx context.Context, msg *telego.Message) {
+	_ = b.replyPlain(ctx, msg.Chat.ID, "Not yet implemented: /recon")
+}
+
+func (b *Bot) handleStatus(ctx context.Context, msg *telego.Message) {
+	_ = b.replyPlain(ctx, msg.Chat.ID, "Not yet implemented: /status")
+}
+
+func (b *Bot) handleStats(ctx context.Context, msg *telego.Message) {
+	_ = b.replyPlain(ctx, msg.Chat.ID, "Not yet implemented: /stats")
+}
+
+func (b *Bot) handleProviders(ctx context.Context, msg *telego.Message) {
+	_ = b.replyPlain(ctx, msg.Chat.ID, "Not yet implemented: /providers")
+}
+
+func (b *Bot) handleHelp(ctx context.Context, msg *telego.Message) {
+	_ = b.replyPlain(ctx, msg.Chat.ID, "Not yet implemented: /help")
+}
+
+func (b *Bot) handleKey(ctx context.Context, msg *telego.Message) {
+	_ = b.replyPlain(ctx, msg.Chat.ID, "Not yet implemented: /key")
+}
+
+// handleSubscribe and handleUnsubscribe are implemented in subscribe.go.

pkg/bot/notify.go

@@ -0,0 +1,124 @@
+package bot
+
+import (
+	"context"
+	"fmt"
+	"log"
+
+	"github.com/mymmrac/telego"
+	"github.com/mymmrac/telego/telegoutil"
+	"github.com/salvacybersec/keyhunter/pkg/engine"
+	"github.com/salvacybersec/keyhunter/pkg/scheduler"
+)
+
+// NotifyNewFindings sends a notification to all subscribers about scan results.
+// It returns the number of messages successfully sent and any per-subscriber errors.
+// If FindingCount is 0 and Error is nil, no notification is sent (silent success).
+// If Error is non-nil, an error notification is sent instead.
+func (b *Bot) NotifyNewFindings(result scheduler.JobResult) (int, []error) {
+	// No notification for zero-finding success.
+	if result.FindingCount == 0 && result.Error == nil {
+		return 0, nil
+	}
+
+	subs, err := b.cfg.DB.ListSubscribers()
+	if err != nil {
+		log.Printf("notify: listing subscribers: %v", err)
+		return 0, []error{fmt.Errorf("listing subscribers: %w", err)}
+	}
+	if len(subs) == 0 {
+		return 0, nil
+	}
+
+	var msg string
+	if result.Error != nil {
+		msg = formatErrorNotification(result)
+	} else {
+		msg = formatNotification(result)
+	}
+
+	var sent int
+	var errs []error
+	for _, sub := range subs {
+		if b.bot == nil {
+			// No telego bot (test mode) -- count as would-send.
+			continue
+		}
+		params := telegoutil.Message(telego.ChatID{ID: sub.ChatID}, msg)
+		if _, sendErr := b.bot.SendMessage(context.Background(), params); sendErr != nil {
+			log.Printf("notify: sending to chat %d: %v", sub.ChatID, sendErr)
+			errs = append(errs, fmt.Errorf("chat %d: %w", sub.ChatID, sendErr))
+			continue
+		}
+		sent++
+	}
+
+	return sent, errs
+}
+
+// NotifyFinding sends a real-time notification about an individual finding
+// to all subscribers. The key is always masked.
+func (b *Bot) NotifyFinding(finding engine.Finding) (int, []error) {
+	subs, err := b.cfg.DB.ListSubscribers()
+	if err != nil {
+		log.Printf("notify: listing subscribers: %v", err)
+		return 0, []error{fmt.Errorf("listing subscribers: %w", err)}
+	}
+	if len(subs) == 0 {
+		return 0, nil
+	}
+
+	msg := formatFindingNotification(finding)
+
+	var sent int
+	var errs []error
+	for _, sub := range subs {
+		if b.bot == nil {
+			continue
+		}
+		params := telegoutil.Message(telego.ChatID{ID: sub.ChatID}, msg)
+		if _, sendErr := b.bot.SendMessage(context.Background(), params); sendErr != nil {
+			log.Printf("notify: sending finding to chat %d: %v", sub.ChatID, sendErr)
+			errs = append(errs, fmt.Errorf("chat %d: %w", sub.ChatID, sendErr))
+			continue
+		}
+		sent++
+	}
+
+	return sent, errs
+}
+
+// formatNotification builds the notification message for a successful scan
+// with findings.
+func formatNotification(result scheduler.JobResult) string {
+	return fmt.Sprintf(
+		"New findings from scheduled scan!\n\nJob: %s\nNew keys found: %d\nDuration: %s\n\nUse /stats for details.",
+		result.JobName,
+		result.FindingCount,
+		result.Duration,
+	)
+}
+
+// formatErrorNotification builds the notification message for a scan that
+// encountered an error.
+func formatErrorNotification(result scheduler.JobResult) string {
+	return fmt.Sprintf(
+		"Scheduled scan error\n\nJob: %s\nDuration: %s\nError: %v",
+		result.JobName,
+		result.Duration,
+		result.Error,
+	)
+}
+
+// formatFindingNotification builds the notification message for an individual
+// finding. Always uses the masked key.
+func formatFindingNotification(finding engine.Finding) string {
+	return fmt.Sprintf(
+		"New key detected!\nProvider: %s\nKey: %s\nSource: %s:%d\nConfidence: %s",
+		finding.ProviderName,
+		finding.KeyMasked,
+		finding.Source,
+		finding.LineNumber,
+		finding.Confidence,
+	)
+}

pkg/bot/subscribe.go

@@ -0,0 +1,59 @@
+package bot
+
+import (
+	"context"
+	"fmt"
+	"log"
+
+	"github.com/mymmrac/telego"
+)
+
+// handleSubscribe adds the requesting chat to the subscribers table.
+// If the chat is already subscribed, it informs the user without error.
+func (b *Bot) handleSubscribe(ctx context.Context, msg *telego.Message) {
+	chatID := msg.Chat.ID
+	var username string
+	if msg.From != nil {
+		username = msg.From.Username
+	}
+
+	subscribed, err := b.cfg.DB.IsSubscribed(chatID)
+	if err != nil {
+		log.Printf("subscribe: checking subscription for chat %d: %v", chatID, err)
+		_ = b.replyPlain(ctx, chatID, "Error checking subscription status. Please try again.")
+		return
+	}
+
+	if subscribed {
+		_ = b.replyPlain(ctx, chatID, "You are already subscribed to notifications.")
+		return
+	}
+
+	if err := b.cfg.DB.AddSubscriber(chatID, username); err != nil {
+		log.Printf("subscribe: adding subscriber chat %d: %v", chatID, err)
+		_ = b.replyPlain(ctx, chatID, fmt.Sprintf("Error subscribing: %v", err))
+		return
+	}
+
+	_ = b.replyPlain(ctx, chatID, "Subscribed! You will receive notifications when new API keys are found.")
+}
+
+// handleUnsubscribe removes the requesting chat from the subscribers table.
+// If the chat was not subscribed, it informs the user without error.
+func (b *Bot) handleUnsubscribe(ctx context.Context, msg *telego.Message) {
+	chatID := msg.Chat.ID
+
+	rows, err := b.cfg.DB.RemoveSubscriber(chatID)
+	if err != nil {
+		log.Printf("unsubscribe: removing subscriber chat %d: %v", chatID, err)
+		_ = b.replyPlain(ctx, chatID, fmt.Sprintf("Error unsubscribing: %v", err))
+		return
+	}
+
+	if rows == 0 {
+		_ = b.replyPlain(ctx, chatID, "You are not subscribed.")
+		return
+	}
+
+	_ = b.replyPlain(ctx, chatID, "Unsubscribed. You will no longer receive notifications.")
+}

pkg/bot/subscribe_test.go

@@ -0,0 +1,121 @@
+package bot
+
+import (
+	"testing"
+	"time"
+
+	"github.com/salvacybersec/keyhunter/pkg/engine"
+	"github.com/salvacybersec/keyhunter/pkg/scheduler"
+	"github.com/salvacybersec/keyhunter/pkg/storage"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func openTestDB(t *testing.T) *storage.DB {
+	t.Helper()
+	db, err := storage.Open(":memory:")
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = db.Close() })
+	return db
+}
+
+func TestSubscribeUnsubscribe(t *testing.T) {
+	db := openTestDB(t)
+
+	// Initially not subscribed.
+	ok, err := db.IsSubscribed(12345)
+	require.NoError(t, err)
+	assert.False(t, ok, "should not be subscribed initially")
+
+	// Subscribe.
+	err = db.AddSubscriber(12345, "testuser")
+	require.NoError(t, err)
+
+	ok, err = db.IsSubscribed(12345)
+	require.NoError(t, err)
+	assert.True(t, ok, "should be subscribed after AddSubscriber")
+
+	// Unsubscribe.
+	rows, err := db.RemoveSubscriber(12345)
+	require.NoError(t, err)
+	assert.Equal(t, int64(1), rows, "should have removed 1 row")
+
+	ok, err = db.IsSubscribed(12345)
+	require.NoError(t, err)
+	assert.False(t, ok, "should not be subscribed after RemoveSubscriber")
+
+	// Unsubscribe again returns 0 rows.
+	rows, err = db.RemoveSubscriber(12345)
+	require.NoError(t, err)
+	assert.Equal(t, int64(0), rows, "should have removed 0 rows when not subscribed")
+}
+
+func TestNotifyNewFindings_NoSubscribers(t *testing.T) {
+	db := openTestDB(t)
+
+	b := &Bot{cfg: Config{DB: db}}
+	sent, errs := b.NotifyNewFindings(scheduler.JobResult{
+		JobName:      "nightly-scan",
+		FindingCount: 5,
+		Duration:     10 * time.Second,
+	})
+	assert.Equal(t, 0, sent, "should send 0 messages with no subscribers")
+	assert.Empty(t, errs, "should have no errors with no subscribers")
+}
+
+func TestNotifyNewFindings_ZeroFindings(t *testing.T) {
+	db := openTestDB(t)
+	_ = db.AddSubscriber(12345, "user1")
+
+	b := &Bot{cfg: Config{DB: db}}
+	sent, errs := b.NotifyNewFindings(scheduler.JobResult{
+		JobName:      "nightly-scan",
+		FindingCount: 0,
+		Duration:     3 * time.Second,
+	})
+	assert.Equal(t, 0, sent, "should not notify for zero findings")
+	assert.Empty(t, errs, "should have no errors for zero findings")
+}
+
+func TestFormatNotification(t *testing.T) {
+	result := scheduler.JobResult{
+		JobName:      "nightly-scan",
+		FindingCount: 7,
+		Duration:     2*time.Minute + 30*time.Second,
+	}
+	msg := formatNotification(result)
+	assert.Contains(t, msg, "nightly-scan", "message should contain job name")
+	assert.Contains(t, msg, "7", "message should contain finding count")
+	assert.Contains(t, msg, "2m30s", "message should contain duration")
+	assert.Contains(t, msg, "/stats", "message should reference /stats command")
+}
+
+func TestFormatNotification_Error(t *testing.T) {
+	result := scheduler.JobResult{
+		JobName:      "daily-scan",
+		FindingCount: 0,
+		Duration:     5 * time.Second,
+		Error:        assert.AnError,
+	}
+	msg := formatErrorNotification(result)
+	assert.Contains(t, msg, "daily-scan", "error message should contain job name")
+	assert.Contains(t, msg, "error", "error message should indicate error")
+}
+
+func TestFormatFindingNotification(t *testing.T) {
+	finding := engine.Finding{
+		ProviderName: "OpenAI",
+		KeyValue:     "sk-proj-1234567890abcdef",
+		KeyMasked:    "sk-proj-...cdef",
+		Confidence:   "high",
+		Source:       "/tmp/test.py",
+		LineNumber:   42,
+	}
+	msg := formatFindingNotification(finding)
+	assert.Contains(t, msg, "OpenAI", "should contain provider name")
+	assert.Contains(t, msg, "sk-proj-...cdef", "should contain masked key")
+	assert.NotContains(t, msg, "sk-proj-1234567890abcdef", "should NOT contain full key")
+	assert.Contains(t, msg, "/tmp/test.py", "should contain source path")
+	assert.Contains(t, msg, "42", "should contain line number")
+	assert.Contains(t, msg, "high", "should contain confidence")
+}