From c87507e619f413da55f9a39a5a9c72d69937f9c2 Mon Sep 17 00:00:00 2001 From: Manuel Raimo Date: Thu, 15 Jun 2023 19:38:54 +0200 Subject: [PATCH] feat: dj mode If enabled, only the users with the DJ role (and the owner) can add songs to the queue (via commands like stream, play, custom). Closes #14 --- commands.go | 130 ++++++++++++++++++++++++++++++++++++ constant.go | 8 +++ database/common/common.go | 38 +++++++---- database/mysql/database.go | 16 ++++- database/sqlite/database.go | 18 ++++- database/structure.go | 8 +++ main.go | 15 +++++ structures.go | 4 ++ utilities.go | 9 +++ 9 files changed, 231 insertions(+), 15 deletions(-) diff --git a/commands.go b/commands.go index 0eca629..65d2040 100644 --- a/commands.go +++ b/commands.go @@ -232,6 +232,34 @@ var ( }, }, }, + { + Name: "dj", + Description: "Toggles DJ mode, which allows only users with the DJ role to control the bot.", + }, + { + Name: "djrole", + Description: "Adds or removes a role from the DJ role list.", + Options: []*discordgo.ApplicationCommandOption{ + { + Type: discordgo.ApplicationCommandOptionRole, + Name: "role", + Description: "Role to remove or add to the DJ role list", + Required: true, + }, + }, + }, + { + Name: "djroletoggle", + Description: "Adds or removes the DJ role from the specified user", + Options: []*discordgo.ApplicationCommandOption{ + { + Type: discordgo.ApplicationCommandOptionUser, + Name: "user", + Description: "User to add or remove the DJ role from", + Required: true, + }, + }, + }, } // Handler @@ -430,6 +458,10 @@ var ( }, // Calls a custom command "custom": func(s *discordgo.Session, i *discordgo.InteractionCreate) { + if djModeCheck(s, i) { + return + } + command := strings.ToLower(i.ApplicationCommandData().Options[0].Value.(string)) if server[i.GuildID].custom[command] != nil { @@ -566,6 +598,10 @@ var ( }, // Streams a song from the given URL, useful for radios "stream": func(s *discordgo.Session, i *discordgo.InteractionCreate) { + if djModeCheck(s, i) { + return + } + if vs := findUserVoiceState(s, i.GuildID, i.Member.User.ID); vs != nil { url := i.ApplicationCommandData().Options[0].Value.(string) if !strings.HasPrefix(url, "file") && isValidURL(url) { @@ -607,10 +643,95 @@ var ( SetColor(0x7289DA).MessageEmbed, i.Interaction, time.Second*5) } }, + // Enables or disables DJ mode + "dj": func(s *discordgo.Session, i *discordgo.InteractionCreate) { + if i.Member.User.ID == owner { + if server[i.GuildID].djMode { + server[i.GuildID].djMode = false + err := db.SetDJSettings(i.GuildID, false) + if err != nil { + lit.Error("Error while disabling DJ mode, %s", err) + } + + sendAndDeleteEmbedInteraction(s, NewEmbed().SetTitle(s.State.User.Username).AddField(djTitle, djDisabled). + SetColor(0x7289DA).MessageEmbed, i.Interaction, time.Second*5) + } else { + server[i.GuildID].djMode = true + err := db.SetDJSettings(i.GuildID, true) + if err != nil { + lit.Error("Error while enabling DJ mode, %s", err) + } + + sendAndDeleteEmbedInteraction(s, NewEmbed().SetTitle(s.State.User.Username).AddField(djTitle, djEnabled). + SetColor(0x7289DA).MessageEmbed, i.Interaction, time.Second*5) + } + } else { + sendAndDeleteEmbedInteraction(s, NewEmbed().SetTitle(s.State.User.Username).AddField(errorTitle, + "Only the owner of the bot can use this command!"). + SetColor(0x7289DA).MessageEmbed, i.Interaction, time.Second*3) + } + }, + // Sets the DJ role + "djrole": func(s *discordgo.Session, i *discordgo.InteractionCreate) { + if i.Member.User.ID == owner { + role := i.ApplicationCommandData().Options[0].RoleValue(nil, "") + if role.ID != server[i.GuildID].djRole { + server[i.GuildID].djRole = role.ID + err := db.UpdateDJRole(i.GuildID, role.ID) + if err != nil { + lit.Error("Error updating DJ role", err) + } + + sendAndDeleteEmbedInteraction(s, NewEmbed().SetTitle(s.State.User.Username).AddField(djTitle, djRoleChanged). + SetColor(0x7289DA).MessageEmbed, i.Interaction, time.Second*5) + } else { + sendAndDeleteEmbedInteraction(s, NewEmbed().SetTitle(s.State.User.Username).AddField(djTitle, djRoleEqual). + SetColor(0x7289DA).MessageEmbed, i.Interaction, time.Second*5) + } + } else { + sendAndDeleteEmbedInteraction(s, NewEmbed().SetTitle(s.State.User.Username).AddField(errorTitle, + "Only the owner of the bot can use this command!"). + SetColor(0x7289DA).MessageEmbed, i.Interaction, time.Second*3) + } + }, + // Adds or removes the DJ role from a user + "djroletoggle": func(s *discordgo.Session, i *discordgo.InteractionCreate) { + if i.Member.User.ID == owner { + var err error + var action string + + user, _ := s.GuildMember(i.GuildID, i.ApplicationCommandData().Options[0].UserValue(nil).ID) + if !hasRole(user.Roles, server[i.GuildID].djRole) { + err = s.GuildMemberRoleAdd(i.GuildID, user.User.ID, server[i.GuildID].djRole) + action = "added!" + } else { + err = s.GuildMemberRoleRemove(i.GuildID, user.User.ID, server[i.GuildID].djRole) + action = "removed!" + } + + if err != nil { + sendAndDeleteEmbedInteraction(s, NewEmbed().SetTitle(s.State.User.Username).AddField(errorTitle, + "The bot doesn't have the necessary permission!"). + SetColor(0x7289DA).MessageEmbed, i.Interaction, time.Second*3) + } else { + sendAndDeleteEmbedInteraction(s, NewEmbed().SetTitle(s.State.User.Username).AddField(djTitle, + "The role has been succefully "+action). + SetColor(0x7289DA).MessageEmbed, i.Interaction, time.Second*5) + } + } else { + sendAndDeleteEmbedInteraction(s, NewEmbed().SetTitle(s.State.User.Username).AddField(errorTitle, + "Only the owner of the bot can use this command!"). + SetColor(0x7289DA).MessageEmbed, i.Interaction, time.Second*3) + } + }, } ) func playCommand(s *discordgo.Session, i *discordgo.InteractionCreate, playlist bool) { + if djModeCheck(s, i) { + return + } + // Check if user is not in a voice channel if vs := findUserVoiceState(s, i.GuildID, i.Member.User.ID); vs != nil { if joinVC(i.Interaction, vs.ChannelID) { @@ -651,3 +772,12 @@ func playCommand(s *discordgo.Session, i *discordgo.InteractionCreate, playlist SetColor(0x7289DA).MessageEmbed, i.Interaction, time.Second*5) } } + +func djModeCheck(s *discordgo.Session, i *discordgo.InteractionCreate) bool { + if server[i.GuildID].djMode && i.Member.User.ID != owner && !hasRole(i.Member.Roles, server[i.GuildID].djRole) { + sendAndDeleteEmbedInteraction(s, NewEmbed().SetTitle(s.State.User.Username).AddField(errorTitle, djNot). + SetColor(0x7289DA).MessageEmbed, i.Interaction, time.Second*3) + return true + } + return false +} diff --git a/constant.go b/constant.go index 2889e7e..154e381 100644 --- a/constant.go +++ b/constant.go @@ -27,6 +27,7 @@ const ( commandsTitle = "Commands" blacklistTitle = "Blacklist" resumeTitle = "Resume" + djTitle = "DJ" ) // Messages for embeds @@ -64,4 +65,11 @@ const ( // Feedback disconnected = "Bye-bye!" + + // DJ + djEnabled = "DJ mode enabled!" + djDisabled = "DJ mode disabled!" + djNot = "User is not a DJ, and DJ mode is enabled!" + djRoleChanged = "DJ role changed!" + djRoleEqual = "DJ role is already that role!" ) diff --git a/database/common/common.go b/database/common/common.go index 09478aa..f154e30 100644 --- a/database/common/common.go +++ b/database/common/common.go @@ -36,11 +36,7 @@ func (c Common) CheckInDb(link string) (queue.Element, error) { func (c Common) AddCommand(command string, song string, guild string, loop bool) error { // Else, we add it to the database _, err := c.db.Exec("INSERT INTO customCommands (`guild`, `command`, `song`, `loop`) VALUES(?, ?, ?, ?)", guild, command, song, loop) - if err != nil { - return err - } - - return nil + return err } // RemoveCustom removes a custom command from the DB and from the command map @@ -98,20 +94,38 @@ func (c Common) RemoveFromDB(el queue.Element) { func (c Common) AddToBlacklist(id string) error { _, err := c.db.Exec("INSERT INTO blacklist (id) VALUES(?)", id) - if err != nil { - return err - } - - return nil + return err } func (c Common) RemoveFromBlacklist(id string) error { _, err := c.db.Exec("DELETE FROM blacklist WHERE id=?", id) + return err +} + +func (c Common) GetDJ() (map[string]database.DJ, error) { + roles := make(map[string]database.DJ) + + rows, err := c.db.Query("SELECT guild, role, enabled FROM dj") if err != nil { - return err + return nil, err } - return nil + for rows.Next() { + var role, guild string + var enabled bool + + err = rows.Scan(&guild, &role, &enabled) + if err != nil { + lit.Error("Error scanning rows from query, %s", err) + continue + } + + roles[guild] = database.DJ{Role: role, Enabled: enabled} + } + + return roles, nil +} + func (c Common) GetBlacklist() (map[string]bool, error) { ids := make(map[string]bool) diff --git a/database/mysql/database.go b/database/mysql/database.go index eb525c8..0de41dd 100644 --- a/database/mysql/database.go +++ b/database/mysql/database.go @@ -14,6 +14,7 @@ const ( tblCommands = "CREATE TABLE IF NOT EXISTS `customCommands` (`guild` varchar(18) NOT NULL, `command` varchar(100) NOT NULL, `song` varchar(100) NOT NULL, `loop` tinyint(1) NOT NULL DEFAULT 0, PRIMARY KEY (`guild`,`command`));" tblBlacklist = "CREATE TABLE IF NOT EXISTS `blacklist`(`id` VARCHAR(20) NOT NULL, PRIMARY KEY (`id`));" tblLink = "CREATE TABLE IF NOT EXISTS `link` ( `songID` varchar(200) NOT NULL, `link` varchar(500) NOT NULL, PRIMARY KEY (`link`), KEY `FK__song2` (`songID`), CONSTRAINT `FK__song2` FOREIGN KEY (`songID`) REFERENCES `song` (`id`) );" + tblDJ = "CREATE TABLE IF NOT EXISTS `dj` ( `guild` VARCHAR(20) NOT NULL, `role` VARCHAR(20) NULL, `enabled` TINYINT(1) NOT NULL DEFAULT '0', PRIMARY KEY (`guild`));" ) var db *sql.DB @@ -30,7 +31,7 @@ func NewDatabase(dsn string) *database.Database { } // Create tables if they don't exist - database.ExecQuery(db, tblSong, tblCommands, tblBlacklist, tblLink) + database.ExecQuery(db, tblSong, tblCommands, tblBlacklist, tblLink, tblDJ) c := common.NewCommon(db) @@ -44,7 +45,10 @@ func NewDatabase(dsn string) *database.Database { Close: c.Close, AddToBlacklist: c.AddToBlacklist, RemoveFromBlacklist: c.RemoveFromBlacklist, + GetDJ: c.GetDJ, + UpdateDJRole: updateDJRole, GetBlacklist: c.GetBlacklist, + SetDJSettings: setDJSettings, } } @@ -66,3 +70,13 @@ func addToDb(el queue.Element, exist bool) { } } } + +func setDJSettings(guild string, enabled bool) error { + _, err := db.Exec("INSERT INTO dj (guild, enabled) VALUES (?, ?) ON DUPLICATE KEY UPDATE enabled = ?", guild, enabled, enabled) + return err +} + +func updateDJRole(guild string, role string) error { + _, err := db.Exec("INSERT INTO dj (guild, role) VALUES (?, ?) ON DUPLICATE KEY UPDATE role = ?", guild, role, role) + return err +} diff --git a/database/sqlite/database.go b/database/sqlite/database.go index dcd694a..2e1677e 100644 --- a/database/sqlite/database.go +++ b/database/sqlite/database.go @@ -13,7 +13,8 @@ const ( tblSong = "CREATE TABLE IF NOT EXISTS `song` (`id` varchar(200) NOT NULL, `title` varchar(200) NOT NULL, `duration` varchar(20) NOT NULL, `thumbnail` varchar(500) NOT NULL, `segments` varchar(1000) NOT NULL, PRIMARY KEY (`id`));" tblCommands = "CREATE TABLE IF NOT EXISTS `customCommands` (`guild` varchar(18) NOT NULL, `command` varchar(100) NOT NULL, `song` varchar(100) NOT NULL, `loop` tinyint(1) NOT NULL DEFAULT 0, PRIMARY KEY (`guild`,`command`));" tblBlacklist = "CREATE TABLE IF NOT EXISTS `blacklist`(`id` VARCHAR(20) NOT NULL, PRIMARY KEY (`id`));" - tblLink = "create table if not exists link ( songID varchar(200) not null references song, link varchar(500) not null constraint link_pk primary key );" + tblLink = "create table IF NOT EXISTS link ( songID varchar(200) not null references song, link varchar(500) not null constraint link_pk primary key );" + tblDJ = "CREATE TABLE IF NOT EXISTS `dj` ( `guild` VARCHAR(20) NOT NULL, `role` VARCHAR(20) NULL, `enabled` TINYINT(1) NOT NULL DEFAULT '0', PRIMARY KEY (`guild`) );" ) var db *sql.DB @@ -30,7 +31,7 @@ func NewDatabase(dsn string) *database.Database { } // Create tables if they don't exist - database.ExecQuery(db, tblSong, tblCommands, tblBlacklist, tblLink) + database.ExecQuery(db, tblSong, tblCommands, tblBlacklist, tblLink, tblDJ) c := common.NewCommon(db) @@ -44,7 +45,10 @@ func NewDatabase(dsn string) *database.Database { Close: c.Close, AddToBlacklist: c.AddToBlacklist, RemoveFromBlacklist: c.RemoveFromBlacklist, + GetDJ: c.GetDJ, + UpdateDJRole: updateDJRole, GetBlacklist: c.GetBlacklist, + SetDJSettings: setDJSettings, } } @@ -66,3 +70,13 @@ func addToDb(el queue.Element, exist bool) { } } } + +func setDJSettings(guild string, enabled bool) error { + _, err := db.Exec("INSERT INTO dj (guild, enabled) VALUES (?, ?) ON CONFLICT(guild) DO UPDATE SET enabled = ?", guild, enabled, enabled) + return err +} + +func updateDJRole(guild string, role string) error { + _, err := db.Exec("INSERT INTO dj (guild, role) VALUES (?, ?) ON CONFLICT(guild) DO UPDATE SET role = ?", guild, role, role) + return err +} diff --git a/database/structure.go b/database/structure.go index 831f3e6..5a1de2d 100644 --- a/database/structure.go +++ b/database/structure.go @@ -12,7 +12,10 @@ type Database struct { AddToBlacklist func(id string) error RemoveFromBlacklist func(id string) error Close func() + UpdateDJRole func(guild string, role string) error + GetDJ func() (map[string]DJ, error) GetBlacklist func() (map[string]bool, error) + SetDJSettings func(guild string, enabled bool) error } // CustomCommand holds data about a custom command @@ -20,3 +23,8 @@ type CustomCommand struct { Link string Loop bool } + +type DJ struct { + Enabled bool + Role string +} diff --git a/main.go b/main.go index de078a4..2bf2ec5 100644 --- a/main.go +++ b/main.go @@ -95,6 +95,21 @@ func init() { lit.Error("Error loading blacklist: %s", err) } + // Load the DJ settings + dj, err := db.GetDJ() + if err != nil { + lit.Error("Error loading DJ settings: %s", err) + } + + for k := range dj { + if server[k] == nil { + initializeServer(k) + } + + server[k].djMode = dj[k].Enabled + server[k].djRole = dj[k].Role + } + // Create folders used by the bot if _, err = os.Stat(cachePath); err != nil { if err = os.Mkdir(cachePath, 0755); err != nil { diff --git a/structures.go b/structures.go index 88ac990..606a0a0 100644 --- a/structures.go +++ b/structures.go @@ -38,6 +38,10 @@ type Server struct { resume chan struct{} // Wait group for waiting for spotify to finish before lowering the clear flag wg *sync.WaitGroup + // Role ID for the DJ role + djRole string + // Whether the DJ mode is enabled + djMode bool } // YtDLP structure for holding yt-dlp data diff --git a/utilities.go b/utilities.go index a1e3857..09eb2de 100644 --- a/utilities.go +++ b/utilities.go @@ -285,3 +285,12 @@ func clearAndExit(guildID string) { server[guildID].voiceChannel = "" } } + +func hasRole(roles []string, role string) bool { + for _, r := range roles { + if r == role { + return true + } + } + return false +}