Framework blocks, disabled commands, improvements

4 files changed, 205 insertions, 48 deletions

diff --git a/src/ext/framework/command.rs b/src/ext/framework/command.rs
index c1e9197..acddca9 100644
--- a/src/ext/framework/command.rs
+++ b/src/ext/framework/command.rs
@@ -8,7 +8,7 @@ use std::collections::HashMap;
 pub type Check = Fn(&Context, &Message) -> bool + Send + Sync + 'static;
 pub type Exec = Fn(&Context, &Message, Vec<String>) -> Result<(), String> + Send + Sync + 'static;
 pub type Help = Fn(&Context, &Message, HashMap<String, Arc<CommandGroup>>, Vec<String>) -> Result<(), String> + Send + Sync + 'static;
-pub type Hook = Fn(&Context, &Message, &String) + Send + Sync + 'static;
+pub type BeforeHook = Fn(&Context, &Message, &String) -> bool + Send + Sync + 'static;
 pub type AfterHook = Fn(&Context, &Message, &String, Result<(), String>) + Send + Sync + 'static;
 #[doc(hidden)]
 pub type InternalCommand = Arc<Command>;
@@ -55,6 +55,8 @@ pub struct Command {
     pub dm_only: bool,
     /// Whether command can be used only in guilds or not.
     pub guild_only: bool,
+    /// Whether command can only be used by owners or not.
+    pub owners_only: bool,
 }
 
 impl Command {
@@ -73,6 +75,7 @@ impl Command {
             min_args: None,
             max_args: None,
             required_permissions: Permissions::empty(),
+            owners_only: false,
         }
     }
 }
diff --git a/src/ext/framework/configuration.rs b/src/ext/framework/configuration.rs
index cd48e44..404b248 100644
--- a/src/ext/framework/configuration.rs
+++ b/src/ext/framework/configuration.rs
@@ -1,7 +1,8 @@
+use std::collections::HashSet;
 use std::default::Default;
 use super::command::PrefixCheck;
-use ::client::rest;
-use ::client::Context;
+use ::client::{Context, rest};
+use ::model::{GuildId, UserId};
 
 /// Account type used for configuration.
 pub enum AccountType {
@@ -59,11 +60,27 @@ pub struct Configuration {
     #[doc(hidden)]
     pub no_guild_message: Option<String>,
     #[doc(hidden)]
+    pub blocked_user_message: Option<String>,
+    #[doc(hidden)]
+    pub blocked_guild_message: Option<String>,
+    #[doc(hidden)]
     pub too_many_args_message: Option<String>,
     #[doc(hidden)]
     pub not_enough_args_message: Option<String>,
     #[doc(hidden)]
+    pub command_disabled_message: Option<String>,
+    #[doc(hidden)]
     pub account_type: AccountType,
+    #[doc(hidden)]
+    pub blocked_users: HashSet<UserId>,
+    #[doc(hidden)]
+    pub blocked_guilds: HashSet<GuildId>,
+    #[doc(hidden)]
+    pub owners: HashSet<UserId>,
+    #[doc(hidden)]
+    pub disabled_commands: HashSet<String>,
+    #[doc(hidden)]
+    pub allow_dm: bool,
 }
 
 impl Configuration {
@@ -98,13 +115,27 @@ impl Configuration {
     ///
     /// // bot processes and executes the "about" command if it exists
     /// ```
-    pub fn allow_whitespace(mut self, allow_whitespace: bool)
-        -> Self {
+    pub fn allow_whitespace(mut self, allow_whitespace: bool) -> Self {
         self.allow_whitespace = allow_whitespace;
 
         self
     }
 
+    /// HashSet of user Ids whose commands will be ignored.
+    /// Guilds owned by user Ids will also be ignored.
+    pub fn blocked_users(mut self, users: HashSet<UserId>) -> Self {
+        self.blocked_users = users;
+
+        self
+    }
+
+    /// HashSet of guild Ids where commands will be ignored.
+    pub fn blocked_guilds(mut self, guilds: HashSet<GuildId>) -> Self {
+        self.blocked_guilds = guilds;
+
+        self
+    }
+
     /// The default depth of the message to check for commands. Defaults to 5.
     /// This determines how "far" into a message to check for a valid command.
     ///
@@ -119,7 +150,14 @@ impl Configuration {
         self
     }
 
-    /// Sets the prefix to respond to. This can either be a single-char or
+    /// HashSet of command names that won't be run.
+    pub fn disabled_commands(mut self, commands: HashSet<String>) -> Self {
+        self.disabled_commands = commands;
+
+        self
+    }
+
+    /// Sets the prefix to respond to. This can either be a single- or
     /// multi-char string.
     pub fn dynamic_prefix<F>(mut self, dynamic_prefix: F) -> Self
         where F: Fn(&Context) -> Option<String> + Send + Sync + 'static {
@@ -135,6 +173,30 @@ impl Configuration {
         self
     }
 
+    /// Message that's sent if a command is disabled.
+    ///
+    /// %command% will be replaced with command name.
+    pub fn command_disabled_message(mut self, content: &str) -> Self {
+        self.command_disabled_message = Some(content.to_owned());
+
+        self
+    }
+
+    /// Message that's sent if a blocked user attempts to use a command.
+    pub fn blocked_user_message(mut self, content: &str) -> Self {
+        self.blocked_user_message = Some(content.to_owned());
+
+        self
+    }
+
+    /// Message that's sent if a command issued within a guild owned by a
+    /// blocked user.
+    pub fn blocked_guild_message(mut self, content: &str) -> Self {
+        self.blocked_guild_message = Some(content.to_owned());
+
+        self
+    }
+
     /// Message that's sent when a user with wrong permissions calls a command.
     pub fn invalid_permission_message(mut self, content: &str) -> Self {
         self.invalid_permission_message = Some(content.to_owned());
@@ -213,6 +275,13 @@ impl Configuration {
         self
     }
 
+    /// HashSet of user Ids checks won't apply to.
+    pub fn owners(mut self, user_ids: HashSet<UserId>) -> Self {
+        self.owners = user_ids;
+
+        self
+    }
+
     /// Sets the prefix to respond to. This can either be a single-char or
     /// multi-char string.
     pub fn prefix(mut self, prefix: &str) -> Self {
@@ -221,8 +290,8 @@ impl Configuration {
         self
     }
 
-    /// Sets the prefix to respond to. This can either be a single-char or
-    /// multi-char string.
+    /// Sets the prefixes to respond to. Those can either be single-chararacter or
+    /// multi-chararacter strings.
     pub fn prefixes(mut self, prefixes: Vec<&str>) -> Self {
         self.prefixes = prefixes.iter().map(|x| x.to_string()).collect();
 
@@ -252,17 +321,25 @@ impl Default for Configuration {
         Configuration {
             depth: 5,
             on_mention: None,
+            dynamic_prefix: None,
             allow_whitespace: false,
             prefixes: vec![],
-            dynamic_prefix: None,
-            rate_limit_message: None,
-            invalid_permission_message: None,
-            invalid_check_message: None,
             no_dm_message: None,
             no_guild_message: None,
+            rate_limit_message: None,
+            blocked_user_message: None,
             too_many_args_message: None,
+            invalid_check_message: None,
+            blocked_guild_message: None,
             not_enough_args_message: None,
-            account_type: AccountType::Automatic
+            command_disabled_message: None,
+            invalid_permission_message: None,
+            account_type: AccountType::Automatic,
+            owners: HashSet::default(),
+            blocked_users: HashSet::default(),
+            blocked_guilds: HashSet::default(),
+            disabled_commands: HashSet::default(),
+            allow_dm: true,
         }
     }
 }
diff --git a/src/ext/framework/create_command.rs b/src/ext/framework/create_command.rs
index 7a1d744..99505e5 100644
--- a/src/ext/framework/create_command.rs
+++ b/src/ext/framework/create_command.rs
@@ -140,9 +140,17 @@ impl CreateCommand {
         self
     }
 
-    /// Maximum amount of arguments that can be passed.
-    pub fn required_permissions(mut self, required_permissions: Permissions) -> Self {
-        self.0.required_permissions = required_permissions;
+    /// Whether command can be used only privately or not.
+    pub fn owners_only(mut self, owners_only: bool) -> Self {
+        self.0.owners_only = owners_only;
+
+        self
+    }
+
+    /// The permissions that a user must have in the contextual channel in order
+    /// for the command to be processed.
+    pub fn required_permissions(mut self, permissions: Permissions) -> Self {
+        self.0.required_permissions = permissions;
 
         self
     }
@@ -186,7 +194,8 @@ impl Default for Command {
             required_permissions: Permissions::empty(),
             dm_only: false,
             guild_only: false,
-            help_available: true
+            help_available: true,
+            owners_only: false
         }
     }
 }
diff --git a/src/ext/framework/mod.rs b/src/ext/framework/mod.rs
index 796f1a3..e08140d 100644
--- a/src/ext/framework/mod.rs
+++ b/src/ext/framework/mod.rs
@@ -67,7 +67,7 @@ pub use self::configuration::{AccountType, Configuration};
 pub use self::create_command::CreateCommand;
 pub use self::create_group::CreateGroup;
 
-use self::command::{AfterHook, Hook};
+use self::command::{AfterHook, BeforeHook};
 use std::collections::HashMap;
 use std::default::Default;
 use std::sync::Arc;
@@ -163,7 +163,7 @@ macro_rules! command {
 pub struct Framework {
     configuration: Configuration,
     groups: HashMap<String, Arc<CommandGroup>>,
-    before: Option<Arc<Hook>>,
+    before: Option<Arc<BeforeHook>>,
     buckets: HashMap<String, Bucket>,
     after: Option<Arc<AfterHook>>,
     /// Whether the framework has been "initialized".
@@ -321,43 +321,106 @@ impl Framework {
                     };
 
                     if let Some(command) = group.commands.get(&to_check) {
-                        if let Some(ref bucket_name) = command.bucket {
-                            let rate_limit = self.ratelimit_time(bucket_name, message.author.id.0);
-
-                            if rate_limit > 0 {
-                                if let Some(ref message) = self.configuration.rate_limit_message {
-                                    let _ = context.say(
-                                        &message.replace("%time%", &rate_limit.to_string()));
+                        let is_owner = self.configuration.owners.contains(&message.author.id);
+                        // Most of the checks don't apply to owners.
+                        if !is_owner {
+                            if !self.configuration.allow_dm {
+                                if let Some(ref message) = self.configuration.no_dm_message {
+                                    let _ = context.say(message);
                                 }
 
                                 return;
                             }
-                        }
 
-                        if message.is_private() {
-                            if command.guild_only {
-                                if let Some(ref message) = self.configuration.no_guild_message {
+                            if let Some(guild_id) = message.guild_id() {
+                                if self.configuration.blocked_guilds.contains(&guild_id) {
+                                    if let Some(ref message) = self.configuration.blocked_guild_message {
+                                        let _ = context.say(message);
+                                    }
+
+                                    return;
+                                }
+                            }
+
+                            #[cfg(feature="cache")]
+                            {
+                                if let Some(guild_id) = message.guild_id() {
+                                    if let Some(guild) = guild_id.find() {
+                                        if self.configuration.blocked_users.contains(&guild.owner_id) {
+                                            if let Some(ref message) = self.configuration.blocked_guild_message {
+                                                let _ = context.say(message);
+                                            }
+
+                                            return;
+                                        }
+                                    }
+                                }
+                            }
+
+                            if self.configuration.blocked_users.contains(&message.author.id) {
+                                if let Some(ref message) = self.configuration.blocked_user_message {
                                     let _ = context.say(message);
                                 }
 
                                 return;
                             }
-                        } else if command.dm_only {
-                            if let Some(ref message) = self.configuration.no_dm_message {
-                                let _ = context.say(message);
+
+                            if self.configuration.disabled_commands.contains(&to_check) ||
+                               self.configuration.disabled_commands.contains(&built) {
+                                if let Some(ref message) = self.configuration.command_disabled_message {
+                                    let msg = message.replace("%command%", &to_check);
+
+                                    let _ = context.say(&msg);
+                                }
+
+                                return;
                             }
 
-                            return;
-                        }
+                            if let Some(ref bucket_name) = command.bucket {
+                                let rate_limit = self.ratelimit_time(bucket_name, message.author.id.0);
+
+                                if rate_limit > 0 {
+                                    if let Some(ref message) = self.configuration.rate_limit_message {
+                                        let msg = message.replace("%time%", &rate_limit.to_string());
 
-                        for check in &command.checks {
-                            if !(check)(&context, &message) {
-                                if let Some(ref message) = self.configuration.invalid_check_message {
+                                        let _ = context.say(&msg);
+                                    }
+
+                                    return;
+                                }
+                            }
+
+                            if message.is_private() {
+                                if command.guild_only {
+                                    if let Some(ref message) = self.configuration.no_guild_message {
+                                        let _ = context.say(message);
+                                    }
+
+                                    return;
+                                }
+                            } else if command.dm_only {
+                                if let Some(ref message) = self.configuration.no_dm_message {
                                     let _ = context.say(message);
                                 }
 
-                                continue 'outer;
+                                return;
                             }
+
+                            for check in &command.checks {
+                                if !(check)(&context, &message) {
+                                    if let Some(ref message) = self.configuration.invalid_check_message {
+                                        let _ = context.say(message);
+                                    }
+
+                                    continue 'outer;
+                                }
+                            }
+                        } else if command.owners_only {
+                            if let Some(ref message) = self.configuration.invalid_permission_message {
+                                let _ = context.say(message);
+                            }
+
+                            return;
                         }
 
                         let before = self.before.clone();
@@ -377,9 +440,10 @@ impl Framework {
                         if let Some(x) = command.min_args {
                             if args.len() < x as usize {
                                 if let Some(ref message) = self.configuration.not_enough_args_message {
-                                    let _ = context.say(
-                                        &message.replace("%min%", &x.to_string())
-                                                .replace("%given%", &args.len().to_string()));
+                                    let msg = message.replace("%min%", &x.to_string())
+                                        .replace("%given%", &args.len().to_string());
+
+                                    let _ = context.say(&msg);
                                 }
 
                                 return;
@@ -389,16 +453,17 @@ impl Framework {
                         if let Some(x) = command.max_args {
                             if args.len() > x as usize {
                                 if let Some(ref message) = self.configuration.too_many_args_message {
-                                    let _ = context.say(
-                                        &message.replace("%max%", &x.to_string())
-                                                .replace("%given%", &args.len().to_string()));
+                                    let msg = message.replace("%max%", &x.to_string())
+                                        .replace("%given%", &args.len().to_string());
+
+                                    let _ = context.say(&msg);
                                 }
 
                                 return;
                             }
                         }
 
-                        if !command.required_permissions.is_empty() {
+                        if !is_owner && !command.required_permissions.is_empty() {
                             let mut permissions_fulfilled = false;
 
                             if let Some(member) = message.get_member() {
@@ -424,7 +489,9 @@ impl Framework {
 
                         thread::spawn(move || {
                             if let Some(before) = before {
-                                (before)(&context, &message, &built);
+                                if !is_owner && !(before)(&context, &message, &built) {
+                                    return;
+                                }
                             }
 
                             let result = match command.exec {
@@ -530,8 +597,9 @@ impl Framework {
     }
 
     /// Specify the function to be called prior to every command's execution.
+    /// If that function returns true, the command will be executed.
     pub fn before<F>(mut self, f: F) -> Self
-        where F: Fn(&Context, &Message, &String) + Send + Sync + 'static {
+        where F: Fn(&Context, &Message, &String) -> bool + Send + Sync + 'static {
         self.before = Some(Arc::new(f));
 
         self