mirror of
https://github.com/HugeFrog24/go-telegram-bot.git
synced 2026-05-01 07:42:18 +00:00
Security/quality
This commit is contained in:
HugeFrog24
Executable → Regular
+323
-85
@@ -2,6 +2,8 @@ package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/go-telegram/bot"
|
||||
@@ -29,53 +31,32 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
|
||||
businessConnectionID = message.BusinessConnectionID
|
||||
}
|
||||
|
||||
if message.From == nil {
|
||||
// Channel posts and some automated messages have no sender — ignore them.
|
||||
// see: https://core.telegram.org/bots/api#message
|
||||
return
|
||||
}
|
||||
|
||||
chatID := message.Chat.ID
|
||||
userID := message.From.ID
|
||||
username := message.From.Username
|
||||
firstName := message.From.FirstName
|
||||
lastName := message.From.LastName
|
||||
languageCode := message.From.LanguageCode
|
||||
isPremium := message.From.IsPremium
|
||||
messageTime := message.Date
|
||||
text := message.Text
|
||||
|
||||
// Pass the incoming message through the centralized screen for storage
|
||||
_, err := b.screenIncomingMessage(message)
|
||||
// Check if it's a new chat
|
||||
isNewChatFlag := b.isNewChat(chatID)
|
||||
|
||||
// Screen incoming message (store to DB + add to chat memory)
|
||||
userMsg, err := b.screenIncomingMessage(message)
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("Error storing user message: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Check if the message is a command
|
||||
if message.Entities != nil {
|
||||
for _, entity := range message.Entities {
|
||||
if entity.Type == "bot_command" {
|
||||
command := strings.TrimSpace(message.Text[entity.Offset : entity.Offset+entity.Length])
|
||||
switch command {
|
||||
case "/stats":
|
||||
b.sendStats(ctx, chatID, userID, username, businessConnectionID)
|
||||
return
|
||||
case "/whoami":
|
||||
b.sendWhoAmI(ctx, chatID, userID, username, businessConnectionID)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check if the message contains a sticker
|
||||
if message.Sticker != nil {
|
||||
b.handleStickerMessage(ctx, chatID, userID, message, businessConnectionID)
|
||||
return
|
||||
}
|
||||
|
||||
// Rate limit check
|
||||
if !b.checkRateLimits(userID) {
|
||||
b.sendRateLimitExceededMessage(ctx, chatID, businessConnectionID)
|
||||
return
|
||||
}
|
||||
|
||||
// Proceed only if the message contains text
|
||||
if text == "" {
|
||||
InfoLogger.Printf("Received a non-text message from user %d in chat %d", userID, chatID)
|
||||
return
|
||||
}
|
||||
|
||||
// Determine if the user is the owner
|
||||
var isOwner bool
|
||||
err = b.db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", userID, b.botID, true).First(&User{}).Error
|
||||
@@ -83,39 +64,181 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U
|
||||
isOwner = true
|
||||
}
|
||||
|
||||
user, err := b.getOrCreateUser(userID, username, isOwner)
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("Error getting or creating user: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Update the username if it's empty or has changed
|
||||
if user.Username != username {
|
||||
user.Username = username
|
||||
if err := b.db.Save(&user).Error; err != nil {
|
||||
ErrorLogger.Printf("Error updating user username: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Determine if the text contains only emojis
|
||||
isEmojiOnly := isOnlyEmojis(text)
|
||||
|
||||
// Prepare context messages for Anthropic
|
||||
// Get the chat memory which now contains the user's message
|
||||
chatMemory := b.getOrCreateChatMemory(chatID)
|
||||
b.addMessageToChatMemory(chatMemory, b.createMessage(chatID, userID, username, user.Role.Name, text, true))
|
||||
contextMessages := b.prepareContextMessages(chatMemory)
|
||||
|
||||
// Get response from Anthropic
|
||||
response, err := b.getAnthropicResponse(ctx, contextMessages, b.isNewChat(chatID), isOwner, isEmojiOnly)
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("Error getting Anthropic response: %v", err)
|
||||
response = "I'm sorry, I'm having trouble processing your request right now."
|
||||
}
|
||||
if isNewChatFlag {
|
||||
|
||||
// Send the response through the centralized screen
|
||||
if err := b.sendResponse(ctx, chatID, response, businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
return
|
||||
// Get response from Anthropic using the context messages
|
||||
response, err := b.getAnthropicResponse(ctx, contextMessages, true, isOwner, false, username, firstName, lastName, isPremium, languageCode, messageTime)
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("Error getting Anthropic response: %v", err)
|
||||
// Use the same error message as in the non-new chat case
|
||||
response = "I'm sorry, I'm having trouble processing your request right now."
|
||||
}
|
||||
|
||||
// Send the AI-generated response or error message
|
||||
if err := b.sendResponse(ctx, chatID, response, businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
return
|
||||
}
|
||||
} else {
|
||||
user, err := b.getOrCreateUser(userID, username, isOwner)
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("Error getting or creating user: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
// Update the username if it's empty or has changed
|
||||
if user.Username != username {
|
||||
user.Username = username
|
||||
if err := b.db.Save(&user).Error; err != nil {
|
||||
ErrorLogger.Printf("Error updating user username: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Check if the message is a command
|
||||
if message.Entities != nil {
|
||||
for _, entity := range message.Entities {
|
||||
if entity.Type == "bot_command" {
|
||||
command := strings.TrimSpace(message.Text[entity.Offset : entity.Offset+entity.Length])
|
||||
switch command {
|
||||
case "/stats":
|
||||
// Parse command parameters
|
||||
parts := strings.Fields(message.Text)
|
||||
|
||||
// Default: show global stats
|
||||
if len(parts) == 1 {
|
||||
b.sendStats(ctx, chatID, userID, 0, businessConnectionID)
|
||||
return
|
||||
}
|
||||
|
||||
// Check for "user" parameter
|
||||
if len(parts) >= 2 && parts[1] == "user" {
|
||||
targetUserID := userID // Default to current user
|
||||
|
||||
// If a user ID is provided, parse it
|
||||
if len(parts) >= 3 {
|
||||
var parseErr error
|
||||
targetUserID, parseErr = strconv.ParseInt(parts[2], 10, 64)
|
||||
if parseErr != nil {
|
||||
InfoLogger.Printf("User %d provided invalid user ID format: %s", userID, parts[2])
|
||||
if err := b.sendResponse(ctx, chatID, "Invalid user ID format. Usage: /stats user [user_id]", businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
b.sendStats(ctx, chatID, userID, targetUserID, businessConnectionID)
|
||||
return
|
||||
}
|
||||
|
||||
// Invalid parameter
|
||||
if err := b.sendResponse(ctx, chatID, "Invalid command format. Usage: /stats or /stats user [user_id]", businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
}
|
||||
return
|
||||
case "/whoami":
|
||||
b.sendWhoAmI(ctx, chatID, userID, username, businessConnectionID)
|
||||
return
|
||||
case "/clear":
|
||||
parts := strings.Fields(message.Text)
|
||||
var targetUserID, targetChatID int64
|
||||
if len(parts) > 1 {
|
||||
var parseErr error
|
||||
targetUserID, parseErr = strconv.ParseInt(parts[1], 10, 64)
|
||||
if parseErr != nil {
|
||||
InfoLogger.Printf("User %d provided invalid user ID format: %s", userID, parts[1])
|
||||
if err := b.sendResponse(ctx, chatID, "Invalid user ID format. Usage: /clear [user_id] [chat_id]", businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
if len(parts) > 2 {
|
||||
var parseErr error
|
||||
targetChatID, parseErr = strconv.ParseInt(parts[2], 10, 64)
|
||||
if parseErr != nil {
|
||||
InfoLogger.Printf("User %d provided invalid chat ID format: %s", userID, parts[2])
|
||||
if err := b.sendResponse(ctx, chatID, "Invalid chat ID format. Usage: /clear [user_id] [chat_id]", businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
b.clearChatHistory(ctx, chatID, userID, targetUserID, targetChatID, businessConnectionID, false)
|
||||
return
|
||||
case "/clear_hard":
|
||||
parts := strings.Fields(message.Text)
|
||||
var targetUserID, targetChatID int64
|
||||
if len(parts) > 1 {
|
||||
var parseErr error
|
||||
targetUserID, parseErr = strconv.ParseInt(parts[1], 10, 64)
|
||||
if parseErr != nil {
|
||||
InfoLogger.Printf("User %d provided invalid user ID format: %s", userID, parts[1])
|
||||
if err := b.sendResponse(ctx, chatID, "Invalid user ID format. Usage: /clear_hard [user_id] [chat_id]", businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
if len(parts) > 2 {
|
||||
var parseErr error
|
||||
targetChatID, parseErr = strconv.ParseInt(parts[2], 10, 64)
|
||||
if parseErr != nil {
|
||||
InfoLogger.Printf("User %d provided invalid chat ID format: %s", userID, parts[2])
|
||||
if err := b.sendResponse(ctx, chatID, "Invalid chat ID format. Usage: /clear_hard [user_id] [chat_id]", businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
b.clearChatHistory(ctx, chatID, userID, targetUserID, targetChatID, businessConnectionID, true)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Rate limit check applies to all message types including stickers
|
||||
if !b.checkRateLimits(userID) {
|
||||
b.sendRateLimitExceededMessage(ctx, chatID, businessConnectionID)
|
||||
return
|
||||
}
|
||||
|
||||
// Check if the message contains a sticker
|
||||
if message.Sticker != nil {
|
||||
b.handleStickerMessage(ctx, chatID, userMsg, message, businessConnectionID)
|
||||
return
|
||||
}
|
||||
|
||||
// Proceed only if the message contains text
|
||||
if text == "" {
|
||||
InfoLogger.Printf("Received a non-text message from user %d in chat %d", userID, chatID)
|
||||
return
|
||||
}
|
||||
|
||||
// Determine if the text contains only emojis
|
||||
isEmojiOnly := isOnlyEmojis(text)
|
||||
|
||||
// Prepare context messages for Anthropic
|
||||
chatMemory := b.getOrCreateChatMemory(chatID)
|
||||
contextMessages := b.prepareContextMessages(chatMemory)
|
||||
|
||||
// Get response from Anthropic
|
||||
response, err := b.getAnthropicResponse(ctx, contextMessages, false, isOwner, isEmojiOnly, username, firstName, lastName, isPremium, languageCode, messageTime) // isNewChat is false here
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("Error getting Anthropic response: %v", err)
|
||||
response = "I'm sorry, I'm having trouble processing your request right now."
|
||||
}
|
||||
|
||||
// Send the response
|
||||
if err := b.sendResponse(ctx, chatID, response, businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -125,21 +248,8 @@ func (b *Bot) sendRateLimitExceededMessage(ctx context.Context, chatID int64, bu
|
||||
}
|
||||
}
|
||||
|
||||
func (b *Bot) handleStickerMessage(ctx context.Context, chatID, userID int64, message *models.Message, businessConnectionID string) {
|
||||
username := message.From.Username
|
||||
|
||||
// Create the user message (without storing it manually)
|
||||
userMessage := b.createMessage(chatID, userID, username, "user", "Sent a sticker.", true)
|
||||
userMessage.StickerFileID = message.Sticker.FileID
|
||||
|
||||
// Safely store the Thumbnail's FileID if available
|
||||
if message.Sticker.Thumbnail != nil {
|
||||
userMessage.StickerPNGFile = message.Sticker.Thumbnail.FileID
|
||||
}
|
||||
|
||||
// Update chat memory with the user message
|
||||
chatMemory := b.getOrCreateChatMemory(chatID)
|
||||
b.addMessageToChatMemory(chatMemory, userMessage)
|
||||
func (b *Bot) handleStickerMessage(ctx context.Context, chatID int64, userMessage Message, message *models.Message, businessConnectionID string) {
|
||||
// userMessage was already screened (stored + added to memory) by handleUpdate — do not call screenIncomingMessage again.
|
||||
|
||||
// Generate AI response about the sticker
|
||||
response, err := b.generateStickerResponse(ctx, userMessage)
|
||||
@@ -155,7 +265,7 @@ func (b *Bot) handleStickerMessage(ctx context.Context, chatID, userID int64, me
|
||||
}
|
||||
}
|
||||
|
||||
// Send the response through the centralized screen
|
||||
// Send the response
|
||||
if err := b.sendResponse(ctx, chatID, response, businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
return
|
||||
@@ -165,18 +275,28 @@ func (b *Bot) handleStickerMessage(ctx context.Context, chatID, userID int64, me
|
||||
func (b *Bot) generateStickerResponse(ctx context.Context, message Message) (string, error) {
|
||||
// Example: Use the sticker type to generate a response
|
||||
if message.StickerFileID != "" {
|
||||
// Create message content with emoji information if available
|
||||
var messageContent string
|
||||
if message.StickerEmoji != "" {
|
||||
messageContent = fmt.Sprintf("User sent a sticker: %s", message.StickerEmoji)
|
||||
} else {
|
||||
messageContent = "User sent a sticker."
|
||||
}
|
||||
|
||||
// Prepare context with information about the sticker
|
||||
contextMessages := []anthropic.Message{
|
||||
{
|
||||
Role: anthropic.RoleUser,
|
||||
Content: []anthropic.MessageContent{
|
||||
anthropic.NewTextMessageContent("User sent a sticker."),
|
||||
anthropic.NewTextMessageContent(messageContent),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// Since this is a sticker message, isEmojiOnly is false
|
||||
response, err := b.getAnthropicResponse(ctx, contextMessages, false, false, false)
|
||||
// Treat sticker messages like emoji messages to get emoji responses
|
||||
// Convert the timestamp to Unix time for the messageTime parameter
|
||||
messageTime := int(message.Timestamp.Unix())
|
||||
response, err := b.getAnthropicResponse(ctx, contextMessages, false, false, true, message.Username, "", "", false, "", messageTime)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
@@ -186,3 +306,121 @@ func (b *Bot) generateStickerResponse(ctx context.Context, message Message) (str
|
||||
|
||||
return "Hmm, that's interesting!", nil
|
||||
}
|
||||
|
||||
func (b *Bot) clearChatHistory(ctx context.Context, chatID int64, currentUserID int64, targetUserID int64, targetChatID int64, businessConnectionID string, hardDelete bool) {
|
||||
// If targetUserID is provided and different from currentUserID, check permissions
|
||||
if targetUserID != 0 && targetUserID != currentUserID {
|
||||
// Check if the current user is an admin or owner
|
||||
if !b.isAdminOrOwner(currentUserID) {
|
||||
InfoLogger.Printf("User %d attempted to clear history for user %d without permission", currentUserID, targetUserID)
|
||||
if err := b.sendResponse(ctx, chatID, "Permission denied. Only admins and owners can clear other users' histories.", businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// Check if the target user exists
|
||||
var targetUser User
|
||||
err := b.db.Where("telegram_id = ? AND bot_id = ?", targetUserID, b.botID).First(&targetUser).Error
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("Error finding target user %d: %v", targetUserID, err)
|
||||
if err := b.sendResponse(ctx, chatID, fmt.Sprintf("User with ID %d not found.", targetUserID), businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
} else {
|
||||
// If no targetUserID is provided, set it to currentUserID
|
||||
targetUserID = currentUserID
|
||||
}
|
||||
|
||||
// Delete messages from the database
|
||||
//
|
||||
// Assumption: this bot is primarily used in private DMs, where each user's messages
|
||||
// are stored with chat_id == their own user_id — not the caller's chat_id. Scoping
|
||||
// a cross-user delete by the caller's chatID would therefore match 0 rows.
|
||||
//
|
||||
// When clearing another user's history the default (targetChatID == 0) deletes all
|
||||
// of that user's messages across every chat for this bot — the natural meaning of
|
||||
// "/clear <userID>" (wipe their entire history with the bot).
|
||||
//
|
||||
// When targetChatID != 0 the deletion is scoped to that specific chat, which is
|
||||
// useful for group moderation ("/clear <userID> <chatID>").
|
||||
var err error
|
||||
if hardDelete {
|
||||
// Permanently delete messages
|
||||
if targetUserID == currentUserID {
|
||||
// Deleting own messages — scope to the current chat only.
|
||||
err = b.db.Unscoped().Where("chat_id = ? AND bot_id = ? AND user_id = ?", chatID, b.botID, targetUserID).Delete(&Message{}).Error
|
||||
InfoLogger.Printf("User %d permanently deleted their own chat history in chat %d", currentUserID, chatID)
|
||||
} else {
|
||||
// Deleting another user's messages — scope bot-wide by default; chat-scoped if targetChatID given (see above).
|
||||
if targetChatID != 0 {
|
||||
err = b.db.Unscoped().Where("chat_id = ? AND bot_id = ? AND user_id = ?", targetChatID, b.botID, targetUserID).Delete(&Message{}).Error
|
||||
InfoLogger.Printf("Admin/owner %d permanently deleted chat history for user %d in chat %d", currentUserID, targetUserID, targetChatID)
|
||||
} else {
|
||||
err = b.db.Unscoped().Where("bot_id = ? AND user_id = ?", b.botID, targetUserID).Delete(&Message{}).Error
|
||||
InfoLogger.Printf("Admin/owner %d permanently deleted all chat history for user %d", currentUserID, targetUserID)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Soft delete messages
|
||||
if targetUserID == currentUserID {
|
||||
// Deleting own messages — scope to the current chat only.
|
||||
err = b.db.Where("chat_id = ? AND bot_id = ? AND user_id = ?", chatID, b.botID, targetUserID).Delete(&Message{}).Error
|
||||
InfoLogger.Printf("User %d soft deleted their own chat history in chat %d", currentUserID, chatID)
|
||||
} else {
|
||||
// Deleting another user's messages — scope bot-wide by default; chat-scoped if targetChatID given (see above).
|
||||
if targetChatID != 0 {
|
||||
err = b.db.Where("chat_id = ? AND bot_id = ? AND user_id = ?", targetChatID, b.botID, targetUserID).Delete(&Message{}).Error
|
||||
InfoLogger.Printf("Admin/owner %d soft deleted chat history for user %d in chat %d", currentUserID, targetUserID, targetChatID)
|
||||
} else {
|
||||
err = b.db.Where("bot_id = ? AND user_id = ?", b.botID, targetUserID).Delete(&Message{}).Error
|
||||
InfoLogger.Printf("Admin/owner %d soft deleted all chat history for user %d", currentUserID, targetUserID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
ErrorLogger.Printf("Error clearing chat history: %v", err)
|
||||
if err := b.sendResponse(ctx, chatID, "Sorry, I couldn't clear the chat history.", businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// Evict the relevant in-memory cache entry so the next access rebuilds from
|
||||
// the now-clean DB. Applies to all cases: own history, cross-user
|
||||
// scoped to a specific chat, and bot-wide cross-user clear.
|
||||
b.chatMemoriesMu.Lock()
|
||||
if targetUserID == currentUserID {
|
||||
// Own history is always scoped to the current chat.
|
||||
delete(b.chatMemories, chatID)
|
||||
} else if targetChatID != 0 {
|
||||
// Admin cleared a specific chat — evict that chat's cache.
|
||||
delete(b.chatMemories, targetChatID)
|
||||
} else {
|
||||
// Bot-wide clear: primary use-case is DMs where chatID == userID.
|
||||
delete(b.chatMemories, targetUserID)
|
||||
}
|
||||
b.chatMemoriesMu.Unlock()
|
||||
|
||||
// Send a confirmation message
|
||||
var confirmationMessage string
|
||||
if targetUserID == currentUserID {
|
||||
confirmationMessage = "Your chat history has been cleared."
|
||||
} else {
|
||||
// Get the username of the target user if available
|
||||
var targetUser User
|
||||
err := b.db.Where("telegram_id = ? AND bot_id = ?", targetUserID, b.botID).First(&targetUser).Error
|
||||
if err == nil && targetUser.Username != "" {
|
||||
confirmationMessage = fmt.Sprintf("Chat history for user @%s (ID: %d) has been cleared.", targetUser.Username, targetUserID)
|
||||
} else {
|
||||
confirmationMessage = fmt.Sprintf("Chat history for user with ID %d has been cleared.", targetUserID)
|
||||
}
|
||||
}
|
||||
|
||||
if err := b.sendResponse(ctx, chatID, confirmationMessage, businessConnectionID); err != nil {
|
||||
ErrorLogger.Printf("Error sending response: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user