Allow commands to be limited to certain roles (#157)

4 files changed, 90 insertions, 9 deletions

diff --git a/src/framework/standard/command.rs b/src/framework/standard/command.rs
index 22a82ce..e869bf5 100644
--- a/src/framework/standard/command.rs
+++ b/src/framework/standard/command.rs
@@ -57,6 +57,8 @@ pub struct Command {
     pub max_args: Option<i32>,
     /// Permissions required to use this command.
     pub required_permissions: Permissions,
+    /// Roles allowed to use this command.
+    pub allowed_roles: Vec<String>,
     /// Whether command should be displayed in help list or not, used by other commands.
     pub help_available: bool,
     /// Whether command can be used only privately or not.
@@ -86,6 +88,7 @@ impl Command {
             max_args: None,
             owners_only: false,
             required_permissions: Permissions::empty(),
+            allowed_roles: Vec::new(),
         }
     }
 }
diff --git a/src/framework/standard/create_command.rs b/src/framework/standard/create_command.rs
index 4fe11f2..546447e 100644
--- a/src/framework/standard/create_command.rs
+++ b/src/framework/standard/create_command.rs
@@ -207,6 +207,13 @@ impl CreateCommand {
 
         self
     }
+
+    /// Sets roles that are allowed to use the command.
+    pub fn allowed_roles(mut self, allowed_roles: Vec<&str>) -> Self {
+        self.0.allowed_roles = allowed_roles.iter().map(|x| x.to_string()).collect();
+
+        self
+    }
 }
 
 impl Default for Command {
@@ -226,6 +233,7 @@ impl Default for Command {
             guild_only: false,
             help_available: true,
             owners_only: false,
+            allowed_roles: Vec::new(),
         }
     }
 }
diff --git a/src/framework/standard/help_commands.rs b/src/framework/standard/help_commands.rs
index f081e1a..c7004d5 100644
--- a/src/framework/standard/help_commands.rs
+++ b/src/framework/standard/help_commands.rs
@@ -29,7 +29,7 @@ use std::fmt::Write;
 use super::command::InternalCommand;
 use super::{Args, Command, CommandGroup, CommandOrAlias};
 use client::Context;
-use model::{ChannelId, Message};
+use model::{ChannelId, Guild, Member, Message};
 use utils::Colour;
 
 fn error_embed(channel_id: &ChannelId, input: &str) {
@@ -50,6 +50,17 @@ fn remove_aliases(cmds: &HashMap<String, CommandOrAlias>) -> HashMap<&String, &I
     result
 }
 
+fn right_roles(cmd: &Command, guild: &Guild, member: &Member) -> bool {
+    if cmd.allowed_roles.len() > 0 {
+        cmd.allowed_roles
+            .iter()
+            .flat_map(|r| guild.role_by_name(&r))
+            .any(|g| member.roles.contains(&g.id))
+    } else {
+        true
+    }
+}
+
 /// Posts an embed showing each individual command group and its commands.
 ///
 /// # Examples
@@ -89,6 +100,19 @@ pub fn with_embeds(_: &mut Context,
                 if name == with_prefix || name == *command_name {
                     match *command {
                         CommandOrAlias::Command(ref cmd) => {
+                            if cmd.allowed_roles.len() > 0 {
+                                if let Some(guild) = msg.guild() {
+                                    let guild = guild.read().unwrap();
+                                    if let Some(member) = guild.members.get(&msg.author.id) {
+                                        if let Ok(permissions) = member.permissions() {
+                                            if !permissions.administrator() &&
+                                               !right_roles(&cmd, &guild, &member) {
+                                                break;
+                                            }
+                                        }
+                                    }
+                                }
+                            }
                             found = Some((command_name, cmd));
                         },
                         CommandOrAlias::Alias(ref name) => {
@@ -188,17 +212,25 @@ pub fn with_embeds(_: &mut Context,
                     let cmd = &commands[name];
 
                     if cmd.help_available {
-                        let _ = write!(desc, "`{}`\n", name);
-
-                        has_commands = true;
+                        if let Some(guild) = msg.guild() {
+                            let guild = guild.read().unwrap();
+                            if let Some(member) = guild.members.get(&msg.author.id) {
+                                if let Ok(permissions) = member.permissions() {
+                                    if cmd.help_available &&
+                                       (right_roles(&cmd, &guild, &member) ||
+                                        permissions.administrator()) {
+                                        let _ = write!(desc, "`{}`\n", name);
+                                        has_commands = true;
+                                    }
+                                }
+                            }
+                        }
                     }
                 }
-
                 if has_commands {
                     e = e.field(|f| f.name(group_name).value(&desc));
                 }
             }
-
             e
         })
     });
@@ -245,6 +277,19 @@ pub fn plain(_: &mut Context,
                 if name == with_prefix || name == *command_name {
                     match *command {
                         CommandOrAlias::Command(ref cmd) => {
+                            if cmd.allowed_roles.len() > 0 {
+                                if let Some(guild) = msg.guild() {
+                                    let guild = guild.read().unwrap();
+                                    if let Some(member) = guild.members.get(&msg.author.id) {
+                                        if let Ok(permissions) = member.permissions() {
+                                            if !permissions.administrator() &&
+                                               !right_roles(&cmd, &guild, &member) {
+                                                break;
+                                            }
+                                        }
+                                    }
+                                }
+                            }
                             found = Some((command_name, cmd));
                         },
                         CommandOrAlias::Alias(ref name) => {
@@ -322,9 +367,16 @@ pub fn plain(_: &mut Context,
 
         for name in command_names {
             let cmd = &commands[name];
-
-            if cmd.help_available {
-                let _ = write!(group_help, "`{}` ", name);
+            if let Some(guild) = msg.guild() {
+                let guild = guild.read().unwrap();
+                if let Some(member) = guild.members.get(&msg.author.id) {
+                    if let Ok(permissions) = member.permissions() {
+                        if cmd.help_available &&
+                           (permissions.administrator() || right_roles(&cmd, &guild, &member)) {
+                            let _ = write!(group_help, "`{}` ", name);
+                        }
+                    }
+                }
             }
         }
 
diff --git a/src/framework/standard/mod.rs b/src/framework/standard/mod.rs
index a8570e6..4fc728d 100644
--- a/src/framework/standard/mod.rs
+++ b/src/framework/standard/mod.rs
@@ -126,6 +126,8 @@ pub enum DispatchError {
     OnlyForGuilds,
     /// When the requested command can only be used by bot owners.
     OnlyForOwners,
+    /// When the requested command requires one role.
+    LackingRole,
     /// When there are too few arguments.
     NotEnoughArguments { min: i32, given: usize },
     /// When there are too many arguments.
@@ -521,6 +523,22 @@ impl StandardFramework {
             } else if self.configuration.disabled_commands.contains(built) {
                 Some(DispatchError::CommandDisabled(built.to_owned()))
             } else {
+                if command.allowed_roles.len() > 0 {
+                    if let Some(guild) = message.guild() {
+                        let guild = guild.read().unwrap();
+                        if let Some(member) = guild.members.get(&message.author.id) {
+                            let right_role = command
+                                .allowed_roles
+                                .iter()
+                                .flat_map(|r| guild.role_by_name(&r))
+                                .any(|g| member.roles.contains(&g.id));
+                            if !right_role {
+                                return Some(DispatchError::LackingRole);
+                            }
+                        }
+                    }
+                }
+
                 let all_passed = command.checks.iter().all(|check| {
                     check(&mut context, message, args, command)
                 });