From 26e931f44f0b7f683faf8021739bc618d362f4eb Mon Sep 17 00:00:00 2001 From: Atridad Lahiji Date: Tue, 22 Oct 2024 12:24:02 -0600 Subject: [PATCH] Re-write the libs with discordgo --- .env.example | 2 - command/hs.go | 46 ++++++++++++++-------- command/markov.go | 93 +++++++++++++++++++------------------------- command/ping.go | 25 ++++++++++-- lib/errors.go | 39 +++++++++---------- lib/helpers.go | 89 ------------------------------------------ lib/member.go | 47 +++++++++++++++++++++++ lib/timer.go | 98 ++++++++++++++++++++++++----------------------- main.go | 14 +++---- 9 files changed, 215 insertions(+), 238 deletions(-) delete mode 100644 lib/helpers.go create mode 100644 lib/member.go diff --git a/.env.example b/.env.example index c1de146..21acad3 100644 --- a/.env.example +++ b/.env.example @@ -1,4 +1,2 @@ # Tokens DISCORD_TOKEN="" -# Comma separated -COOLDOWN_ALLOW_LIST="" diff --git a/command/hs.go b/command/hs.go index 454411d..0ebc92a 100644 --- a/command/hs.go +++ b/command/hs.go @@ -2,29 +2,41 @@ package command import ( "fmt" + "himbot/lib" + "time" "github.com/bwmarrin/discordgo" ) func HsCommand(s *discordgo.Session, i *discordgo.InteractionCreate) { - options := i.ApplicationCommandData().Options - nickname := options[0].StringValue() + if !lib.CheckAndApplyCooldown(s, i, "hs", 10*time.Second) { + return + } - var username string - if i.Member != nil { - username = i.Member.User.Username - } else if i.User != nil { - username = i.User.Username - } else { - username = "User" - } + options := i.ApplicationCommandData().Options + if len(options) == 0 || options[0].Type != discordgo.ApplicationCommandOptionString { + lib.RespondWithError(s, i, "Please provide a nickname.") + return + } + nickname := options[0].StringValue() - response := fmt.Sprintf("%s was %s's nickname in highschool!", nickname, username) + user, err := lib.GetUser(i) + if err != nil { + lib.RespondWithError(s, i, "Error processing command: "+err.Error()) + return + } - s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{ - Type: discordgo.InteractionResponseChannelMessageWithSource, - Data: &discordgo.InteractionResponseData{ - Content: response, - }, - }) + response := fmt.Sprintf("%s was %s's nickname in high school!", nickname, user.Username) + + err = s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{ + Type: discordgo.InteractionResponseChannelMessageWithSource, + Data: &discordgo.InteractionResponseData{ + Content: response, + }, + }) + + if err != nil { + fmt.Println("Error responding to interaction:", err) + lib.RespondWithError(s, i, "An error occurred while processing the command") + } } diff --git a/command/markov.go b/command/markov.go index 6ad795c..689e26c 100644 --- a/command/markov.go +++ b/command/markov.go @@ -1,20 +1,22 @@ package command import ( + "himbot/lib" "log" "math/rand" "strings" + "time" "github.com/bwmarrin/discordgo" ) -// MarkovCommand generates a random message using Markov chains func MarkovCommand(s *discordgo.Session, i *discordgo.InteractionCreate) { - log.Println("MarkovCommand called") + if !lib.CheckAndApplyCooldown(s, i, "markov", 30*time.Second) { + return + } // Get the channel ID from the interaction channelID := i.ChannelID - log.Printf("Channel ID: %s", channelID) // Get the number of messages to fetch from the option numMessages := 100 // Default value @@ -28,9 +30,40 @@ func MarkovCommand(s *discordgo.Session, i *discordgo.InteractionCreate) { } } } - log.Printf("Fetching up to %d messages", numMessages) - // Fetch messages in batches + // Fetch messages + allMessages, err := fetchMessages(s, channelID, numMessages) + if err != nil { + lib.RespondWithError(s, i, "Failed to fetch messages: "+err.Error()) + return + } + + // Build the Markov chain from the fetched messages + chain := buildMarkovChain(allMessages) + + // Generate a new message using the Markov chain + newMessage := generateMessage(chain) + + // Check if the generated message is empty and provide a fallback message + if newMessage == "" { + newMessage = "I couldn't generate a message. The channel might be empty or contain no usable text." + } + + // Respond to the interaction with the generated message + err = s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{ + Type: discordgo.InteractionResponseChannelMessageWithSource, + Data: &discordgo.InteractionResponseData{ + Content: newMessage, + }, + }) + + if err != nil { + log.Printf("Error responding to interaction: %v", err) + lib.RespondWithError(s, i, "An error occurred while processing the command") + } +} + +func fetchMessages(s *discordgo.Session, channelID string, numMessages int) ([]*discordgo.Message, error) { var allMessages []*discordgo.Message var lastMessageID string @@ -42,9 +75,7 @@ func MarkovCommand(s *discordgo.Session, i *discordgo.InteractionCreate) { batch, err := s.ChannelMessages(channelID, batchSize, lastMessageID, "", "") if err != nil { - log.Printf("Error fetching messages: %v", err) - respondWithError(s, i, "Failed to fetch messages") - return + return nil, err } if len(batch) == 0 { @@ -59,34 +90,7 @@ func MarkovCommand(s *discordgo.Session, i *discordgo.InteractionCreate) { } } - log.Printf("Fetched %d messages", len(allMessages)) - - // Build the Markov chain from the fetched messages - chain := buildMarkovChain(allMessages) - log.Printf("Built Markov chain with %d entries", len(chain)) - - // Generate a new message using the Markov chain - newMessage := generateMessage(chain) - log.Printf("Generated message: %s", newMessage) - - // Check if the generated message is empty and provide a fallback message - if newMessage == "" { - newMessage = "I couldn't generate a message. The channel might be empty or contain no usable text." - } - - // Respond to the interaction with the generated message - err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{ - Type: discordgo.InteractionResponseChannelMessageWithSource, - Data: &discordgo.InteractionResponseData{ - Content: newMessage, - }, - }) - - if err != nil { - log.Printf("Error responding to interaction: %v", err) - return - } - log.Println("Successfully responded to interaction") + return allMessages, nil } // buildMarkovChain creates a Markov chain from a list of messages @@ -94,13 +98,11 @@ func buildMarkovChain(messages []*discordgo.Message) map[string][]string { chain := make(map[string][]string) for _, msg := range messages { words := strings.Fields(msg.Content) - log.Printf("Processing message: %s", msg.Content) // Build the chain by associating each word with the word that follows it for i := 0; i < len(words)-1; i++ { chain[words[i]] = append(chain[words[i]], words[i+1]) } } - log.Printf("Built chain with %d entries", len(chain)) return chain } @@ -132,18 +134,3 @@ func generateMessage(chain map[string][]string) string { return strings.Join(words, " ") } - -// respondWithError sends an error message as a response to the interaction -func respondWithError(s *discordgo.Session, i *discordgo.InteractionCreate, message string) { - log.Printf("Responding with error: %s", message) - err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{ - Type: discordgo.InteractionResponseChannelMessageWithSource, - Data: &discordgo.InteractionResponseData{ - Content: message, - Flags: discordgo.MessageFlagsEphemeral, - }, - }) - if err != nil { - log.Printf("Error sending error response: %v", err) - } -} diff --git a/command/ping.go b/command/ping.go index d579492..34ab1c5 100644 --- a/command/ping.go +++ b/command/ping.go @@ -1,12 +1,31 @@ package command -import "github.com/bwmarrin/discordgo" +import ( + "fmt" + "himbot/lib" + "time" + + "github.com/bwmarrin/discordgo" +) func PingCommand(s *discordgo.Session, i *discordgo.InteractionCreate) { - s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{ + if !lib.CheckAndApplyCooldown(s, i, "ping", 5*time.Second) { + return + } + + // Customize the response based on whether it's a guild or DM + responseContent := "Pong!" + + err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{ Type: discordgo.InteractionResponseChannelMessageWithSource, Data: &discordgo.InteractionResponseData{ - Content: "Pong!", + Content: responseContent, }, }) + + if err != nil { + fmt.Println("Error responding to interaction:", err) + // Optionally, you could try to send an error message to the user + lib.RespondWithError(s, i, "An error occurred while processing the command") + } } diff --git a/lib/errors.go b/lib/errors.go index a491ebb..411cfd8 100644 --- a/lib/errors.go +++ b/lib/errors.go @@ -1,28 +1,27 @@ package lib import ( - "net" - "os" + "fmt" + "log" - "github.com/diamondburned/arikawa/v3/api" - "github.com/diamondburned/arikawa/v3/discord" - "github.com/diamondburned/arikawa/v3/utils/json/option" + "github.com/bwmarrin/discordgo" ) -func ErrorResponse(err error) *api.InteractionResponseData { - var content string - switch e := err.(type) { - case *net.OpError: - content = "**Network Error:** " + e.Error() - case *os.PathError: - content = "**File Error:** " + e.Error() - default: - content = "**Error:** " + err.Error() - } - - return &api.InteractionResponseData{ - Content: option.NewNullableString(content), - Flags: discord.EphemeralMessage, - AllowedMentions: &api.AllowedMentions{}, +// respondWithError sends an error message as a response to the interaction +func RespondWithError(s *discordgo.Session, i *discordgo.InteractionCreate, message string) { + log.Printf("Responding with error: %s", message) + err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{ + Type: discordgo.InteractionResponseChannelMessageWithSource, + Data: &discordgo.InteractionResponseData{ + Content: message, + Flags: discordgo.MessageFlagsEphemeral, + }, + }) + if err != nil { + log.Printf("Error sending error response: %v", err) } } + +func ThrowWithError(command, message string) error { + return fmt.Errorf("error in command '%s': %s", command, message) +} diff --git a/lib/helpers.go b/lib/helpers.go deleted file mode 100644 index 2434fb5..0000000 --- a/lib/helpers.go +++ /dev/null @@ -1,89 +0,0 @@ -package lib - -import ( - "fmt" - "os" - "strings" - "time" - - "github.com/diamondburned/arikawa/v3/discord" -) - -var manager = NewTimerManager() - -// Userish is an interface that captures the common methods you may want to call -// on either a discord.Member or discord.User, including a display name. -type Userish interface { - ID() discord.UserID - Username() string - DisplayName() string -} - -// memberUser adapts a discord.Member to the Userish interface. -type memberUser struct { - *discord.Member -} - -func (mu memberUser) ID() discord.UserID { - return mu.User.ID -} - -func (mu memberUser) Username() string { - return mu.User.Username -} - -func (mu memberUser) DisplayName() string { - // If Nick is set, return it as the display name, otherwise return Username - if mu.Member.Nick != "" { - return mu.Member.Nick - } - return mu.User.Username -} - -// directUser adapts a discord.User to the Userish interface. -type directUser struct { - *discord.User -} - -func (du directUser) ID() discord.UserID { - return du.User.ID -} - -func (du directUser) Username() string { - return du.User.Username -} - -func (du directUser) DisplayName() string { - // For a direct user, the display name is just the username since no nickname is available. - return du.User.Username -} - -// GetUserObject takes an interaction event and returns a Userish, which may be -// either a discord.Member or a discord.User, but exposes it through a consistent interface. -func GetUserObject(event discord.InteractionEvent) Userish { - if event.Member != nil { - return memberUser{event.Member} - } else { - return directUser{event.User} - } -} - -func CooldownHandler(event discord.InteractionEvent, key string, duration time.Duration) (bool, string) { - user := GetUserObject(event) - allowList := strings.Split(os.Getenv("COOLDOWN_ALLOW_LIST"), ",") - - // Check if the user ID is in the allowList - for _, id := range allowList { - if id == user.ID().String() { - return true, "" - } - } - - isOnCooldown, remaining := manager.TimerRunning(user.ID().String(), key) - if isOnCooldown { - return false, fmt.Sprintf("You are on cooldown. Please wait for %v", remaining) - } - - manager.StartTimer(user.ID().String(), key, duration) - return true, "" -} diff --git a/lib/member.go b/lib/member.go new file mode 100644 index 0000000..f386732 --- /dev/null +++ b/lib/member.go @@ -0,0 +1,47 @@ +package lib + +import ( + "github.com/bwmarrin/discordgo" +) + +// InteractionUser represents a user from an interaction, abstracting away the differences +// between guild members and DM users. +type InteractionUser struct { + ID string + Username string + Bot bool +} + +// GetUser extracts user information from an interaction, handling both guild and DM cases. +func GetUser(i *discordgo.InteractionCreate) (*InteractionUser, error) { + if i.Member != nil && i.Member.User != nil { + // Guild interaction + return &InteractionUser{ + ID: i.Member.User.ID, + Username: i.Member.User.Username, + Bot: i.Member.User.Bot, + }, nil + } else if i.User != nil { + // DM interaction + return &InteractionUser{ + ID: i.User.ID, + Username: i.User.Username, + Bot: i.User.Bot, + }, nil + } + + return nil, ThrowWithError("GetUser", "Unable to extract user information from interaction") +} + +// IsInGuild checks if the interaction occurred in a guild. +func IsInGuild(i *discordgo.InteractionCreate) bool { + return i.Member != nil +} + +// GetGuildID safely retrieves the guild ID if the interaction is from a guild. +func GetGuildID(i *discordgo.InteractionCreate) string { + if i.GuildID != "" { + return i.GuildID + } + return "" +} diff --git a/lib/timer.go b/lib/timer.go index d296c78..7f82090 100644 --- a/lib/timer.go +++ b/lib/timer.go @@ -1,72 +1,76 @@ package lib import ( + "fmt" "sync" "time" + + "github.com/bwmarrin/discordgo" ) var ( - mu sync.Mutex - instance *TimerManager + once sync.Once + instance *CooldownManager ) -type TimerManager struct { - timers map[string]time.Time - mu sync.Mutex +type CooldownManager struct { + cooldowns map[string]time.Time + mu sync.Mutex } -func NewTimerManager() *TimerManager { - return &TimerManager{ - timers: make(map[string]time.Time), - } -} - -func GetInstance() *TimerManager { - mu.Lock() - defer mu.Unlock() - - if instance == nil { - instance = &TimerManager{ - timers: make(map[string]time.Time), +func GetCooldownManager() *CooldownManager { + once.Do(func() { + instance = &CooldownManager{ + cooldowns: make(map[string]time.Time), } - } - + }) return instance } -func (m *TimerManager) StartTimer(userID string, key string, duration time.Duration) { - m.mu.Lock() - defer m.mu.Unlock() - - m.timers[userID+":"+key] = time.Now().Add(duration) +func (cm *CooldownManager) SetCooldown(userID, commandName string, duration time.Duration) { + cm.mu.Lock() + defer cm.mu.Unlock() + cm.cooldowns[userID+":"+commandName] = time.Now().Add(duration) } -func (m *TimerManager) TimerRunning(userID string, key string) (bool, time.Duration) { - m.mu.Lock() - defer m.mu.Unlock() +func (cm *CooldownManager) CheckCooldown(userID, commandName string) (bool, time.Duration) { + cm.mu.Lock() + defer cm.mu.Unlock() - timerEnd, exists := m.timers[userID+":"+key] - if !exists { - return false, 0 + key := userID + ":" + commandName + if cooldownEnd, exists := cm.cooldowns[key]; exists { + if time.Now().Before(cooldownEnd) { + return false, time.Until(cooldownEnd) + } + delete(cm.cooldowns, key) } - - if time.Now().After(timerEnd) { - delete(m.timers, userID+":"+key) - return false, 0 - } - - return true, time.Until(timerEnd) + return true, 0 } -func CancelTimer(userID string, key string) { - manager := GetInstance() +func CheckAndApplyCooldown(s *discordgo.Session, i *discordgo.InteractionCreate, commandName string, duration time.Duration) bool { + cooldownManager := GetCooldownManager() + user, err := GetUser(i) + if err != nil { + s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{ + Type: discordgo.InteractionResponseChannelMessageWithSource, + Data: &discordgo.InteractionResponseData{ + Content: "Error processing command: " + err.Error(), + }, + }) + return false + } - // Handle non-existent keys gracefully - if _, exists := manager.timers[userID+":"+key]; !exists { - return - } + canUse, remaining := cooldownManager.CheckCooldown(user.ID, commandName) + if !canUse { + s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{ + Type: discordgo.InteractionResponseChannelMessageWithSource, + Data: &discordgo.InteractionResponseData{ + Content: fmt.Sprintf("You can use this command again in %v", remaining.Round(time.Second)), + }, + }) + return false + } - manager.mu.Lock() - defer manager.mu.Unlock() - delete(manager.timers, userID+":"+key) + cooldownManager.SetCooldown(user.ID, commandName, duration) + return true } diff --git a/main.go b/main.go index 8f65a9e..ddd3cb9 100644 --- a/main.go +++ b/main.go @@ -33,13 +33,13 @@ var ( Name: "markov", Description: "Why did the Markov chain break up with its partner? Because it couldn't handle the past!", Options: []*discordgo.ApplicationCommandOption{ - { - Type: discordgo.ApplicationCommandOptionInteger, - Name: "messages", - Description: "Number of messages to use (default: 100, max: 1000)", - Required: false, - }, - }, + { + Type: discordgo.ApplicationCommandOptionInteger, + Name: "messages", + Description: "Number of messages to use (default: 100, max: 1000)", + Required: false, + }, + }, }, }