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

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

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!"
 )

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)
 

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
+}

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
+}

database/structure.go

@@ -12,11 +12,19 @@ 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
 type CustomCommand struct {
 	Link string
 	Loop bool
 }
+
+type DJ struct {
+	Enabled bool
+	Role    string
+}

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 {

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

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
+}