support for DMing commands. execute command function correctly

This commit is contained in:
jude 2020-08-07 23:45:34 +01:00
parent 305cf79ac8
commit 0bda479d01

View File

@ -2,7 +2,12 @@ use async_trait::async_trait;
use serenity::{ use serenity::{
client::Context, client::Context,
framework::Framework, framework::{
Framework,
standard::{
Args, CommandFn,
},
},
model::{ model::{
guild::Guild, guild::Guild,
channel::{ channel::{
@ -27,7 +32,6 @@ use std::{
fmt, fmt,
}; };
use serenity::framework::standard::CommandFn;
use crate::SQLPool; use crate::SQLPool;
#[derive(Debug)] #[derive(Debug)]
@ -60,6 +64,7 @@ impl fmt::Debug for Command {
pub struct RegexFramework { pub struct RegexFramework {
commands: HashMap<String, &'static Command>, commands: HashMap<String, &'static Command>,
regex_matcher: Regex, regex_matcher: Regex,
dm_regex_matcher: Regex,
default_prefix: String, default_prefix: String,
client_id: u64, client_id: u64,
ignore_bots: bool, ignore_bots: bool,
@ -70,6 +75,7 @@ impl RegexFramework {
Self { Self {
commands: HashMap::new(), commands: HashMap::new(),
regex_matcher: Regex::new(r#"^$"#).unwrap(), regex_matcher: Regex::new(r#"^$"#).unwrap(),
dm_regex_matcher: Regex::new(r#"^$"#).unwrap(),
default_prefix: String::from("$"), default_prefix: String::from("$"),
client_id, client_id,
ignore_bots: true, ignore_bots: true,
@ -95,26 +101,59 @@ impl RegexFramework {
} }
pub fn build(mut self) -> Self { pub fn build(mut self) -> Self {
let command_names;
{ {
let mut command_names_vec = self.commands let command_names;
.keys()
.map(|k| &k[..])
.collect::<Vec<&str>>();
command_names_vec.sort_unstable_by(|a, b| b.len().cmp(&a.len())); {
let mut command_names_vec = self.commands
.keys()
.map(|k| &k[..])
.collect::<Vec<&str>>();
command_names = command_names_vec.join("|"); command_names_vec.sort_unstable_by(|a, b| b.len().cmp(&a.len()));
command_names = command_names_vec.join("|");
}
info!("Command names: {}", command_names);
{
let match_string = r#"^(?:(?:<@ID>\s+)|(?:<@!ID>\s+)|(?P<prefix>\S{1,5}?))(?P<cmd>COMMANDS)(?:$|\s+(?P<args>.*))$"#
.replace("COMMANDS", command_names.as_str())
.replace("ID", self.client_id.to_string().as_str());
self.regex_matcher = Regex::new(match_string.as_str()).unwrap();
}
} }
info!("Command names: {}", command_names); {
let dm_command_names;
let match_string = r#"^(?:(?:<@ID>\s+)|(?:<@!ID>\s+)|(?P<prefix>\S{1,5}?))(?P<cmd>COMMANDS)(?:$|\s+(?P<args>.*))$"# {
.replace("COMMANDS", command_names.as_str()) let mut command_names_vec = self.commands
.replace("ID", self.client_id.to_string().as_str()); .iter()
.filter_map(|(key, command)| {
if command.supports_dm {
Some(&key[..])
} else {
None
}
})
.collect::<Vec<&str>>();
self.regex_matcher = Regex::new(match_string.as_str()).unwrap(); command_names_vec.sort_unstable_by(|a, b| b.len().cmp(&a.len()));
dm_command_names = command_names_vec.join("|");
}
{
let match_string = r#"^(?:(?:<@ID>\s+)|(?:<@!ID>\s+)|(\$)|())(?P<cmd>COMMANDS)(?:$|\s+(?P<args>.*))$"#
.replace("COMMANDS", dm_command_names.as_str())
.replace("ID", self.client_id.to_string().as_str());
self.dm_regex_matcher = Regex::new(match_string.as_str()).unwrap();
}
}
self self
} }
@ -130,7 +169,7 @@ enum PermissionCheck {
impl Framework for RegexFramework { impl Framework for RegexFramework {
async fn dispatch(&self, ctx: Context, msg: Message) { async fn dispatch(&self, ctx: Context, msg: Message) {
async fn check_self_permissions(ctx: &Context, guild: &Guild, channel: &GuildChannel) -> Result<PermissionCheck, Box<dyn std::error::Error>> { async fn check_self_permissions(ctx: &Context, guild: &Guild, channel: &GuildChannel) -> Result<PermissionCheck, Box<dyn std::error::Error + Sync + Send>> {
let user_id = ctx.cache.current_user_id().await; let user_id = ctx.cache.current_user_id().await;
let guild_perms = guild.member_permissions(user_id); let guild_perms = guild.member_permissions(user_id);
@ -149,12 +188,12 @@ impl Framework for RegexFramework {
}) })
} }
async fn check_prefix(ctx: &Context, guild_id: u64, prefix_opt: Option<Match<'_>>) -> bool { async fn check_prefix(ctx: &Context, guild: &Guild, prefix_opt: Option<Match<'_>>) -> bool {
if let Some(prefix) = prefix_opt { if let Some(prefix) = prefix_opt {
let pool = ctx.data.read().await let pool = ctx.data.read().await
.get::<SQLPool>().cloned().expect("Could not get SQLPool from data"); .get::<SQLPool>().cloned().expect("Could not get SQLPool from data");
match sqlx::query!("SELECT prefix FROM guilds WHERE id = ?", guild_id) match sqlx::query!("SELECT prefix FROM guilds WHERE id = ?", guild.id.as_u64())
.fetch_one(&pool) .fetch_one(&pool)
.await { .await {
Ok(row) => { Ok(row) => {
@ -162,6 +201,10 @@ impl Framework for RegexFramework {
} }
Err(sqlx::Error::RowNotFound) => { Err(sqlx::Error::RowNotFound) => {
sqlx::query!("INSERT INTO guilds (guild, name) VALUES (?, ?)", guild.id.as_u64(), guild.name)
.execute(&pool)
.await;
prefix.as_str() == "$" prefix.as_str() == "$"
} }
@ -191,15 +234,27 @@ impl Framework for RegexFramework {
if let Some(full_match) = self.regex_matcher.captures(&msg.content[..]) { if let Some(full_match) = self.regex_matcher.captures(&msg.content[..]) {
if check_prefix(&ctx, *guild.id.as_u64(), full_match.name("prefix")).await { if check_prefix(&ctx, &guild, full_match.name("prefix")).await {
debug!("Prefix matched on {}", msg.content); debug!("Prefix matched on {}", msg.content);
match check_self_permissions(&ctx, &guild, &channel).await { match check_self_permissions(&ctx, &guild, &channel).await {
Ok(perms) => match perms { Ok(perms) => match perms {
PermissionCheck::All => {} PermissionCheck::All => {
let command = self.commands.get(full_match.name("cmd").unwrap().as_str()).unwrap();
let args = Args::new(
full_match.name("args")
.map(|m| m.as_str())
.unwrap_or(""),
&[]
);
PermissionCheck::Basic => {} (command.func)(&ctx, &msg, args).await;
}
PermissionCheck::Basic => {
msg.channel_id.say(&ctx, "Not enough perms").await;
}
PermissionCheck::None => { PermissionCheck::None => {
warn!("Missing enough permissions for guild {}", guild.id); warn!("Missing enough permissions for guild {}", guild.id);
@ -216,7 +271,17 @@ impl Framework for RegexFramework {
// DM Command // DM Command
else { else {
if let Some(full_match) = self.dm_regex_matcher.captures(&msg.content[..]) {
let command = self.commands.get(full_match.name("cmd").unwrap().as_str()).unwrap();
let args = Args::new(
full_match.name("args")
.map(|m| m.as_str())
.unwrap_or(""),
&[]
);
(command.func)(&ctx, &msg, args).await;
}
} }
} }
} }