diff --git a/.gitignore b/.gitignore old mode 100755 new mode 100644 diff --git a/anthropic.go b/anthropic.go old mode 100755 new mode 100644 diff --git a/bot.go b/bot.go old mode 100755 new mode 100644 index b90593c..2cc861f --- a/bot.go +++ b/bot.go @@ -43,6 +43,33 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock) (*Bot, error) { return nil, err } + // Ensure the owner exists in the Users table + var owner User + err = db.Where("telegram_id = ? AND bot_id = ?", config.OwnerTelegramID, botEntry.ID).First(&owner).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + // Assign the "owner" role + var ownerRole Role + err := db.Where("name = ?", "owner").First(&ownerRole).Error + if err != nil { + return nil, fmt.Errorf("owner role not found: %w", err) + } + + owner = User{ + BotID: botEntry.ID, + TelegramID: config.OwnerTelegramID, + Username: "Owner", // You might want to fetch the actual username + RoleID: ownerRole.ID, + IsOwner: true, + } + + if err := db.Create(&owner).Error; err != nil { + return nil, fmt.Errorf("failed to create owner user: %w", err) + } + } else if err != nil { + return nil, err + } + + // Initialize Anthropic client anthropicClient := anthropic.NewClient(os.Getenv("ANTHROPIC_API_KEY")) b := &Bot{ @@ -69,26 +96,70 @@ func (b *Bot) Start(ctx context.Context) { b.tgBot.Start(ctx) } -func (b *Bot) getOrCreateUser(userID int64, username string) (User, error) { +func (b *Bot) getOrCreateUser(userID int64, username string, isOwner bool) (User, error) { var user User - err := b.db.Preload("Role").Where("telegram_id = ?", userID).First(&user).Error + err := b.db.Preload("Role").Where("telegram_id = ? AND bot_id = ?", userID, b.botID).First(&user).Error if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { - var defaultRole Role - if err := b.db.Where("name = ?", "user").First(&defaultRole).Error; err != nil { - return User{}, err + var role Role + if isOwner { + role, err = b.getRoleByName("owner") + if err != nil { + return User{}, err + } + } else { + role, err = b.getRoleByName("admin") + if err != nil { + return User{}, err + } } - user = User{TelegramID: userID, Username: username, RoleID: defaultRole.ID} + + user = User{ + BotID: b.botID, + TelegramID: userID, + Username: username, + RoleID: role.ID, + IsOwner: isOwner, + } + if err := b.db.Create(&user).Error; err != nil { return User{}, err } } else { return User{}, err } + } else { + if isOwner && !user.IsOwner { + // Check if another owner exists + var existingOwner User + err := b.db.Where("bot_id = ? AND is_owner = ?", b.botID, true).First(&existingOwner).Error + if err == nil { + return User{}, fmt.Errorf("a bot can have only one owner") + } else if !errors.Is(err, gorm.ErrRecordNotFound) { + return User{}, err + } + // Promote to owner + user.Role, err = b.getRoleByName("owner") + if err != nil { + return User{}, err + } + user.RoleID = user.Role.ID + user.IsOwner = true + if err := b.db.Save(&user).Error; err != nil { + return User{}, err + } + } } + return user, nil } +func (b *Bot) getRoleByName(roleName string) (Role, error) { + var role Role + err := b.db.Where("name = ?", roleName).First(&role).Error + return role, err +} + func (b *Bot) createMessage(chatID, userID int64, username, userRole, text string, isUser bool) Message { message := Message{ ChatID: chatID, diff --git a/clock.go b/clock.go old mode 100755 new mode 100644 diff --git a/config.go b/config.go old mode 100755 new mode 100644 index 0ff9436..2d1e556 --- a/config.go +++ b/config.go @@ -19,6 +19,7 @@ type BotConfig struct { Model anthropic.Model `json:"model"` // Changed from string to anthropic.Model SystemPrompts map[string]string `json:"system_prompts"` Active bool `json:"active"` // New field to control bot activity + OwnerTelegramID int64 `json:"owner_telegram_id"` } // Custom unmarshalling to handle anthropic.Model diff --git a/config/default.json b/config/default.json old mode 100755 new mode 100644 index 9df35f7..fc0e75a --- a/config/default.json +++ b/config/default.json @@ -2,6 +2,7 @@ "id": "default_bot", "active": false, "telegram_token": "YOUR_TELEGRAM_BOT_TOKEN", + "owner_telegram_id": 000000000, "memory_size": 10, "messages_per_hour": 20, "messages_per_day": 100, diff --git a/database.go b/database.go old mode 100755 new mode 100644 index da7a29a..d35a25d --- a/database.go +++ b/database.go @@ -27,11 +27,15 @@ func initDB() (*gorm.DB, error) { return nil, fmt.Errorf("failed to connect to database: %w", err) } + // AutoMigrate with unique constraint for owners err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}) if err != nil { return nil, fmt.Errorf("failed to migrate database schema: %w", err) } + // Add unique index for owners per bot + db.SetupJoinTable(&BotModel{}, "Users", &User{}) + err = createDefaultRoles(db) if err != nil { return nil, err diff --git a/go.mod b/go.mod old mode 100755 new mode 100644 diff --git a/go.sum b/go.sum old mode 100755 new mode 100644 diff --git a/handlers.go b/handlers.go old mode 100755 new mode 100644 index a7d9473..4d01477 --- a/handlers.go +++ b/handlers.go @@ -69,7 +69,14 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U return } - user, err := b.getOrCreateUser(userID, username) + // 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 + if err == nil { + isOwner = true + } + + user, err := b.getOrCreateUser(userID, username, isOwner) if err != nil { log.Printf("Error getting or creating user: %v", err) return @@ -84,8 +91,8 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U contextMessages := b.prepareContextMessages(chatMemory) - isEmojiOnly := isOnlyEmojis(text) // Ensure you have this variable defined - response, err := b.getAnthropicResponse(ctx, contextMessages, b.isNewChat(chatID), b.isAdminOrOwner(userID), isEmojiOnly) + isEmojiOnly := isOnlyEmojis(text) + response, err := b.getAnthropicResponse(ctx, contextMessages, b.isNewChat(chatID), isOwner, isEmojiOnly) if err != nil { log.Printf("Error getting Anthropic response: %v", err) response = "I'm sorry, I'm having trouble processing your request right now." diff --git a/main.go b/main.go old mode 100755 new mode 100644 diff --git a/models.go b/models.go old mode 100755 new mode 100644 index f518754..3c2247f --- a/models.go +++ b/models.go @@ -11,7 +11,7 @@ type BotModel struct { Identifier string `gorm:"uniqueIndex"` // Renamed from ID to Identifier Name string Configs []ConfigModel `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"` - Users []User `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"` // Added foreign key + Users []User `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"` // Associated users Messages []Message `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"` } @@ -54,9 +54,10 @@ type Role struct { type User struct { gorm.Model - BotID uint `gorm:"index"` // Added foreign key to BotModel - TelegramID int64 `gorm:"uniqueIndex"` // Consider composite unique index if TelegramID is unique per Bot + BotID uint `gorm:"index"` // Foreign key to BotModel + TelegramID int64 `gorm:"uniqueIndex;not null"` // Unique per user Username string RoleID uint Role Role `gorm:"foreignKey:RoleID"` + IsOwner bool `gorm:"default:false"` // Indicates if the user is the owner } diff --git a/rate_limiter.go b/rate_limiter.go old mode 100755 new mode 100644 diff --git a/rate_limiter_test.go b/rate_limiter_test.go old mode 100755 new mode 100644 index 6e4f54a..a2871f3 --- a/rate_limiter_test.go +++ b/rate_limiter_test.go @@ -22,6 +22,7 @@ func TestCheckRateLimits(t *testing.T) { TempBanDuration: "1m", // Temporary ban duration of 1 minute for testing SystemPrompts: make(map[string]string), TelegramToken: "YOUR_TELEGRAM_BOT_TOKEN", + OwnerTelegramID: 123456789, } // Initialize the Bot with mock data and MockClock