Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions pkg/telegram/auth.go
Original file line number Diff line number Diff line change
@@ -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]
}
59 changes: 59 additions & 0 deletions pkg/telegram/bot.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
package telegram

import (
Comment thread
freddycodes23 marked this conversation as resolved.
"context"
"log"

Comment thread
freddycodes23 marked this conversation as resolved.
tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5"
)
Comment thread
freddycodes23 marked this conversation as resolved.

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():
Comment thread
freddycodes23 marked this conversation as resolved.
return
case <-b.stopCh:
return
}
Comment thread
freddycodes23 marked this conversation as resolved.
}
}

func (b *Bot) Stop() {
close(b.stopCh)
b.api.StopReceivingUpdates()
}
Comment thread
freddycodes23 marked this conversation as resolved.
44 changes: 44 additions & 0 deletions pkg/telegram/formatter.go
Original file line number Diff line number Diff line change
@@ -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 <task> β€” create a new flow
/status <id> β€” check flow status
/stop <id> β€” stop a running flow
/help β€” show this message`
}

func FormatFlowList(flows []Flow) string {
if len(flows) == 0 {
return "No flows found. Use /new <task> 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)
}
Comment thread
freddycodes23 marked this conversation as resolved.
88 changes: 88 additions & 0 deletions pkg/telegram/handler.go
Original file line number Diff line number Diff line change
@@ -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
}
Comment thread
freddycodes23 marked this conversation as resolved.

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 <task description>")
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 <flow_id>")
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 <flow_id>")
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)
Comment thread
freddycodes23 marked this conversation as resolved.
}