feat(17-01): add telego dependency and create Bot package skeleton
- Add telego v1.8.0 as direct dependency for Telegram bot - Create pkg/bot/bot.go with Bot struct, Config, New, Start, Stop - Implement isAllowed chat authorization and per-user rate limiting - Add command dispatch with handler stubs for all 10 commands - Long polling lifecycle with context cancellation for graceful shutdown
This commit is contained in:
salvacybersec
259
pkg/bot/bot.go
Normal file
259
pkg/bot/bot.go
Normal file
@@ -0,0 +1,259 @@
|
||||
// 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
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
// --- 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")
|
||||
}
|
||||
|
||||
func (b *Bot) handleSubscribe(ctx context.Context, msg *telego.Message) {
|
||||
_ = b.replyPlain(ctx, msg.Chat.ID, "Not yet implemented: /subscribe")
|
||||
}
|
||||
|
||||
func (b *Bot) handleUnsubscribe(ctx context.Context, msg *telego.Message) {
|
||||
_ = b.replyPlain(ctx, msg.Chat.ID, "Not yet implemented: /unsubscribe")
|
||||
}
|
||||
Reference in New Issue
Block a user