From 63d094aa55a1f2813d202609f97588ea0bb649ba Mon Sep 17 00:00:00 2001 From: Frederick Date: Thu, 23 Apr 2026 07:16:34 +0200 Subject: [PATCH] feat: add Telegram bot integration --- pkg/telegram/auth.go | 17 ++++++++ pkg/telegram/bot.go | 59 ++++++++++++++++++++++++++ pkg/telegram/formatter.go | 44 ++++++++++++++++++++ pkg/telegram/handler.go | 88 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 208 insertions(+) create mode 100644 pkg/telegram/auth.go create mode 100644 pkg/telegram/bot.go create mode 100644 pkg/telegram/formatter.go create mode 100644 pkg/telegram/handler.go diff --git a/pkg/telegram/auth.go b/pkg/telegram/auth.go new file mode 100644 index 000000000..65c3bf94c --- /dev/null +++ b/pkg/telegram/auth.go @@ -0,0 +1,17 @@ +package telegram + +type Auth struct { + allowed map[int64]bool +} + +func NewAuth(ids []int64) *Auth { + m := make(map[int64]bool, len(ids)) + for _, id := range ids { + m[id] = true + } + return &Auth{allowed: m} +} + +func (a *Auth) IsAllowed(userID int64) bool { + return a.allowed[userID] +} \ No newline at end of file diff --git a/pkg/telegram/bot.go b/pkg/telegram/bot.go new file mode 100644 index 000000000..b4ce5db61 --- /dev/null +++ b/pkg/telegram/bot.go @@ -0,0 +1,59 @@ +package telegram + +import ( + "context" + "log" + + tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5" +) + +type Bot struct { + api *tgbotapi.BotAPI + handler *Handler + auth *Auth + stopCh chan struct{} +} + +func New(token string, allowedIDs []int64, svc FlowService) (*Bot, error) { + api, err := tgbotapi.NewBotAPI(token) + if err != nil { + return nil, err + } + api.Debug = false + log.Printf("Telegram bot authorized as @%s", api.Self.UserName) + return &Bot{ + api: api, + handler: NewHandler(api, svc), + auth: NewAuth(allowedIDs), + stopCh: make(chan struct{}), + }, nil +} + +func (b *Bot) Start(ctx context.Context) { + u := tgbotapi.NewUpdate(0) + u.Timeout = 60 + updates := b.api.GetUpdatesChan(u) + log.Println("Telegram bot polling for updates...") + for { + select { + case update := <-updates: + if update.Message == nil { + continue + } + if !b.auth.IsAllowed(update.Message.From.ID) { + log.Printf("Telegram: blocked user %d", update.Message.From.ID) + continue + } + go b.handler.Handle(update.Message) + case <-ctx.Done(): + return + case <-b.stopCh: + return + } + } +} + +func (b *Bot) Stop() { + close(b.stopCh) + b.api.StopReceivingUpdates() +} \ No newline at end of file diff --git a/pkg/telegram/formatter.go b/pkg/telegram/formatter.go new file mode 100644 index 000000000..df2866410 --- /dev/null +++ b/pkg/telegram/formatter.go @@ -0,0 +1,44 @@ +package telegram + +import ( + "fmt" + "strings" +) + +func FormatWelcome() string { + return `*Welcome to PentAGI* + +I am your AI-powered security assistant. + +Send /help to see available commands.` +} + +func FormatHelp() string { + return `*Available commands* + +/flows — list your recent flows +/new — create a new flow +/status — check flow status +/stop — stop a running flow +/help — show this message` +} + +func FormatFlowList(flows []Flow) string { + if len(flows) == 0 { + return "No flows found. Use /new to create one." + } + var sb strings.Builder + sb.WriteString("*Your flows:*\n\n") + for _, f := range flows { + sb.WriteString(fmt.Sprintf("• `%s` — %s (%s)\n", f.ID[:8], f.Title, f.Status)) + } + return sb.String() +} + +func FormatFlowCreated(f *Flow) string { + return fmt.Sprintf("*Flow created*\n\nID: `%s`\nTask: %s\nStatus: %s", f.ID, f.Title, f.Status) +} + +func FormatFlowStatus(f *Flow) string { + return fmt.Sprintf("*Flow status*\n\nID: `%s`\nTask: %s\nStatus: %s", f.ID, f.Title, f.Status) +} \ No newline at end of file diff --git a/pkg/telegram/handler.go b/pkg/telegram/handler.go new file mode 100644 index 000000000..7f94bd63d --- /dev/null +++ b/pkg/telegram/handler.go @@ -0,0 +1,88 @@ +package telegram + +import ( + "fmt" + + tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5" +) + +type FlowService interface { + ListFlows(userID int64) ([]Flow, error) + CreateFlow(userID int64, task string) (*Flow, error) + GetFlowStatus(flowID string) (*Flow, error) + StopFlow(flowID string) error +} + +type Flow struct { + ID string + Title string + Status string +} + +type Handler struct { + api *tgbotapi.BotAPI + svc FlowService +} + +func NewHandler(api *tgbotapi.BotAPI, svc FlowService) *Handler { + return &Handler{api: api, svc: svc} +} + +func (h *Handler) Handle(msg *tgbotapi.Message) { + switch msg.Command() { + case "start": + h.reply(msg, FormatWelcome()) + case "help": + h.reply(msg, FormatHelp()) + case "flows": + flows, err := h.svc.ListFlows(msg.From.ID) + if err != nil { + h.reply(msg, fmt.Sprintf("Error fetching flows: %v", err)) + return + } + h.reply(msg, FormatFlowList(flows)) + case "new": + task := msg.CommandArguments() + if task == "" { + h.reply(msg, "Usage: /new ") + return + } + flow, err := h.svc.CreateFlow(msg.From.ID, task) + if err != nil { + h.reply(msg, fmt.Sprintf("Error creating flow: %v", err)) + return + } + h.reply(msg, FormatFlowCreated(flow)) + case "status": + id := msg.CommandArguments() + if id == "" { + h.reply(msg, "Usage: /status ") + return + } + flow, err := h.svc.GetFlowStatus(id) + if err != nil { + h.reply(msg, fmt.Sprintf("Error: %v", err)) + return + } + h.reply(msg, FormatFlowStatus(flow)) + case "stop": + id := msg.CommandArguments() + if id == "" { + h.reply(msg, "Usage: /stop ") + return + } + if err := h.svc.StopFlow(id); err != nil { + h.reply(msg, fmt.Sprintf("Error: %v", err)) + return + } + h.reply(msg, "Flow stopped.") + default: + h.reply(msg, "Unknown command. Send /help for available commands.") + } +} + +func (h *Handler) reply(msg *tgbotapi.Message, text string) { + m := tgbotapi.NewMessage(msg.Chat.ID, text) + m.ParseMode = tgbotapi.ModeMarkdown + h.api.Send(m) +} \ No newline at end of file