Re-write the libs with discordgo

This commit is contained in:
Atridad Lahiji 2024-10-22 12:24:02 -06:00
parent 7328a0139e
commit 26e931f44f
Signed by: atridad
SSH key fingerprint: SHA256:LGomp8Opq0jz+7kbwNcdfTcuaLRb5Nh0k5AchDDb438
9 changed files with 215 additions and 238 deletions

View file

@ -1,4 +1,2 @@
# Tokens
DISCORD_TOKEN=""
# Comma separated
COOLDOWN_ALLOW_LIST=""

View file

@ -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()
var username string
if i.Member != nil {
username = i.Member.User.Username
} else if i.User != nil {
username = i.User.Username
} else {
username = "User"
if !lib.CheckAndApplyCooldown(s, i, "hs", 10*time.Second) {
return
}
response := fmt.Sprintf("%s was %s's nickname in highschool!", nickname, username)
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()
s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
user, err := lib.GetUser(i)
if err != nil {
lib.RespondWithError(s, i, "Error processing command: "+err.Error())
return
}
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")
}
}

View file

@ -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)
}
}

View file

@ -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")
}
}

View file

@ -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)
}

View file

@ -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, ""
}

47
lib/member.go Normal file
View file

@ -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 ""
}

View file

@ -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
type CooldownManager struct {
cooldowns map[string]time.Time
mu sync.Mutex
}
func NewTimerManager() *TimerManager {
return &TimerManager{
timers: make(map[string]time.Time),
func GetCooldownManager() *CooldownManager {
once.Do(func() {
instance = &CooldownManager{
cooldowns: make(map[string]time.Time),
}
}
func GetInstance() *TimerManager {
mu.Lock()
defer mu.Unlock()
if instance == nil {
instance = &TimerManager{
timers: 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)
}
if time.Now().After(timerEnd) {
delete(m.timers, userID+":"+key)
return false, 0
delete(cm.cooldowns, key)
}
return true, time.Until(timerEnd)
return true, 0
}
func CancelTimer(userID string, key string) {
manager := GetInstance()
// Handle non-existent keys gracefully
if _, exists := manager.timers[userID+":"+key]; !exists {
return
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
}
manager.mu.Lock()
defer manager.mu.Unlock()
delete(manager.timers, userID+":"+key)
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
}
cooldownManager.SetCooldown(user.ID, commandName, duration)
return true
}