diff options
| -rw-r--r-- | src/bot.rs | 205 | ||||
| -rw-r--r-- | src/config.rs | 10 | ||||
| -rw-r--r-- | src/log_setup.rs (renamed from src/log.rs) | 0 | ||||
| -rw-r--r-- | src/main.rs | 191 | ||||
| -rw-r--r-- | src/util.rs | 7 |
5 files changed, 223 insertions, 190 deletions
diff --git a/src/bot.rs b/src/bot.rs new file mode 100644 index 0000000..3923458 --- /dev/null +++ b/src/bot.rs @@ -0,0 +1,205 @@ +use std::{ + sync::Mutex, + fs::File, + result::Result as StdResult, +}; + +use serenity::{ + prelude::*, + model::{ + gateway::Ready, + id::{ + ChannelId, + MessageId, + }, + channel::Message, + }, + framework::StandardFramework, +}; + +use fnv::{ + FnvHashMap, + FnvHashSet, +}; + +use chrono::Datelike; +use lazy_static::lazy_static; +use log::{ + debug, + info, + error, + trace, + warn, +}; + +use crate::{ + Result, + Error, + audio, + util::OAUTH_URL, + util::CtxExt, + commands::register_commands, + config::CONFIG, +}; + +struct Handler; +impl EventHandler for Handler { + fn ready(&self, ctx: Context, r: Ready) { + let guild = r.guilds.iter() + .find(|g| g.id() == CONFIG.discord.guild()); + + if guild.is_none() { + info!("bot isn't in configured guild. join here: {:?}", OAUTH_URL.as_str()); + } + + #[cfg(debug_assertions)] + let botname = "thulani (dev)"; + + #[cfg(not(debug_assertions))] + let botname = "thulani"; + + guild.iter().for_each(|g| { + if let Err(e) = g.id().edit_nickname(&ctx, Some(botname)) { + error!("changing nickname: {:?}", e); + } + }); + } + + fn message_delete(&self, _ctx: Context, _channel_id: ChannelId, deleted_message_id: MessageId) { + MESSAGE_WATCH.lock() + .unwrap() + .remove(&deleted_message_id); + } +} + +lazy_static! { + static ref MESSAGE_WATCH: Mutex<FnvHashMap<MessageId, MessageId>> = Mutex::new(FnvHashMap::default()); + static ref PREFIXES: Vec<&'static str> = vec!["!thulani ", "!thulan ", "!thulando madando ", "!thulando "]; + static ref RESTRICTED_PREFIXES: Vec<&'static str> = vec!["!todd ", "!toddbert ", "!toddlani "]; + static ref ALL_PREFIXES: Vec<&'static str> = { + let mut all_prefixes: Vec<&'static str> = vec![]; + all_prefixes.extend(PREFIXES.iter()); + all_prefixes.extend(RESTRICTED_PREFIXES.iter()); + all_prefixes + }; + + static ref RESTRICT_IDS: FnvHashSet<u64> = { + let restrict_ids = File::open("restrict.json") + .map_err(Error::from) + .and_then(|f| serde_json::from_reader::<_, Vec<u64>>(f).map_err(Error::from)); + + if let Err(ref e) = restrict_ids { + warn!("opening restrict file: {}", e); + } + + restrict_ids + .unwrap_or_default() + .into_iter() + .collect::<FnvHashSet<_>>() + }; +} + +fn framework() -> StandardFramework { + let framework = StandardFramework::new() + .configure(|c| c + .allow_dm(false) + .with_whitespace(true) + .prefixes(ALL_PREFIXES.iter()) + .ignore_bots(true) + .on_mention(None) + .owners(vec![CONFIG.discord.owner()].into_iter().collect()) + .case_insensitivity(true) + ) + .before(before_handle) + .after(after_handle) + .bucket("Standard", |b| { + b.delay(1).limit(20).time_span(60) + }); + + register_commands(framework) +} + +fn before_handle(ctx: &mut Context, message: &Message, cmd: &str) -> bool { + debug!("got command '{}' from user '{}' ({})", cmd, message.author.name, message.author.id); + + if !message.guild_id.map_or(false, |x| x == CONFIG.discord.guild()) { + info!("rejecting command '{}' from user '{}': wrong guild", cmd, message.author.name); + return false; + } + + if message.author.id == CONFIG.discord.owner() { + return true; + } + + let restricted_prefix = RESTRICTED_PREFIXES.iter() + .any(|prefix| message.content.starts_with(prefix)); + + if !restricted_prefix { + return true; + } + + const PERMITTED_WEEKDAY: chrono::Weekday = chrono::Weekday::Tue; + + let user_is_restricted = RESTRICT_IDS.contains(&message.author.id.0); + let restrictions_flipped = chrono::Local::now().weekday() == PERMITTED_WEEKDAY; + + if user_is_restricted == restrictions_flipped { + return true; + } + + let reason = if !restrictions_flipped { + "restricted prefix".to_owned() + } else { + format!("it is {:?}", PERMITTED_WEEKDAY) + }; + + info!("rejecting command '{}' from user '{}': {}", cmd, message.author.name, reason); + + match ctx.send_result(message.channel_id, "no", message.tts) { + Err(e) => error!("sending restricted prefix response: {}", e), + Ok(msg_id) => { + let mut mp = MESSAGE_WATCH.lock().unwrap(); + mp.insert(message.id, msg_id); + } + } + + false +} + +fn after_handle(ctx: &mut Context, msg: &Message, cmd: &str, err: StdResult<(), Error>) { + match err { + Ok(()) => { + trace!("command '{}' completed successfully", cmd); + }, + + Err(e) => { + if let Err(e) = msg.react(&ctx, "❌") { + error!("reacting to failed message: {}", e); + } + + if let Err(e) = ctx.send(msg.channel_id, "BANIC", msg.tts) { + error!("sending BANIC: {}", e); + } + + error!("error encountered handling command '{}': {:?}", cmd, e); + } + } +} + +pub fn run() -> Result<()> { + let token = &CONFIG.discord.auth.token; + let mut client = Client::new(token, Handler)?; + + audio::VoiceManager::register(&mut client); + audio::PlayQueue::register(&mut client); + + client.with_framework(framework()); + + let shard_manager = client.shard_manager.clone(); + ctrlc::set_handler(move || { + info!("shutting down"); + shard_manager.lock().shutdown_all(); + }).expect("unable to create SIGINT/SIGTERM handlers"); + + client.start() +} diff --git a/src/config.rs b/src/config.rs index 13d015e..53801b6 100644 --- a/src/config.rs +++ b/src/config.rs @@ -6,8 +6,18 @@ use serenity::{ }, }; +use dotenv::dotenv; +use lazy_static::lazy_static; use envconfig::Envconfig; +lazy_static! { + pub static ref CONFIG: Config = { + dotenv().ok(); + + Config::init().unwrap() + }; +} + #[derive(Envconfig)] pub struct Config { #[envconfig(from = "DATABASE_URL")] diff --git a/src/log.rs b/src/log_setup.rs index a51ebcc..a51ebcc 100644 --- a/src/log.rs +++ b/src/log_setup.rs diff --git a/src/main.rs b/src/main.rs index daafbf4..6c2d651 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,8 +12,6 @@ #[macro_use] extern crate envconfig_derive; use std::{ - default::Default, - fs::File, thread, time::{ Duration, @@ -21,29 +19,10 @@ use std::{ }, }; -use chrono::Datelike; -use fnv::{FnvHashMap, FnvHashSet}; use log::{ - debug, error, info, - trace, - warn, }; -use serenity::{ - framework::StandardFramework, - model::{ - gateway::Ready, - id::{ChannelId, MessageId}, - }, - prelude::*, -}; - -use dotenv::dotenv; -use lazy_static::lazy_static; -use envconfig::Envconfig; - -use self::commands::register_commands; pub use self::util::*; pub use self::config::*; @@ -68,174 +47,12 @@ mod commands; mod util; mod audio; mod config; -mod log; +mod log_setup; +mod bot; pub type Error = anyhow::Error; - pub type Result<T> = anyhow::Result<T>; -lazy_static! { - pub static ref CONFIG: Config = { - dotenv().ok(); - - Config::init().unwrap() - }; -} - -struct Handler; -impl EventHandler for Handler { - fn ready(&self, ctx: Context, r: Ready) { - let guild = r.guilds.iter() - .find(|g| g.id() == CONFIG.discord.guild()); - - if guild.is_none() { - info!("bot isn't in configured guild. join here: {:?}", OAUTH_URL.as_str()); - } - - #[cfg(debug_assertions)] { - let _ = guild.map(|g| g.id().edit_nickname(ctx, Some("thulani (dev)"))); - } - - #[cfg(not(debug_assertions))] { - let _ = guild.map(|g| g.id().edit_nickname(ctx, Some("thulani"))); - } - } - - fn message_delete(&self, ctx: Context, channel_id: ChannelId, deleted_message_id: MessageId) { - MESSAGE_WATCH.lock() - .remove(&deleted_message_id) - .iter() - .for_each(|id| { - if let Err(e) = channel_id.delete_message(&ctx, id) { - error!("deleting message: {}", e); - } - }); - } -} - -lazy_static! { - static ref MESSAGE_WATCH: Mutex<FnvHashMap<MessageId, MessageId>> = Mutex::new(FnvHashMap::default()); - static ref PREFIXES: Vec<&'static str> = vec!["!thulani ", "!thulan ", "!thulando madando ", "!thulando "]; - static ref RESTRICTED_PREFIXES: Vec<&'static str> = vec!["!todd ", "!toddbert ", "!toddlani "]; -} - - -fn run() -> Result<()> { - let token = &CONFIG.discord.auth.token; - let mut client = Client::new(token, Handler)?; - - audio::VoiceManager::register(&mut client); - audio::PlayQueue::register(&mut client); - - let all_prefixes = { - let mut all_prefixes: Vec<&'static str> = vec![]; - all_prefixes.extend(PREFIXES.iter()); - all_prefixes.extend(RESTRICTED_PREFIXES.iter()); - all_prefixes - }; - - let restrict_ids = File::open("restrict.json") - .map_err(Error::from) - .and_then(|f| serde_json::from_reader::<_, Vec<u64>>(f).map_err(Error::from)); - - if let Err(ref e) = restrict_ids { - warn!("opening restrict file: {}", e); - } - - let restrict_ids = restrict_ids - .unwrap_or_default() - .into_iter() - .collect::<FnvHashSet<_>>(); - - let mut framework = StandardFramework::new() - .configure(|c| c - .allow_dm(false) - .with_whitespace(true) - .prefixes(all_prefixes) - .ignore_bots(true) - .on_mention(None) - .owners(vec![CONFIG.discord.owner()].into_iter().collect()) - .case_insensitivity(true) - ) - .before(move |ctx, message, cmd| { - debug!("got command '{}' from user '{}' ({})", cmd, message.author.name, message.author.id); - if !message.guild_id.map_or(false, |x| x == CONFIG.discord.guild()) { - info!("rejecting command '{}' from user '{}': wrong guild", cmd, message.author.name); - return false; - } - - if message.author.id == CONFIG.discord.owner() { - return true; - } - - let restricted_prefix = RESTRICTED_PREFIXES.iter().any(|prefix| message.content.starts_with(prefix)); - if !restricted_prefix { - return true; - } - - const PERMITTED_WEEKDAY: chrono::Weekday = chrono::Weekday::Tue; - - let restricted_user = restrict_ids.contains(&message.author.id.0); - let flip_restriction_day = chrono::Local::now().weekday() == PERMITTED_WEEKDAY; - - if restricted_user == flip_restriction_day { - return true; - } - - let reason = if !flip_restriction_day { - "restricted prefix".to_owned() - } else { - format!("it is {:?}", PERMITTED_WEEKDAY) - }; - - info!("rejecting command '{}' from user '{}': {}", cmd, message.author.name, reason); - - match ctx.send_result(message.channel_id, "no", message.tts) { - Err(e) => error!("sending restricted prefix response: {}", e), - Ok(msg_id) => { - let mut mp = MESSAGE_WATCH.lock(); - mp.insert(message.id, msg_id); - } - } - - return false; - }) - .after(|ctx, msg, cmd, err| { - match err { - Ok(()) => { - trace!("command '{}' completed successfully", cmd); - }, - - Err(e) => { - if let Err(e) = msg.react(&ctx, "❌") { - error!("reacting to failed message: {}", e); - } - - if let Err(e) = ctx.send(msg.channel_id, "BANIC", msg.tts) { - error!("sending BANIC: {}", e); - } - - error!("error encountered handling command '{}': {:?}", cmd, e); - } - } - }) - .bucket("Standard", |b| b.delay(1).limit(20).time_span(60)); - - framework = register_commands(framework); - - client.with_framework(framework); - - let shard_manager = client.shard_manager.clone(); - ctrlc::set_handler(move || { - info!("shutting down"); - shard_manager.lock().shutdown_all(); - }).expect("unable to create SIGINT/SIGTERM handlers"); - - client.start()?; - - Ok(()) -} - fn main() { const BACKOFF_FACTOR: f64 = 2.0; const MAX_BACKOFFS: usize = 3; @@ -245,7 +62,7 @@ fn main() { info!("starting"); - log::init().expect("initializing logging"); + log_setup::init().expect("initializing logging"); let mut backoff_count: usize = 0; @@ -253,7 +70,7 @@ fn main() { let start = Instant::now(); info!("starting bot"); - match run() { + match bot::run() { Err(e) => { error!("error encountered running client: {:?}", e); }, diff --git a/src/util.rs b/src/util.rs index 77f2dd4..e9b6203 100644 --- a/src/util.rs +++ b/src/util.rs @@ -8,9 +8,10 @@ use serenity::{ permissions::Permissions,
}
};
-use url::Url;
+use url::Url;
use lazy_static::lazy_static;
+use log::debug;
use crate::{
CONFIG,
@@ -52,8 +53,8 @@ impl CtxExt for Context { #[inline]
fn send_result<A: AsRef<str>>(&self, channel: ChannelId, text: A, tts: bool) -> Result<MessageId> {
let text = text.as_ref();
- debug!("sending message {:?} to channel {:?}", text, channel);
- let result = channel.send_message(self, |m| m.content(text.as_ref()).tts(tts))?;
+ debug!("sending message {:?} to channel {:?} (tts: {})", text, channel, tts);
+ let result = channel.send_message(self, |m| m.content(text).tts(tts))?;
Ok(result.id)
}
}
|