diff options
| author | Jokler <jokler.contact@gmail.com> | 2018-03-12 16:02:51 +0100 |
|---|---|---|
| committer | Jokler <jokler.contact@gmail.com> | 2018-03-12 16:02:51 +0100 |
| commit | 909cabe9280722e43c5fb283f768051bb85e1890 (patch) | |
| tree | 506ac34b7e22cdb95568cef9e649ee64cb3b0fdb /src | |
| parent | 15e855ddecfdac31ddda26b12fcfd1a142a0ec21 (diff) | |
| parent | 8e40e919aca8b8592be43e2c5bbcc0717bf14a6b (diff) | |
| download | frippy-909cabe9280722e43c5fb283f768051bb85e1890.tar.gz frippy-909cabe9280722e43c5fb283f768051bb85e1890.zip | |
Merge branch 'dev'
Diffstat (limited to 'src')
| -rw-r--r-- | src/error.rs | 31 | ||||
| -rw-r--r-- | src/lib.rs | 339 | ||||
| -rw-r--r-- | src/main.rs | 191 | ||||
| -rw-r--r-- | src/plugin.rs | 167 | ||||
| -rw-r--r-- | src/plugins/currency.rs | 131 | ||||
| -rw-r--r-- | src/plugins/emoji.rs | 62 | ||||
| -rw-r--r-- | src/plugins/factoids/database.rs | 163 | ||||
| -rw-r--r-- | src/plugins/factoids/mod.rs | 342 | ||||
| -rw-r--r-- | src/plugins/factoids/sandbox.lua | 86 | ||||
| -rw-r--r-- | src/plugins/factoids/utils.rs | 25 | ||||
| -rw-r--r-- | src/plugins/help.rs | 32 | ||||
| -rw-r--r-- | src/plugins/keepnick.rs | 70 | ||||
| -rw-r--r-- | src/plugins/mod.rs | 15 | ||||
| -rw-r--r-- | src/plugins/tell/database.rs | 150 | ||||
| -rw-r--r-- | src/plugins/tell/mod.rs | 264 | ||||
| -rw-r--r-- | src/plugins/url.rs | 130 | ||||
| -rw-r--r-- | src/utils.rs | 65 |
17 files changed, 1958 insertions, 305 deletions
diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..36d5724 --- /dev/null +++ b/src/error.rs @@ -0,0 +1,31 @@ +//! Errors for `frippy` crate using `failure`. + +use failure::Fail; + +pub fn log_error(e: FrippyError) { + let text = e.causes() + .skip(1) + .fold(format!("{}", e), |acc, err| format!("{}: {}", acc, err)); + error!("{}", text); +} + +/// The main crate-wide error type. +#[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)] +#[error = "FrippyError"] +pub enum ErrorKind { + /// Connection error + #[fail(display = "A connection error occured")] + Connection, + + /// A Url error + #[fail(display = "A Url error has occured")] + Url, + + /// A Tell error + #[fail(display = "A Tell error has occured")] + Tell, + + /// A Factoids error + #[fail(display = "A Factoids error has occured")] + Factoids, +} @@ -1,135 +1,196 @@ -#![cfg_attr(feature="clippy", feature(plugin))] -#![cfg_attr(feature="clippy", plugin(clippy))] +#![cfg_attr(feature = "clippy", feature(plugin))] +#![cfg_attr(feature = "clippy", plugin(clippy))] //! Frippy is an IRC bot that runs plugins on each message //! received. //! -//! ## Example +//! ## Examples //! ```no_run -//! extern crate frippy; +//! # extern crate irc; +//! # extern crate frippy; +//! # fn main() { +//! use frippy::{plugins, Config, Bot}; +//! use irc::client::reactor::IrcReactor; //! -//! frippy::run(); +//! let config = Config::load("config.toml").unwrap(); +//! let mut reactor = IrcReactor::new().unwrap(); +//! let mut bot = Bot::new(); +//! +//! bot.add_plugin(plugins::help::Help::new()); +//! bot.add_plugin(plugins::emoji::Emoji::new()); +//! bot.add_plugin(plugins::currency::Currency::new()); +//! +//! bot.connect(&mut reactor, &config).unwrap(); +//! reactor.run().unwrap(); +//! # } //! ``` //! //! # Logging //! Frippy uses the [log](https://docs.rs/log) crate so you can log events //! which might be of interest. +#[cfg(feature = "mysql")] #[macro_use] -extern crate log; +extern crate diesel; +#[cfg(feature = "mysql")] +extern crate r2d2; +#[cfg(feature = "mysql")] +extern crate r2d2_diesel; + +#[macro_use] +extern crate failure; +#[macro_use] +extern crate frippy_derive; #[macro_use] -extern crate plugin_derive; +extern crate lazy_static; +#[macro_use] +extern crate log; +extern crate chrono; +extern crate humantime; extern crate irc; -extern crate tokio_core; -extern crate futures; -extern crate glob; +extern crate reqwest; +extern crate time; -mod plugin; -mod plugins; +pub mod plugin; +pub mod plugins; +pub mod utils; +pub mod error; +use std::collections::HashMap; +use std::fmt; +use std::thread::spawn; use std::sync::Arc; -use irc::client::prelude::*; -use irc::error::Error as IrcError; - -use tokio_core::reactor::Core; -use futures::future; -use glob::glob; +pub use irc::client::prelude::*; +pub use irc::error::IrcError; +use error::*; +use failure::ResultExt; use plugin::*; -/// Runs the bot -/// -/// # Remarks -/// -/// This blocks the current thread while the bot is running -pub fn run() { - - // Load all toml files in the configs directory - let mut configs = Vec::new(); - for toml in glob("configs/*.toml").unwrap() { - match toml { - Ok(path) => { - info!("Loading {}", path.to_str().unwrap()); - match Config::load(path) { - Ok(v) => configs.push(v), - Err(e) => error!("Incorrect config file {}", e), - } - } - Err(e) => error!("Failed to read path {}", e), +/// The bot which contains the main logic. +#[derive(Default)] +pub struct Bot { + plugins: ThreadedPlugins, +} + +impl Bot { + /// Creates a `Bot`. + /// By itself the bot only responds to a few simple CTCP commands + /// defined per config file. + /// Any other functionality has to be provided by plugins + /// which need to implement [`Plugin`](plugin/trait.Plugin.html). + /// + /// # Examples + /// ``` + /// use frippy::Bot; + /// let mut bot = Bot::new(); + /// ``` + pub fn new() -> Bot { + Bot { + plugins: ThreadedPlugins::new(), } } - // Without configs the bot would just idle - if configs.is_empty() { - error!("No config file found"); - return; + /// Adds the [`Plugin`](plugin/trait.Plugin.html). + /// These plugins will be used to evaluate incoming messages from IRC. + /// + /// # Examples + /// ``` + /// use frippy::{plugins, Bot}; + /// + /// let mut bot = frippy::Bot::new(); + /// bot.add_plugin(plugins::help::Help::new()); + /// ``` + pub fn add_plugin<T: Plugin + 'static>(&mut self, plugin: T) { + self.plugins.add(plugin); } - // The list of plugins in use - let mut plugins = ThreadedPlugins::new(); - plugins.add(plugins::Help::new()); - plugins.add(plugins::Emoji::new()); - plugins.add(plugins::Currency::new()); - info!("Plugins loaded: {}", plugins); - - // Create an event loop to run the connections on. - let mut reactor = Core::new().unwrap(); - - // Open a connection and add work for each config - for config in configs { - let server = - match IrcServer::new_future(reactor.handle(), &config).and_then(|f| reactor.run(f)) { - Ok(v) => v, - Err(e) => { - error!("Failed to connect: {}", e); - return; - } - }; + /// Removes a [`Plugin`](plugin/trait.Plugin.html) based on its name. + /// The binary currently uses this to disable plugins + /// based on user configuration. + /// + /// # Examples + /// ``` + /// use frippy::{plugins, Bot}; + /// + /// let mut bot = frippy::Bot::new(); + /// bot.add_plugin(plugins::help::Help::new()); + /// bot.remove_plugin("Help"); + /// ``` + pub fn remove_plugin(&mut self, name: &str) -> Option<()> { + self.plugins.remove(name) + } - info!("Connected to server"); + /// This connects the `Bot` to IRC and creates a task on the + /// [`IrcReactor`](../irc/client/reactor/struct.IrcReactor.html) + /// which returns an Ok if the connection was cleanly closed and + /// an Err if the connection was lost. + /// + /// You need to run the [`IrcReactor`](../irc/client/reactor/struct.IrcReactor.html), + /// so that the `Bot` + /// can actually do its work. + /// + /// # Examples + /// ```no_run + /// # extern crate irc; + /// # extern crate frippy; + /// # fn main() { + /// use frippy::{Config, Bot}; + /// use irc::client::reactor::IrcReactor; + /// + /// let config = Config::load("config.toml").unwrap(); + /// let mut reactor = IrcReactor::new().unwrap(); + /// let mut bot = Bot::new(); + /// + /// bot.connect(&mut reactor, &config).unwrap(); + /// reactor.run().unwrap(); + /// # } + /// ``` + pub fn connect(&self, reactor: &mut IrcReactor, config: &Config) -> Result<(), FrippyError> { + info!("Plugins loaded: {}", self.plugins); - match server.identify() { - Ok(_) => info!("Identified"), - Err(e) => error!("Failed to identify: {}", e), - }; + let client = reactor + .prepare_client_and_connect(config) + .context(ErrorKind::Connection)?; + + info!("Connected to IRC server"); + + client.identify().context(ErrorKind::Connection)?; + info!("Identified"); // TODO Verify if we actually need to clone plugins twice - let plugins = plugins.clone(); + let plugins = self.plugins.clone(); - let task = server - .stream() - .for_each(move |message| process_msg(&server, plugins.clone(), message)) - .map_err(|e| error!("Failed to process message: {}", e)); + reactor.register_client_with_handler(client, move |client, message| { + process_msg(client, plugins.clone(), message) + }); - reactor.handle().spawn(task); + Ok(()) } - - // Run the main loop forever - reactor.run(future::empty::<(), ()>()).unwrap(); } -fn process_msg(server: &IrcServer, - mut plugins: ThreadedPlugins, - message: Message) - -> Result<(), IrcError> { - +fn process_msg( + client: &IrcClient, + mut plugins: ThreadedPlugins, + message: Message, +) -> Result<(), IrcError> { + // Log any channels we join if let Command::JOIN(ref channel, _, _) = message.command { - if message.source_nickname().unwrap() == server.current_nickname() { + if message.source_nickname().unwrap() == client.current_nickname() { info!("Joined {}", channel); } } // Check for possible command and save the result for later - let command = PluginCommand::from(&server.current_nickname().to_lowercase(), &message); + let command = PluginCommand::from(&client.current_nickname().to_lowercase(), &message); - let message = Arc::new(message); - plugins.execute_plugins(server, message); + plugins.execute_plugins(client, message); // If the message contained a command, handle it if let Some(command) = command { - if let Err(e) = plugins.handle_command(server, command) { + if let Err(e) = plugins.handle_command(client, command) { error!("Failed to handle command: {}", e); } } @@ -137,6 +198,110 @@ fn process_msg(server: &IrcServer, Ok(()) } +#[derive(Clone, Default, Debug)] +struct ThreadedPlugins { + plugins: HashMap<String, Arc<Plugin>>, +} + +impl ThreadedPlugins { + pub fn new() -> ThreadedPlugins { + ThreadedPlugins { + plugins: HashMap::new(), + } + } + + pub fn add<T: Plugin + 'static>(&mut self, plugin: T) { + let name = plugin.name().to_lowercase(); + let safe_plugin = Arc::new(plugin); -#[cfg(test)] -mod tests {} + self.plugins.insert(name, safe_plugin); + } + + pub fn remove(&mut self, name: &str) -> Option<()> { + self.plugins.remove(&name.to_lowercase()).map(|_| ()) + } + + pub fn execute_plugins(&mut self, client: &IrcClient, message: Message) { + let message = Arc::new(message); + + for (name, plugin) in self.plugins.clone() { + // Send the message to the plugin if the plugin needs it + match plugin.execute(client, &message) { + ExecutionStatus::Done => (), + ExecutionStatus::Err(e) => log_error(e), + ExecutionStatus::RequiresThread => { + debug!( + "Spawning thread to execute {} with {}", + name, + message.to_string().replace("\r\n", "") + ); + + // Clone everything before the move - the client uses an Arc internally too + let plugin = Arc::clone(&plugin); + let message = Arc::clone(&message); + let client = client.clone(); + + // Execute the plugin in another thread + spawn(move || { + if let Err(e) = plugin.execute_threaded(&client, &message) { + log_error(e); + }; + }); + } + } + } + } + + pub fn handle_command( + &mut self, + client: &IrcClient, + mut command: PluginCommand, + ) -> Result<(), FrippyError> { + if !command.tokens.iter().any(|s| !s.is_empty()) { + let help = format!("Use \"{} help\" to get help", client.current_nickname()); + client + .send_notice(&command.source, &help) + .context(ErrorKind::Connection)?; + } + + // Check if the command is for this plugin + if let Some(plugin) = self.plugins.get(&command.tokens[0].to_lowercase()) { + // The first token contains the name of the plugin + let name = command.tokens.remove(0); + + debug!("Sending command \"{:?}\" to {}", command, name); + + // Clone for the move - the client uses an Arc internally + let client = client.clone(); + let plugin = Arc::clone(plugin); + spawn(move || { + if let Err(e) = plugin.command(&client, command) { + log_error(e); + }; + }); + + Ok(()) + } else { + let help = format!( + "\"{} {}\" is not a command, \ + try \"{0} help\" instead.", + client.current_nickname(), + command.tokens[0] + ); + + Ok(client + .send_notice(&command.source, &help) + .context(ErrorKind::Connection)?) + } + } +} + +impl fmt::Display for ThreadedPlugins { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let plugin_names = self.plugins + .iter() + .map(|(_, p)| p.name().to_owned()) + .collect::<Vec<String>>(); + write!(f, "{}", plugin_names.join(", ")) + } +} diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..b9a4b8f --- /dev/null +++ b/src/main.rs @@ -0,0 +1,191 @@ +#![cfg_attr(feature = "clippy", feature(plugin))] +#![cfg_attr(feature = "clippy", plugin(clippy))] + +extern crate frippy; +extern crate glob; +extern crate irc; +extern crate time; + +#[cfg(feature = "mysql")] +extern crate diesel; +#[cfg(feature = "mysql")] +#[macro_use] +extern crate diesel_migrations; +#[cfg(feature = "mysql")] +extern crate r2d2; +#[cfg(feature = "mysql")] +extern crate r2d2_diesel; + +#[macro_use] +extern crate failure; +#[macro_use] +extern crate log; + +#[cfg(feature = "mysql")] +use std::sync::Arc; +use std::collections::HashMap; +use log::{Level, LevelFilter, Metadata, Record}; + +use irc::client::reactor::IrcReactor; +use glob::glob; + +pub use frippy::plugins::help::Help; +pub use frippy::plugins::url::Url; +pub use frippy::plugins::emoji::Emoji; +pub use frippy::plugins::tell::Tell; +pub use frippy::plugins::currency::Currency; +pub use frippy::plugins::keepnick::KeepNick; +pub use frippy::plugins::factoids::Factoids; + +use frippy::Config; +use failure::Error; + +#[cfg(feature = "mysql")] +embed_migrations!(); + +struct Logger; + +impl log::Log for Logger { + fn enabled(&self, metadata: &Metadata) -> bool { + metadata.target().contains("frippy") + } + + fn log(&self, record: &Record) { + if self.enabled(record.metadata()) { + if record.metadata().level() >= Level::Debug { + println!( + "[{}]({}) {} -> {}", + time::now().rfc822(), + record.level(), + record.target(), + record.args() + ); + } else { + println!( + "[{}]({}) {}", + time::now().rfc822(), + record.level(), + record.args() + ); + } + } + } + + fn flush(&self) {} +} + +static LOGGER: Logger = Logger; + +fn main() { + // Print any errors that caused frippy to shut down + if let Err(e) = run() { + let text = e.causes() + .skip(1) + .fold(format!("{}", e), |acc, err| format!("{}: {}", acc, err)); + error!("{}", text); + }; +} + +fn run() -> Result<(), Error> { + log::set_max_level(if cfg!(debug_assertions) { + LevelFilter::Debug + } else { + LevelFilter::Info + }); + + log::set_logger(&LOGGER).unwrap(); + + // Load all toml files in the configs directory + let mut configs = Vec::new(); + for toml in glob("configs/*.toml").unwrap() { + match toml { + Ok(path) => { + info!("Loading {}", path.to_str().unwrap()); + match Config::load(path) { + Ok(v) => configs.push(v), + Err(e) => error!("Incorrect config file {}", e), + } + } + Err(e) => error!("Failed to read path {}", e), + } + } + + // Without configs the bot would just idle + if configs.is_empty() { + bail!("No config file was found"); + } + + // Create an event loop to run the connections on. + let mut reactor = IrcReactor::new()?; + + // Open a connection and add work for each config + for config in configs { + let mut disabled_plugins = None; + let mut mysql_url = None; + if let Some(ref options) = config.options { + if let Some(disabled) = options.get("disabled_plugins") { + disabled_plugins = Some(disabled.split(',').map(|p| p.trim()).collect::<Vec<_>>()); + } + + mysql_url = options.get("mysql_url"); + } + + let mut bot = frippy::Bot::new(); + bot.add_plugin(Help::new()); + bot.add_plugin(Url::new(1024)); + bot.add_plugin(Emoji::new()); + bot.add_plugin(Currency::new()); + bot.add_plugin(KeepNick::new()); + + #[cfg(feature = "mysql")] + { + if let Some(url) = mysql_url { + use diesel::MysqlConnection; + use r2d2; + use r2d2_diesel::ConnectionManager; + + let manager = ConnectionManager::<MysqlConnection>::new(url.clone()); + match r2d2::Pool::builder().build(manager) { + Ok(pool) => match embedded_migrations::run(&*pool.get()?) { + Ok(_) => { + let pool = Arc::new(pool); + bot.add_plugin(Factoids::new(pool.clone())); + bot.add_plugin(Tell::new(pool.clone())); + info!("Connected to MySQL server") + } + Err(e) => { + bot.add_plugin(Factoids::new(HashMap::new())); + bot.add_plugin(Tell::new(HashMap::new())); + error!("Failed to run migrations: {}", e); + } + }, + Err(e) => error!("Failed to connect to database: {}", e), + } + } else { + bot.add_plugin(Factoids::new(HashMap::new())); + bot.add_plugin(Tell::new(HashMap::new())); + } + } + #[cfg(not(feature = "mysql"))] + { + if mysql_url.is_some() { + error!("frippy was not built with the mysql feature") + } + bot.add_plugin(Factoids::new(HashMap::new())); + bot.add_plugin(Tell::new(HashMap::new())); + } + + if let Some(disabled_plugins) = disabled_plugins { + for name in disabled_plugins { + if bot.remove_plugin(name).is_none() { + error!("\"{}\" was not found - could not disable", name); + } + } + } + + bot.connect(&mut reactor, &config)?; + } + + // Run the bots until they throw an error - an error could be loss of connection + Ok(reactor.run()?) +} diff --git a/src/plugin.rs b/src/plugin.rs index d1f849a..bc428d5 100644 --- a/src/plugin.rs +++ b/src/plugin.rs @@ -1,49 +1,85 @@ +//! Definitions required for every `Plugin` use std::fmt; -use std::collections::HashMap; -use std::thread::spawn; -use std::sync::Arc; use irc::client::prelude::*; -use irc::error::Error as IrcError; +use error::FrippyError; + +/// Describes if a [`Plugin`](trait.Plugin.html) is done working on a +/// [`Message`](../../irc/proto/message/struct.Message.html) or if another thread is required. +#[derive(Debug)] +pub enum ExecutionStatus { + /// The [`Plugin`](trait.Plugin.html) does not need to do any more work on this + /// [`Message`](../../irc/proto/message/struct.Message.html). + Done, + /// An error occured during the execution. + Err(FrippyError), + /// The execution needs to be done by [`execute_threaded()`](trait.Plugin.html#tymethod.execute_threaded). + RequiresThread, +} +/// `Plugin` has to be implemented for any struct that should be usable +/// as a `Plugin` in frippy. pub trait Plugin: PluginName + Send + Sync + fmt::Debug { - fn is_allowed(&self, server: &IrcServer, message: &Message) -> bool; - fn execute(&self, server: &IrcServer, message: &Message) -> Result<(), IrcError>; - fn command(&self, server: &IrcServer, command: PluginCommand) -> Result<(), IrcError>; + /// Handles messages which are not commands or returns + /// [`RequiresThread`](enum.ExecutionStatus.html#variant.RequiresThread) + /// if [`execute_threaded()`](trait.Plugin.html#tymethod.execute_threaded) should be used instead. + fn execute(&self, client: &IrcClient, message: &Message) -> ExecutionStatus; + /// Handles messages which are not commands in a new thread. + fn execute_threaded(&self, client: &IrcClient, message: &Message) -> Result<(), FrippyError>; + /// Handles any command directed at this plugin. + fn command(&self, client: &IrcClient, command: PluginCommand) -> Result<(), FrippyError>; + /// Similar to [`command()`](trait.Plugin.html#tymethod.command) but return a String instead of + /// sending messages directly to IRC. + fn evaluate(&self, client: &IrcClient, command: PluginCommand) -> Result<String, String>; } -pub trait PluginName: Send + Sync + fmt::Debug { +/// `PluginName` is required by [`Plugin`](trait.Plugin.html). +/// +/// To implement it simply add `#[derive(PluginName)]` +/// above the definition of the struct. +/// +/// # Examples +/// ```ignore +/// #[macro_use] extern crate frippy_derive; +/// +/// #[derive(PluginName)] +/// struct Foo; +/// ``` +pub trait PluginName { + /// Returns the name of the `Plugin`. fn name(&self) -> &str; } +/// Represents a command sent by a user to the bot. #[derive(Clone, Debug)] pub struct PluginCommand { + /// The sender of the command. pub source: String, + /// If the command was sent to a channel, this will be that channel + /// otherwise it is the same as `source`. pub target: String, + /// The remaining part of the message that has not been processed yet - split by spaces. pub tokens: Vec<String>, } impl PluginCommand { + /// Creates a `PluginCommand` from [`Message`](../../irc/proto/message/struct.Message.html) + /// if it contains a [`PRIVMSG`](../../irc/proto/command/enum.Command.html#variant.PRIVMSG) + /// that starts with the provided `nick`. pub fn from(nick: &str, message: &Message) -> Option<PluginCommand> { - // Get the actual message out of PRIVMSG if let Command::PRIVMSG(_, ref content) = message.command { - // Split content by spaces and filter empty tokens let mut tokens: Vec<String> = content.split(' ').map(ToOwned::to_owned).collect(); // Commands start with our name if tokens[0].to_lowercase().starts_with(nick) { - // Remove the bot's name from the first token tokens[0].drain(..nick.len()); // We assume that only ':' and ',' are used as suffixes on IRC // If there are any other chars we assume that it is not ment for the bot - tokens[0] = tokens[0] - .chars() - .filter(|&c| !":,".contains(c)) - .collect(); + tokens[0] = tokens[0].chars().filter(|&c| !":,".contains(c)).collect(); if !tokens[0].is_empty() { return None; } @@ -52,10 +88,10 @@ impl PluginCommand { tokens.remove(0); Some(PluginCommand { - source: message.source_nickname().unwrap().to_string(), - target: message.response_target().unwrap().to_string(), - tokens: tokens, - }) + source: message.source_nickname().unwrap().to_string(), + target: message.response_target().unwrap().to_string(), + tokens: tokens, + }) } else { None } @@ -64,96 +100,3 @@ impl PluginCommand { } } } - -#[derive(Clone, Debug)] -pub struct ThreadedPlugins { - plugins: HashMap<String, Arc<Plugin>>, -} - -impl ThreadedPlugins { - pub fn new() -> ThreadedPlugins { - ThreadedPlugins { plugins: HashMap::new() } - } - - pub fn add<T: Plugin + 'static>(&mut self, plugin: T) { - let name = plugin.name().to_lowercase(); - let safe_plugin = Arc::new(plugin); - - self.plugins.insert(name, safe_plugin); - } - - pub fn execute_plugins(&mut self, server: &IrcServer, message: Arc<Message>) { - - for (name, plugin) in self.plugins.clone() { - // Send the message to the plugin if the plugin needs it - if plugin.is_allowed(server, &message) { - - debug!("Executing {} with {}", - name, - message.to_string().replace("\r\n", "")); - - // Clone everything before the move - // The server uses an Arc internally too - let plugin = Arc::clone(&plugin); - let message = Arc::clone(&message); - let server = server.clone(); - - // Execute the plugin in another thread - spawn(move || { - if let Err(e) = plugin.execute(&server, &message) { - error!("Error in {} - {}", name, e); - }; - }); - } - } - } - - pub fn handle_command(&mut self, - server: &IrcServer, - mut command: PluginCommand) - -> Result<(), IrcError> { - - if !command.tokens.iter().any(|s| !s.is_empty()) { - let help = format!("Use \"{} help\" to get help", server.current_nickname()); - return server.send_notice(&command.source, &help); - } - - // Check if the command is for this plugin - if let Some(plugin) = self.plugins.get(&command.tokens[0].to_lowercase()) { - - // The first token contains the name of the plugin - let name = command.tokens.remove(0); - - debug!("Sending command \"{:?}\" to {}", command, name); - - // Clone for the move - the server uses an Arc internally - let server = server.clone(); - let plugin = Arc::clone(plugin); - spawn(move || { - if let Err(e) = plugin.command(&server, command) { - error!("Error in {} command - {}", name, e); - }; - }); - - Ok(()) - - } else { - let help = format!("\"{} {}\" is not a command, \ - try \"{0} help\" instead.", - server.current_nickname(), - command.tokens[0]); - - server.send_notice(&command.source, &help) - } - } -} - -impl fmt::Display for ThreadedPlugins { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - let plugin_names = self.plugins - .iter() - .map(|(_, p)| p.name().to_string()) - .collect::<Vec<String>>(); - write!(f, "{}", plugin_names.join(", ")) - } -} diff --git a/src/plugins/currency.rs b/src/plugins/currency.rs index d6cf928..53a245c 100644 --- a/src/plugins/currency.rs +++ b/src/plugins/currency.rs @@ -6,7 +6,6 @@ use std::io::Read; use std::num::ParseFloatError; use irc::client::prelude::*; -use irc::error::Error as IrcError; use self::reqwest::Client; use self::reqwest::header::Connection; @@ -14,7 +13,11 @@ use self::serde_json::Value; use plugin::*; -#[derive(PluginName, Debug)] +use error::FrippyError; +use error::ErrorKind as FrippyErrorKind; +use failure::ResultExt; + +#[derive(PluginName, Default, Debug)] pub struct Currency; struct ConvertionRequest<'a> { @@ -23,18 +26,8 @@ struct ConvertionRequest<'a> { target: &'a str, } -macro_rules! try_option { - ($e:expr) => { - match $e { - Some(v) => v, - None => { return None; } - } - } -} - impl<'a> ConvertionRequest<'a> { fn send(&self) -> Option<f64> { - let response = Client::new() .get("https://api.fixer.io/latest") .form(&[("base", self.source)]) @@ -44,16 +37,14 @@ impl<'a> ConvertionRequest<'a> { match response { Ok(mut response) => { let mut body = String::new(); - try_option!(response.read_to_string(&mut body).ok()); + response.read_to_string(&mut body).ok()?; let convertion_rates: Result<Value, _> = serde_json::from_str(&body); match convertion_rates { Ok(convertion_rates) => { - - let rates: &Value = try_option!(convertion_rates.get("rates")); - let target_rate: &Value = - try_option!(rates.get(self.target.to_uppercase())); - Some(self.value * try_option!(target_rate.as_f64())) + let rates: &Value = convertion_rates.get("rates")?; + let target_rate: &Value = rates.get(self.target.to_uppercase())?; + Some(self.value * target_rate.as_f64()?) } Err(_) => None, } @@ -68,7 +59,10 @@ impl Currency { Currency {} } - fn eval_command<'a>(&self, tokens: &'a [String]) -> Result<ConvertionRequest<'a>, ParseFloatError> { + fn eval_command<'a>( + &self, + tokens: &'a [String], + ) -> Result<ConvertionRequest<'a>, ParseFloatError> { Ok(ConvertionRequest { value: tokens[0].parse()?, source: &tokens[1], @@ -76,76 +70,97 @@ impl Currency { }) } - fn convert(&self, server: &IrcServer, command: PluginCommand) -> Result<(), IrcError> { - + fn convert(&self, client: &IrcClient, command: &mut PluginCommand) -> Result<String, String> { if command.tokens.len() < 3 { - return self.invalid_command(server, &command); + return Err(self.invalid_command(client)); } let request = match self.eval_command(&command.tokens) { Ok(request) => request, Err(_) => { - return self.invalid_command(server, &command); + return Err(self.invalid_command(client)); } }; match request.send() { Some(response) => { - let response = format!("{} {} => {:.4} {}", - request.value, - request.source.to_lowercase(), - response / 1.00000000, - request.target.to_lowercase()); - - server.send_privmsg(&command.target, &response) + let response = format!( + "{} {} => {:.4} {}", + request.value, + request.source.to_lowercase(), + response / 1.00000000, + request.target.to_lowercase() + ); + + Ok(response) } - None => server.send_notice(&command.source, "Error while converting given currency"), + None => Err(String::from( + "An error occured during the conversion of the given currency", + )), } } - fn help(&self, server: &IrcServer, command: &mut PluginCommand) -> Result<(), IrcError> { - let help = format!("usage: {} currency value from_currency to_currency\r\n\ - example: 1.5 eur usd\r\n\ - available currencies: AUD, BGN, BRL, CAD, \ - CHF, CNY, CZK, DKK, GBP, HKD, HRK, HUF, \ - IDR, ILS, INR, JPY, KRW, MXN, MYR, NOK, \ - NZD, PHP, PLN, RON, RUB, SEK, SGD, THB, \ - TRY, USD, ZAR", - server.current_nickname()); - - server.send_notice(&command.source, &help) + fn help(&self, client: &IrcClient) -> String { + format!( + "usage: {} currency value from_currency to_currency\r\n\ + example: {0} currency 1.5 eur usd\r\n\ + available currencies: AUD, BGN, BRL, CAD, \ + CHF, CNY, CZK, DKK, GBP, HKD, HRK, HUF, \ + IDR, ILS, INR, JPY, KRW, MXN, MYR, NOK, \ + NZD, PHP, PLN, RON, RUB, SEK, SGD, THB, \ + TRY, USD, ZAR", + client.current_nickname() + ) } - fn invalid_command(&self, server: &IrcServer, command: &PluginCommand) -> Result<(), IrcError> { - let help = format!("Incorrect Command. \ - Send \"{} currency help\" for help.", - server.current_nickname()); - - server.send_notice(&command.source, &help) + fn invalid_command(&self, client: &IrcClient) -> String { + format!( + "Incorrect Command. \ + Send \"{} currency help\" for help.", + client.current_nickname() + ) } } impl Plugin for Currency { - fn is_allowed(&self, _: &IrcServer, _: &Message) -> bool { - false + fn execute(&self, _: &IrcClient, _: &Message) -> ExecutionStatus { + ExecutionStatus::Done } - fn execute(&self, _: &IrcServer, _: &Message) -> Result<(), IrcError> { + fn execute_threaded(&self, _: &IrcClient, _: &Message) -> Result<(), FrippyError> { panic!("Currency does not implement the execute function!") } - fn command(&self, server: &IrcServer, mut command: PluginCommand) -> Result<(), IrcError> { + fn command(&self, client: &IrcClient, mut command: PluginCommand) -> Result<(), FrippyError> { + if command.tokens.is_empty() { + return Ok(client + .send_notice(&command.source, &self.invalid_command(client)) + .context(FrippyErrorKind::Connection)?); + } + + match command.tokens[0].as_ref() { + "help" => Ok(client + .send_notice(&command.source, &self.help(client)) + .context(FrippyErrorKind::Connection)?), + _ => match self.convert(client, &mut command) { + Ok(msg) => Ok(client + .send_privmsg(&command.target, &msg) + .context(FrippyErrorKind::Connection)?), + Err(msg) => Ok(client + .send_notice(&command.source, &msg) + .context(FrippyErrorKind::Connection)?), + }, + } + } + fn evaluate(&self, client: &IrcClient, mut command: PluginCommand) -> Result<String, String> { if command.tokens.is_empty() { - return self.invalid_command(server, &command); + return Err(self.invalid_command(client)); } match command.tokens[0].as_ref() { - "help" => self.help(server, &mut command), - _ => self.convert(server, command), + "help" => Ok(self.help(client)), + _ => self.convert(client, &mut command), } } } - -#[cfg(test)] -mod tests {} diff --git a/src/plugins/emoji.rs b/src/plugins/emoji.rs index 1bb714c..f1d9376 100644 --- a/src/plugins/emoji.rs +++ b/src/plugins/emoji.rs @@ -3,10 +3,14 @@ extern crate unicode_names; use std::fmt; use irc::client::prelude::*; -use irc::error::Error as IrcError; use plugin::*; +use error::FrippyError; +use error::ErrorKind as FrippyErrorKind; +use failure::Fail; +use failure::ResultExt; + struct EmojiHandle { symbol: char, count: i32, @@ -14,7 +18,6 @@ struct EmojiHandle { impl fmt::Display for EmojiHandle { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - let name = match unicode_names::name(self.symbol) { Some(sym) => sym.to_string().to_lowercase(), None => String::from("UNKNOWN"), @@ -28,7 +31,7 @@ impl fmt::Display for EmojiHandle { } } -#[derive(PluginName, Debug)] +#[derive(PluginName, Default, Debug)] pub struct Emoji; impl Emoji { @@ -36,13 +39,12 @@ impl Emoji { Emoji {} } - fn emoji(&self, server: &IrcServer, content: &str, target: &str) -> Result<(), IrcError> { - let names = self.return_emojis(content) + fn emoji(&self, content: &str) -> String { + self.return_emojis(content) .iter() .map(|e| e.to_string()) - .collect::<Vec<String>>(); - - server.send_privmsg(target, &names.join(", ")) + .collect::<Vec<String>>() + .join(", ") } fn return_emojis(&self, string: &str) -> Vec<EmojiHandle> { @@ -53,7 +55,6 @@ impl Emoji { count: 0, }; - for c in string.chars() { if !self.is_emoji(&c) { continue; @@ -61,7 +62,6 @@ impl Emoji { if current.symbol == c { current.count += 1; - } else { if current.count > 0 { emojis.push(current); @@ -98,27 +98,37 @@ impl Emoji { } impl Plugin for Emoji { - fn is_allowed(&self, _: &IrcServer, message: &Message) -> bool { + fn execute(&self, client: &IrcClient, message: &Message) -> ExecutionStatus { match message.command { - Command::PRIVMSG(_, _) => true, - _ => false, + Command::PRIVMSG(_, ref content) => match client + .send_privmsg(message.response_target().unwrap(), &self.emoji(content)) + { + Ok(_) => ExecutionStatus::Done, + Err(e) => ExecutionStatus::Err(e.context(FrippyErrorKind::Connection).into()), + }, + _ => ExecutionStatus::Done, } } - fn execute(&self, server: &IrcServer, message: &Message) -> Result<(), IrcError> { - match message.command { - Command::PRIVMSG(_, ref content) => { - self.emoji(server, content, message.response_target().unwrap()) - } - _ => Ok(()), - } + fn execute_threaded(&self, _: &IrcClient, _: &Message) -> Result<(), FrippyError> { + panic!("Emoji should not use threading") } - fn command(&self, server: &IrcServer, command: PluginCommand) -> Result<(), IrcError> { - server.send_notice(&command.source, - "This Plugin does not implement any commands.") + fn command(&self, client: &IrcClient, command: PluginCommand) -> Result<(), FrippyError> { + Ok(client + .send_notice( + &command.source, + "This Plugin does not implement any commands.", + ) + .context(FrippyErrorKind::Connection)?) } -} -#[cfg(test)] -mod tests {} + fn evaluate(&self, _: &IrcClient, command: PluginCommand) -> Result<String, String> { + let emojis = self.emoji(&command.tokens[0]); + if emojis.is_empty() { + Ok(emojis) + } else { + Err(String::from("No emojis were found.")) + } + } +} diff --git a/src/plugins/factoids/database.rs b/src/plugins/factoids/database.rs new file mode 100644 index 0000000..b1fe8dd --- /dev/null +++ b/src/plugins/factoids/database.rs @@ -0,0 +1,163 @@ +#[cfg(feature = "mysql")] +extern crate dotenv; + +#[cfg(feature = "mysql")] +use std::sync::Arc; +use std::collections::HashMap; + +#[cfg(feature = "mysql")] +use diesel::prelude::*; +#[cfg(feature = "mysql")] +use diesel::mysql::MysqlConnection; +#[cfg(feature = "mysql")] +use r2d2::Pool; +#[cfg(feature = "mysql")] +use r2d2_diesel::ConnectionManager; +#[cfg(feature = "mysql")] +use failure::ResultExt; + +use chrono::NaiveDateTime; + +use super::error::*; + +#[cfg_attr(feature = "mysql", derive(Queryable))] +#[derive(Clone, Debug)] +pub struct Factoid { + pub name: String, + pub idx: i32, + pub content: String, + pub author: String, + pub created: NaiveDateTime, +} + +#[cfg_attr(feature = "mysql", derive(Insertable))] +#[cfg_attr(feature = "mysql", table_name = "factoids")] +pub struct NewFactoid<'a> { + pub name: &'a str, + pub idx: i32, + pub content: &'a str, + pub author: &'a str, + pub created: NaiveDateTime, +} + +pub trait Database: Send { + fn insert_factoid(&mut self, factoid: &NewFactoid) -> Result<(), FactoidsError>; + fn get_factoid(&self, name: &str, idx: i32) -> Result<Factoid, FactoidsError>; + fn delete_factoid(&mut self, name: &str, idx: i32) -> Result<(), FactoidsError>; + fn count_factoids(&self, name: &str) -> Result<i32, FactoidsError>; +} + +// HashMap +impl Database for HashMap<(String, i32), Factoid> { + fn insert_factoid(&mut self, factoid: &NewFactoid) -> Result<(), FactoidsError> { + let factoid = Factoid { + name: String::from(factoid.name), + idx: factoid.idx, + content: factoid.content.to_string(), + author: factoid.author.to_string(), + created: factoid.created, + }; + + let name = factoid.name.clone(); + match self.insert((name, factoid.idx), factoid) { + None => Ok(()), + Some(_) => Err(ErrorKind::Duplicate)?, + } + } + + fn get_factoid(&self, name: &str, idx: i32) -> Result<Factoid, FactoidsError> { + Ok(self.get(&(String::from(name), idx)) + .cloned() + .ok_or(ErrorKind::NotFound)?) + } + + fn delete_factoid(&mut self, name: &str, idx: i32) -> Result<(), FactoidsError> { + match self.remove(&(String::from(name), idx)) { + Some(_) => Ok(()), + None => Err(ErrorKind::NotFound)?, + } + } + + fn count_factoids(&self, name: &str) -> Result<i32, FactoidsError> { + Ok(self.iter().filter(|&(&(ref n, _), _)| n == name).count() as i32) + } +} + +// Diesel automatically defines the factoids module as public. +// We create a schema module to keep it private. +#[cfg(feature = "mysql")] +mod schema { + table! { + factoids (name, idx) { + name -> Varchar, + idx -> Integer, + content -> Text, + author -> Varchar, + created -> Timestamp, + } + } +} + +#[cfg(feature = "mysql")] +use self::schema::factoids; + +#[cfg(feature = "mysql")] +impl Database for Arc<Pool<ConnectionManager<MysqlConnection>>> { + fn insert_factoid(&mut self, factoid: &NewFactoid) -> Result<(), FactoidsError> { + use diesel; + + let conn = &*self.get().context(ErrorKind::NoConnection)?; + diesel::insert_into(factoids::table) + .values(factoid) + .execute(conn) + .context(ErrorKind::MysqlError)?; + + Ok(()) + } + + fn get_factoid(&self, name: &str, idx: i32) -> Result<Factoid, FactoidsError> { + let conn = &*self.get().context(ErrorKind::NoConnection)?; + Ok(factoids::table + .find((name, idx)) + .first(conn) + .context(ErrorKind::MysqlError)?) + } + + fn delete_factoid(&mut self, name: &str, idx: i32) -> Result<(), FactoidsError> { + use diesel; + use self::factoids::columns; + + let conn = &*self.get().context(ErrorKind::NoConnection)?; + match diesel::delete( + factoids::table + .filter(columns::name.eq(name)) + .filter(columns::idx.eq(idx)), + ).execute(conn) + { + Ok(v) => { + if v > 0 { + Ok(()) + } else { + Err(ErrorKind::NotFound)? + } + } + Err(e) => Err(e).context(ErrorKind::MysqlError)?, + } + } + + fn count_factoids(&self, name: &str) -> Result<i32, FactoidsError> { + use diesel; + + let conn = &*self.get().context(ErrorKind::NoConnection)?; + let count: Result<i64, _> = factoids::table + .filter(factoids::columns::name.eq(name)) + .count() + .get_result(conn); + + match count { + Ok(c) => Ok(c as i32), + Err(diesel::NotFound) => Ok(0), + Err(e) => Err(e).context(ErrorKind::MysqlError)?, + } + } +} diff --git a/src/plugins/factoids/mod.rs b/src/plugins/factoids/mod.rs new file mode 100644 index 0000000..2f3690f --- /dev/null +++ b/src/plugins/factoids/mod.rs @@ -0,0 +1,342 @@ +extern crate rlua; + +use std::fmt; +use std::str::FromStr; +use std::sync::Mutex; +use self::rlua::prelude::*; +use irc::client::prelude::*; + +use time; +use chrono::NaiveDateTime; + +use plugin::*; +pub mod database; +use self::database::Database; + +mod utils; +use self::utils::*; + +use failure::ResultExt; +use error::ErrorKind as FrippyErrorKind; +use error::FrippyError; +use self::error::*; + +static LUA_SANDBOX: &'static str = include_str!("sandbox.lua"); + +#[derive(PluginName)] +pub struct Factoids<T: Database> { + factoids: Mutex<T>, +} + +macro_rules! try_lock { + ( $m:expr ) => { + match $m.lock() { + Ok(guard) => guard, + Err(poisoned) => poisoned.into_inner(), + } + } +} + +impl<T: Database> Factoids<T> { + pub fn new(db: T) -> Factoids<T> { + Factoids { + factoids: Mutex::new(db), + } + } + + fn create_factoid( + &self, + name: &str, + content: &str, + author: &str, + ) -> Result<&str, FactoidsError> { + let count = try_lock!(self.factoids).count_factoids(name)?; + let tm = time::now().to_timespec(); + + let factoid = database::NewFactoid { + name: name, + idx: count, + content: content, + author: author, + created: NaiveDateTime::from_timestamp(tm.sec, 0u32), + }; + + Ok(try_lock!(self.factoids) + .insert_factoid(&factoid) + .map(|()| "Successfully added!")?) + } + + fn add(&self, command: &mut PluginCommand) -> Result<&str, FactoidsError> { + if command.tokens.len() < 2 { + Err(ErrorKind::InvalidCommand)?; + } + + let name = command.tokens.remove(0); + let content = command.tokens.join(" "); + + Ok(self.create_factoid(&name, &content, &command.source)?) + } + + fn add_from_url(&self, command: &mut PluginCommand) -> Result<&str, FactoidsError> { + if command.tokens.len() < 2 { + Err(ErrorKind::InvalidCommand)?; + } + + let name = command.tokens.remove(0); + let url = &command.tokens[0]; + let content = ::utils::download(url, Some(1024)).context(ErrorKind::Download)?; + + Ok(self.create_factoid(&name, &content, &command.source)?) + } + + fn remove(&self, command: &mut PluginCommand) -> Result<&str, FactoidsError> { + if command.tokens.len() < 1 { + Err(ErrorKind::InvalidCommand)?; + } + + let name = command.tokens.remove(0); + let count = try_lock!(self.factoids).count_factoids(&name)?; + + match try_lock!(self.factoids).delete_factoid(&name, count - 1) { + Ok(()) => Ok("Successfully removed"), + Err(e) => Err(e)?, + } + } + + fn get(&self, command: &PluginCommand) -> Result<String, FactoidsError> { + let (name, idx) = match command.tokens.len() { + 0 => Err(ErrorKind::InvalidCommand)?, + 1 => { + let name = &command.tokens[0]; + let count = try_lock!(self.factoids).count_factoids(name)?; + + if count < 1 { + Err(ErrorKind::NotFound)?; + } + + (name, count - 1) + } + _ => { + let name = &command.tokens[0]; + let idx = match i32::from_str(&command.tokens[1]) { + Ok(i) => i, + Err(_) => Err(ErrorKind::InvalidCommand)?, + }; + + (name, idx) + } + }; + + let factoid = try_lock!(self.factoids) + .get_factoid(name, idx) + .context(ErrorKind::NotFound)?; + + let message = factoid.content.replace("\n", "|").replace("\r", ""); + + Ok(format!("{}: {}", factoid.name, message)) + } + + fn info(&self, command: &PluginCommand) -> Result<String, FactoidsError> { + match command.tokens.len() { + 0 => Err(ErrorKind::InvalidCommand)?, + 1 => { + let name = &command.tokens[0]; + let count = try_lock!(self.factoids).count_factoids(name)?; + + Ok(match count { + 0 => Err(ErrorKind::NotFound)?, + 1 => format!("There is 1 version of {}", name), + _ => format!("There are {} versions of {}", count, name), + }) + } + _ => { + let name = &command.tokens[0]; + let idx = i32::from_str(&command.tokens[1]).context(ErrorKind::InvalidIndex)?; + let factoid = try_lock!(self.factoids).get_factoid(name, idx)?; + + Ok(format!( + "{}: Added by {} at {} UTC", + name, factoid.author, factoid.created + )) + } + } + } + + fn exec(&self, mut command: PluginCommand) -> Result<String, FactoidsError> { + if command.tokens.len() < 1 { + Err(ErrorKind::InvalidIndex)? + } else { + let name = command.tokens.remove(0); + let count = try_lock!(self.factoids).count_factoids(&name)?; + let factoid = try_lock!(self.factoids).get_factoid(&name, count - 1)?; + + let content = factoid.content; + let value = if content.starts_with('>') { + let content = String::from(&content[1..]); + + if content.starts_with('>') { + content + } else { + match self.run_lua(&name, &content, &command) { + Ok(v) => v, + Err(e) => format!("\"{}\"", e), + } + } + } else { + content + }; + + Ok(value.replace("\n", "|").replace("\r", "")) + } + } + + fn run_lua( + &self, + name: &str, + code: &str, + command: &PluginCommand, + ) -> Result<String, rlua::Error> { + let args = command + .tokens + .iter() + .filter(|x| !x.is_empty()) + .map(ToOwned::to_owned) + .collect::<Vec<String>>(); + + let lua = unsafe { Lua::new_with_debug() }; + let globals = lua.globals(); + + globals.set("factoid", code)?; + globals.set("download", lua.create_function(download)?)?; + globals.set("sleep", lua.create_function(sleep)?)?; + globals.set("args", args)?; + globals.set("input", command.tokens.join(" "))?; + globals.set("user", command.source.clone())?; + globals.set("channel", command.target.clone())?; + globals.set("output", lua.create_table()?)?; + + lua.exec::<()>(LUA_SANDBOX, Some(name))?; + let output: Vec<String> = globals.get::<_, Vec<String>>("output")?; + + Ok(output.join("|")) + } +} + +impl<T: Database> Plugin for Factoids<T> { + fn execute(&self, _: &IrcClient, message: &Message) -> ExecutionStatus { + match message.command { + Command::PRIVMSG(_, ref content) => if content.starts_with('!') { + ExecutionStatus::RequiresThread + } else { + ExecutionStatus::Done + }, + _ => ExecutionStatus::Done, + } + } + + fn execute_threaded(&self, client: &IrcClient, message: &Message) -> Result<(), FrippyError> { + if let Command::PRIVMSG(_, mut content) = message.command.clone() { + content.remove(0); + + let t: Vec<String> = content.split(' ').map(ToOwned::to_owned).collect(); + + let c = PluginCommand { + source: message.source_nickname().unwrap().to_owned(), + target: message.response_target().unwrap().to_owned(), + tokens: t, + }; + + Ok(match self.exec(c) { + Ok(f) => client + .send_privmsg(&message.response_target().unwrap(), &f) + .context(FrippyErrorKind::Connection)?, + Err(_) => (), + }) + } else { + Ok(()) + } + } + + fn command(&self, client: &IrcClient, mut command: PluginCommand) -> Result<(), FrippyError> { + if command.tokens.is_empty() { + return Ok(client + .send_notice(&command.target, "Invalid command") + .context(FrippyErrorKind::Connection)?); + } + + let target = command.target.clone(); + let source = command.source.clone(); + + let sub_command = command.tokens.remove(0); + let result = match sub_command.as_ref() { + "add" => self.add(&mut command).map(|s| s.to_owned()), + "fromurl" => self.add_from_url(&mut command).map(|s| s.to_owned()), + "remove" => self.remove(&mut command).map(|s| s.to_owned()), + "get" => self.get(&command), + "info" => self.info(&command), + "exec" => self.exec(command), + _ => Err(ErrorKind::InvalidCommand.into()), + }; + + Ok(match result { + Ok(v) => client + .send_privmsg(&target, &v) + .context(FrippyErrorKind::Connection)?, + Err(e) => { + let message = e.to_string(); + client + .send_notice(&source, &message) + .context(FrippyErrorKind::Connection)?; + Err(e).context(FrippyErrorKind::Factoids)? + } + }) + } + + fn evaluate(&self, _: &IrcClient, _: PluginCommand) -> Result<String, String> { + Err(String::from( + "Evaluation of commands is not implemented for Factoids at this time", + )) + } +} + +impl<T: Database> fmt::Debug for Factoids<T> { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "Factoids {{ ... }}") + } +} + +pub mod error { + #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)] + #[error = "FactoidsError"] + pub enum ErrorKind { + /// Invalid command error + #[fail(display = "Invalid Command")] + InvalidCommand, + + /// Invalid index error + #[fail(display = "Invalid index")] + InvalidIndex, + + /// Download error + #[fail(display = "Download failed")] + Download, + + /// Duplicate error + #[fail(display = "Entry already exists")] + Duplicate, + + /// Not found error + #[fail(display = "Factoid was not found")] + NotFound, + + /// MySQL error + #[cfg(feature = "mysql")] + #[fail(display = "Failed to execute MySQL Query")] + MysqlError, + + /// No connection error + #[cfg(feature = "mysql")] + #[fail(display = "No connection to the database")] + NoConnection, + } +} diff --git a/src/plugins/factoids/sandbox.lua b/src/plugins/factoids/sandbox.lua new file mode 100644 index 0000000..3fc74cd --- /dev/null +++ b/src/plugins/factoids/sandbox.lua @@ -0,0 +1,86 @@ +function send(text) + local text = tostring(text) + local len = #output + if len < 1 then + output = { text } + else + output[len] = output[len] .. text + end +end + +function sendln(text) + send(text) + table.insert(output, "") +end + +local sandbox_env = { + print = send, + println = sendln, + eval = nil, + args = args, + input = input, + user = user, + channel = channel, + request = download, + string = string, + math = math, + table = table, + pairs = pairs, + ipairs = ipairs, + next = next, + select = select, + unpack = unpack, + tostring = tostring, + tonumber = tonumber, + type = type, + assert = assert, + error = error, + pcall = pcall, + xpcall = xpcall, + _VERSION = _VERSION +} + +sandbox_env.os = { + clock = os.clock, + time = os.time, + difftime = os.difftime +} + +sandbox_env.string.rep = nil +sandbox_env.string.dump = nil +sandbox_env.math.randomseed = nil + +-- Temporary evaluation function +function eval(code) + local c, e = load(code, nil, nil, sandbox_env) + if c then + return c() + else + error(e) + end +end + +sandbox_env.eval = eval + +-- Check if the factoid timed out +function checktime(event, line) + if os.time() - time >= timeout then + error("Timed out after " .. timeout .. " seconds", 0) + else + -- Limit the cpu usage of factoids + sleep(1) + end +end + +local f, e = load(factoid, nil, nil, sandbox_env) + +-- Add timeout hook +time = os.time() +timeout = 30 +debug.sethook(checktime, "l") + +if f then + f() +else + error(e) +end diff --git a/src/plugins/factoids/utils.rs b/src/plugins/factoids/utils.rs new file mode 100644 index 0000000..70ac8a7 --- /dev/null +++ b/src/plugins/factoids/utils.rs @@ -0,0 +1,25 @@ +extern crate reqwest; + +use std::thread; +use std::time::Duration; + +use utils; +use super::rlua::prelude::*; + +use self::LuaError::RuntimeError; + +pub fn download(_: &Lua, url: String) -> Result<String, LuaError> { + match utils::download(&url, Some(1024)) { + Ok(v) => Ok(v), + Err(e) => Err(RuntimeError(format!( + "Failed to download {} - {}", + url, + e.to_string() + ))), + } +} + +pub fn sleep(_: &Lua, dur: u64) -> Result<(), LuaError> { + thread::sleep(Duration::from_millis(dur)); + Ok(()) +} diff --git a/src/plugins/help.rs b/src/plugins/help.rs index 8f3fb4d..7e3658d 100644 --- a/src/plugins/help.rs +++ b/src/plugins/help.rs @@ -1,34 +1,36 @@ use irc::client::prelude::*; -use irc::error::Error as IrcError; use plugin::*; -#[derive(PluginName, Debug)] +use error::FrippyError; +use error::ErrorKind as FrippyErrorKind; +use failure::ResultExt; + +#[derive(PluginName, Default, Debug)] pub struct Help; impl Help { pub fn new() -> Help { Help {} } - - fn help(&self, server: &IrcServer, command: PluginCommand) -> Result<(), IrcError> { - server.send_notice(&command.source, "Help has not been added yet.") - } } impl Plugin for Help { - fn is_allowed(&self, _: &IrcServer, _: &Message) -> bool { - false + fn execute(&self, _: &IrcClient, _: &Message) -> ExecutionStatus { + ExecutionStatus::Done } - fn execute(&self, _: &IrcServer, _: &Message) -> Result<(), IrcError> { - panic!("Help does not implement the execute function!") + fn execute_threaded(&self, _: &IrcClient, _: &Message) -> Result<(), FrippyError> { + panic!("Help should not use threading") } - fn command(&self, server: &IrcServer, command: PluginCommand) -> Result<(), IrcError> { - self.help(server, command) + fn command(&self, client: &IrcClient, command: PluginCommand) -> Result<(), FrippyError> { + Ok(client + .send_notice(&command.source, "Help has not been added yet.") + .context(FrippyErrorKind::Connection)?) } -} -#[cfg(test)] -mod tests {} + fn evaluate(&self, _: &IrcClient, _: PluginCommand) -> Result<String, String> { + Err(String::from("Help has not been added yet.")) + } +} diff --git a/src/plugins/keepnick.rs b/src/plugins/keepnick.rs new file mode 100644 index 0000000..58ac167 --- /dev/null +++ b/src/plugins/keepnick.rs @@ -0,0 +1,70 @@ +use irc::client::prelude::*; + +use plugin::*; + +use error::FrippyError; +use error::ErrorKind as FrippyErrorKind; +use failure::ResultExt; + +#[derive(PluginName, Default, Debug)] +pub struct KeepNick; + +impl KeepNick { + pub fn new() -> KeepNick { + KeepNick {} + } + + fn check_nick(&self, client: &IrcClient, leaver: &str) -> ExecutionStatus { + let cfg_nick = match client.config().nickname { + Some(ref nick) => nick.clone(), + None => return ExecutionStatus::Done, + }; + + if leaver != cfg_nick { + return ExecutionStatus::Done; + } + + let client_nick = client.current_nickname(); + + if client_nick != cfg_nick { + info!("Trying to switch nick from {} to {}", client_nick, cfg_nick); + match client + .send(Command::NICK(cfg_nick)) + .context(FrippyErrorKind::Connection) + { + Ok(_) => ExecutionStatus::Done, + Err(e) => ExecutionStatus::Err(e.into()), + } + } else { + ExecutionStatus::Done + } + } +} + +impl Plugin for KeepNick { + fn execute(&self, client: &IrcClient, message: &Message) -> ExecutionStatus { + match message.command { + Command::QUIT(ref nick) => { + self.check_nick(client, &nick.clone().unwrap_or_else(String::new)) + } + _ => ExecutionStatus::Done, + } + } + + fn execute_threaded(&self, _: &IrcClient, _: &Message) -> Result<(), FrippyError> { + panic!("Tell should not use threading") + } + + fn command(&self, client: &IrcClient, command: PluginCommand) -> Result<(), FrippyError> { + Ok(client + .send_notice( + &command.source, + "This Plugin does not implement any commands.", + ) + .context(FrippyErrorKind::Connection)?) + } + + fn evaluate(&self, _: &IrcClient, _: PluginCommand) -> Result<String, String> { + Err(String::from("This Plugin does not implement any commands.")) + } +} diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index 0dea596..9a3ba2f 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -1,7 +1,8 @@ -mod help; -mod emoji; -mod currency; - -pub use self::help::Help; -pub use self::emoji::Emoji; -pub use self::currency::Currency; +//! Collection of plugins included +pub mod help; +pub mod url; +pub mod emoji; +pub mod tell; +pub mod currency; +pub mod factoids; +pub mod keepnick; diff --git a/src/plugins/tell/database.rs b/src/plugins/tell/database.rs new file mode 100644 index 0000000..98e9fb3 --- /dev/null +++ b/src/plugins/tell/database.rs @@ -0,0 +1,150 @@ +#[cfg(feature = "mysql")] +extern crate dotenv; + +#[cfg(feature = "mysql")] +use std::sync::Arc; +use std::collections::HashMap; + +#[cfg(feature = "mysql")] +use diesel::prelude::*; +#[cfg(feature = "mysql")] +use diesel::mysql::MysqlConnection; +#[cfg(feature = "mysql")] +use r2d2::Pool; +#[cfg(feature = "mysql")] +use r2d2_diesel::ConnectionManager; + +use chrono::NaiveDateTime; + +#[cfg(feature = "mysql")] +use failure::ResultExt; + +use super::error::*; + +#[cfg_attr(feature = "mysql", derive(Queryable))] +#[derive(PartialEq, Clone, Debug)] +pub struct TellMessage { + pub id: i64, + pub sender: String, + pub receiver: String, + pub time: NaiveDateTime, + pub message: String, +} + +#[cfg_attr(feature = "mysql", derive(Insertable))] +#[cfg_attr(feature = "mysql", table_name = "tells")] +pub struct NewTellMessage<'a> { + pub sender: &'a str, + pub receiver: &'a str, + pub time: NaiveDateTime, + pub message: &'a str, +} + +pub trait Database: Send { + fn insert_tell(&mut self, tell: &NewTellMessage) -> Result<(), TellError>; + fn get_tells(&self, receiver: &str) -> Result<Vec<TellMessage>, TellError>; + fn get_receivers(&self) -> Result<Vec<String>, TellError>; + fn delete_tells(&mut self, receiver: &str) -> Result<(), TellError>; +} + +// HashMap +impl Database for HashMap<String, Vec<TellMessage>> { + fn insert_tell(&mut self, tell: &NewTellMessage) -> Result<(), TellError> { + let tell = TellMessage { + id: 0, + sender: tell.sender.to_string(), + receiver: tell.receiver.to_string(), + time: tell.time, + message: tell.message.to_string(), + }; + + let receiver = tell.receiver.clone(); + let tell_messages = self.entry(receiver) + .or_insert_with(|| Vec::with_capacity(3)); + (*tell_messages).push(tell); + + Ok(()) + } + + fn get_tells(&self, receiver: &str) -> Result<Vec<TellMessage>, TellError> { + Ok(self.get(receiver).cloned().ok_or(ErrorKind::NotFound)?) + } + + fn get_receivers(&self) -> Result<Vec<String>, TellError> { + Ok(self.iter() + .map(|(receiver, _)| receiver.to_owned()) + .collect::<Vec<_>>()) + } + + fn delete_tells(&mut self, receiver: &str) -> Result<(), TellError> { + match self.remove(receiver) { + Some(_) => Ok(()), + None => Err(ErrorKind::NotFound)?, + } + } +} + +// Diesel automatically defines the tells module as public. +// We create a schema module to keep it private. +#[cfg(feature = "mysql")] +mod schema { + table! { + tells (id) { + id -> Bigint, + sender -> Varchar, + receiver -> Varchar, + time -> Timestamp, + message -> Varchar, + } + } +} + +#[cfg(feature = "mysql")] +use self::schema::tells; + +#[cfg(feature = "mysql")] +impl Database for Arc<Pool<ConnectionManager<MysqlConnection>>> { + fn insert_tell(&mut self, tell: &NewTellMessage) -> Result<(), TellError> { + use diesel; + + let conn = &*self.get().expect("Failed to get connection"); + diesel::insert_into(tells::table) + .values(tell) + .execute(conn) + .context(ErrorKind::MysqlError)?; + + Ok(()) + } + + fn get_tells(&self, receiver: &str) -> Result<Vec<TellMessage>, TellError> { + use self::tells::columns; + + let conn = &*self.get().context(ErrorKind::NoConnection)?; + Ok(tells::table + .filter(columns::receiver.eq(receiver)) + .order(columns::time.asc()) + .load::<TellMessage>(conn) + .context(ErrorKind::MysqlError)?) + } + + fn get_receivers(&self) -> Result<Vec<String>, TellError> { + use self::tells::columns; + + let conn = &*self.get().context(ErrorKind::NoConnection)?; + Ok(tells::table + .select(columns::receiver) + .load::<String>(conn) + .context(ErrorKind::MysqlError)?) + } + + fn delete_tells(&mut self, receiver: &str) -> Result<(), TellError> { + use diesel; + use self::tells::columns; + + let conn = &*self.get().context(ErrorKind::NoConnection)?; + diesel::delete(tells::table.filter(columns::receiver.eq(receiver))) + .execute(conn) + .context(ErrorKind::MysqlError)?; + Ok(()) + } +} diff --git a/src/plugins/tell/mod.rs b/src/plugins/tell/mod.rs new file mode 100644 index 0000000..bdfb55c --- /dev/null +++ b/src/plugins/tell/mod.rs @@ -0,0 +1,264 @@ +use irc::client::prelude::*; + +use std::time::Duration; +use std::sync::Mutex; + +use time; +use chrono::NaiveDateTime; +use humantime::format_duration; + +use plugin::*; + +use failure::Fail; +use failure::ResultExt; +use error::ErrorKind as FrippyErrorKind; +use error::FrippyError; +use self::error::*; + +pub mod database; +use self::database::Database; + +macro_rules! try_lock { + ( $m:expr ) => { + match $m.lock() { + Ok(guard) => guard, + Err(poisoned) => poisoned.into_inner(), + } + } +} + +#[derive(PluginName, Default)] +pub struct Tell<T: Database> { + tells: Mutex<T>, +} + +impl<T: Database> Tell<T> { + pub fn new(db: T) -> Tell<T> { + Tell { + tells: Mutex::new(db), + } + } + + fn tell_command( + &self, + client: &IrcClient, + command: PluginCommand, + ) -> Result<String, TellError> { + if command.tokens.len() < 2 { + return Ok(self.invalid_command(client)); + } + + let receiver = &command.tokens[0]; + let sender = command.source; + + if receiver.eq_ignore_ascii_case(client.current_nickname()) { + return Ok(String::from("I am right here!")); + } + + if receiver.eq_ignore_ascii_case(&sender) { + return Ok(String::from("That's your name!")); + } + + if let Some(channels) = client.list_channels() { + for channel in channels { + if let Some(users) = client.list_users(&channel) { + if users + .iter() + .any(|u| u.get_nickname().eq_ignore_ascii_case(&receiver)) + { + return Ok(format!("{} is currently online.", receiver)); + } + } + } + } + + let tm = time::now().to_timespec(); + let message = command.tokens[1..].join(" "); + let tell = database::NewTellMessage { + sender: &sender, + receiver: &receiver.to_lowercase(), + time: NaiveDateTime::from_timestamp(tm.sec, 0u32), + message: &message, + }; + + try_lock!(self.tells).insert_tell(&tell)?; + + Ok(String::from("Got it!")) + } + + fn on_namelist( + &self, + client: &IrcClient, + channel: &str, + ) -> Result<(), FrippyError> { + let receivers = try_lock!(self.tells) + .get_receivers() + .context(FrippyErrorKind::Tell)?; + + if let Some(users) = client.list_users(channel) { + debug!("Outstanding tells for {:?}", receivers); + + for receiver in users + .iter() + .map(|u| u.get_nickname()) + .filter(|u| receivers.iter().any(|r| r == &u.to_lowercase())) + { + self.send_tells(client, receiver)?; + } + + Ok(()) + } else { + Ok(()) + } + } + fn send_tells(&self, client: &IrcClient, receiver: &str) -> Result<(), FrippyError> { + if client.current_nickname() == receiver { + return Ok(()); + } + + let mut tells = try_lock!(self.tells); + + let tell_messages = match tells.get_tells(&receiver.to_lowercase()) { + Ok(t) => t, + Err(e) => { + // This warning only occurs if frippy is built without a database + #[allow(unreachable_patterns)] + return match e.kind() { + ErrorKind::NotFound => Ok(()), + _ => Err(e.context(FrippyErrorKind::Tell))?, + }; + } + }; + + for tell in tell_messages { + let now = Duration::new(time::now().to_timespec().sec as u64, 0); + let dur = now - Duration::new(tell.time.timestamp() as u64, 0); + let human_dur = format_duration(dur); + + client + .send_notice( + receiver, + &format!( + "Tell from {} {} ago: {}", + tell.sender, human_dur, tell.message + ), + ) + .context(FrippyErrorKind::Connection)?; + + debug!( + "Sent {:?} from {:?} to {:?}", + tell.message, tell.sender, receiver + ); + } + + tells + .delete_tells(&receiver.to_lowercase()) + .context(FrippyErrorKind::Tell)?; + + Ok(()) + } + + fn invalid_command(&self, client: &IrcClient) -> String { + format!( + "Incorrect Command. \ + Send \"{} tell help\" for help.", + client.current_nickname() + ) + } + + fn help(&self, client: &IrcClient) -> String { + format!( + "usage: {} tell user message\r\n\ + example: {0} tell Foobar Hello!", + client.current_nickname() + ) + } +} + +impl<T: Database> Plugin for Tell<T> { + fn execute(&self, client: &IrcClient, message: &Message) -> ExecutionStatus { + let res = match message.command { + Command::JOIN(_, _, _) => self.send_tells(client, message.source_nickname().unwrap()), + Command::NICK(ref nick) => self.send_tells(client, nick), + Command::Response(resp, ref chan_info, _) => { + if resp == Response::RPL_NAMREPLY { + debug!("NAMREPLY info: {:?}", chan_info); + + self.on_namelist( + client, + &chan_info[chan_info.len() - 1], + ) + } else { + Ok(()) + } + } + _ => Ok(()), + }; + + match res { + Ok(_) => ExecutionStatus::Done, + Err(e) => ExecutionStatus::Err(e), + } + } + + fn execute_threaded(&self, _: &IrcClient, _: &Message) -> Result<(), FrippyError> { + panic!("Tell should not use threading") + } + + fn command(&self, client: &IrcClient, command: PluginCommand) -> Result<(), FrippyError> { + if command.tokens.is_empty() { + return Ok(client + .send_notice(&command.source, &self.invalid_command(client)) + .context(FrippyErrorKind::Connection)?); + } + + let sender = command.source.to_owned(); + + Ok(match command.tokens[0].as_ref() { + "help" => client + .send_notice(&command.source, &self.help(client)) + .context(FrippyErrorKind::Connection) + .into(), + _ => match self.tell_command(client, command) { + Ok(msg) => client + .send_notice(&sender, &msg) + .context(FrippyErrorKind::Connection), + Err(e) => client + .send_notice(&sender, &e.to_string()) + .context(FrippyErrorKind::Connection) + .into(), + }, + }?) + } + + fn evaluate(&self, _: &IrcClient, _: PluginCommand) -> Result<String, String> { + Err(String::from("This Plugin does not implement any commands.")) + } +} + +use std::fmt; +impl<T: Database> fmt::Debug for Tell<T> { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "Tell {{ ... }}") + } +} + +pub mod error { + #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)] + #[error = "TellError"] + pub enum ErrorKind { + /// Not found command error + #[fail(display = "Tell was not found")] + NotFound, + + /// MySQL error + #[cfg(feature = "mysql")] + #[fail(display = "Failed to execute MySQL Query")] + MysqlError, + + /// No connection error + #[cfg(feature = "mysql")] + #[fail(display = "No connection to the database")] + NoConnection, + } +} diff --git a/src/plugins/url.rs b/src/plugins/url.rs new file mode 100644 index 0000000..bff840f --- /dev/null +++ b/src/plugins/url.rs @@ -0,0 +1,130 @@ +extern crate htmlescape; +extern crate regex; + +use irc::client::prelude::*; + +use self::regex::Regex; + +use plugin::*; +use utils; + +use self::error::*; +use error::FrippyError; +use error::ErrorKind as FrippyErrorKind; +use failure::Fail; +use failure::ResultExt; + +lazy_static! { + static ref RE: Regex = Regex::new(r"(^|\s)(https?://\S+)").unwrap(); +} + +#[derive(PluginName, Debug)] +pub struct Url { + max_kib: usize, +} + +impl Url { + /// If a file is larger than `max_kib` KiB the download is stopped + pub fn new(max_kib: usize) -> Url { + Url { max_kib: max_kib } + } + + fn grep_url(&self, msg: &str) -> Option<String> { + let captures = RE.captures(msg)?; + debug!("Url captures: {:?}", captures); + + Some(captures.get(2)?.as_str().to_owned()) + } + + fn get_title<'a>(&self, body: &str) -> Result<String, UrlError> { + let title = body.find("<title") + .map(|tag| { + body[tag..] + .find('>') + .map(|offset| tag + offset + 1) + .map(|start| { + body[start..] + .find("</title>") + .map(|offset| start + offset) + .map(|end| &body[start..end]) + }) + }) + .and_then(|s| s.and_then(|s| s)) + .ok_or(ErrorKind::MissingTitle)?; + + debug!("Title: {:?}", title); + + htmlescape::decode_html(title).map_err(|_| ErrorKind::HtmlDecoding.into()) + } + + fn url(&self, text: &str) -> Result<String, UrlError> { + let url = self.grep_url(text).ok_or(ErrorKind::MissingUrl)?; + let body = utils::download(&url, Some(self.max_kib)).context(ErrorKind::Download)?; + + let title = self.get_title(&body)?; + + Ok(title.replace('\n', "|").replace('\r', "|")) + } +} + +impl Plugin for Url { + fn execute(&self, _: &IrcClient, message: &Message) -> ExecutionStatus { + match message.command { + Command::PRIVMSG(_, ref msg) => if RE.is_match(msg) { + ExecutionStatus::RequiresThread + } else { + ExecutionStatus::Done + }, + _ => ExecutionStatus::Done, + } + } + + fn execute_threaded(&self, client: &IrcClient, message: &Message) -> Result<(), FrippyError> { + Ok(match message.command { + Command::PRIVMSG(_, ref content) => match self.url(content) { + Ok(title) => client + .send_privmsg(message.response_target().unwrap(), &title) + .context(FrippyErrorKind::Connection)?, + Err(e) => Err(e).context(FrippyErrorKind::Url)?, + }, + _ => (), + }) + } + + fn command(&self, client: &IrcClient, command: PluginCommand) -> Result<(), FrippyError> { + Ok(client + .send_notice( + &command.source, + "This Plugin does not implement any commands.", + ) + .context(FrippyErrorKind::Connection)?) + } + + fn evaluate(&self, _: &IrcClient, command: PluginCommand) -> Result<String, String> { + self.url(&command.tokens[0]) + .map_err(|e| e.cause().unwrap().to_string()) + } +} + +pub mod error { + /// A URL plugin error + #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)] + #[error = "UrlError"] + pub enum ErrorKind { + /// A download error + #[fail(display = "A download error occured")] + Download, + + /// Missing URL error + #[fail(display = "No URL was found")] + MissingUrl, + + /// Missing title error + #[fail(display = "No title was found")] + MissingTitle, + + /// Html decoding error + #[fail(display = "Failed to decode Html characters")] + HtmlDecoding, + } +} diff --git a/src/utils.rs b/src/utils.rs new file mode 100644 index 0000000..06156be --- /dev/null +++ b/src/utils.rs @@ -0,0 +1,65 @@ +use std::str; +use std::io::{self, Read}; + +use reqwest::Client; +use reqwest::header::Connection; + +use failure::ResultExt; +use self::error::{DownloadError, ErrorKind}; + +/// Downloads the file and converts it to a String. +/// Any invalid bytes are converted to a replacement character. +/// +/// The error indicated either a failed download or that the DownloadLimit was reached +pub fn download(url: &str, max_kib: Option<usize>) -> Result<String, DownloadError> { + let mut response = Client::new() + .get(url) + .header(Connection::close()) + .send() + .context(ErrorKind::Connection)?; + + // 100 kibibyte buffer + let mut buf = [0; 100 * 1024]; + let mut written = 0; + let mut bytes = Vec::new(); + + // Read until we reach EOF or max_kib KiB + loop { + let len = match response.read(&mut buf) { + Ok(0) => break, + Ok(len) => len, + Err(ref e) if e.kind() == io::ErrorKind::Interrupted => continue, + Err(e) => Err(e).context(ErrorKind::Read)?, + }; + + bytes.extend_from_slice(&buf); + written += len; + + // Check if the file is too large to download + if let Some(max_kib) = max_kib { + if written > max_kib * 1024 { + Err(ErrorKind::DownloadLimit)?; + } + } + } + + Ok(String::from_utf8_lossy(&bytes).into_owned()) +} + +pub mod error { + #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)] + #[error = "DownloadError"] + pub enum ErrorKind { + /// Connection Error + #[fail(display = "A connection error has occured")] + Connection, + + /// Read Error + #[fail(display = "A read error has occured")] + Read, + + /// Reached download limit error + #[fail(display = "Reached download limit")] + DownloadLimit, + } +} |
