diff options
Diffstat (limited to 'src/plugins')
| -rw-r--r-- | src/plugins/factoids/database.rs | 151 | ||||
| -rw-r--r-- | src/plugins/factoids/mod.rs | 330 | ||||
| -rw-r--r-- | src/plugins/factoids/sandbox.lua | 51 | ||||
| -rw-r--r-- | src/plugins/factoids/utils.rs | 21 | ||||
| -rw-r--r-- | src/plugins/mod.rs | 3 | ||||
| -rw-r--r-- | src/plugins/url.rs | 61 |
6 files changed, 559 insertions, 58 deletions
diff --git a/src/plugins/factoids/database.rs b/src/plugins/factoids/database.rs new file mode 100644 index 0000000..b612d6f --- /dev/null +++ b/src/plugins/factoids/database.rs @@ -0,0 +1,151 @@ +#[cfg(feature = "mysql")] +extern crate dotenv; + +use std::collections::HashMap; + +#[cfg(feature = "mysql")] +use diesel::prelude::*; + +#[cfg(feature = "mysql")] +use diesel::mysql::MysqlConnection; + +use chrono::NaiveDateTime; + +pub enum DbResponse { + Success, + Failed(&'static str), +} + +#[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(feature = "mysql")] +use self::mysql::factoids; +#[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(&mut self, factoid: &NewFactoid) -> DbResponse; + fn get(&self, name: &str, idx: i32) -> Option<Factoid>; + fn delete(&mut self, name: &str, idx: i32) -> DbResponse; + fn count(&self, name: &str) -> Result<i32, &'static str>; +} + +// HashMap +impl Database for HashMap<(String, i32), Factoid> { + fn insert(&mut self, factoid: &NewFactoid) -> DbResponse { + 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 = String::from(factoid.name.clone()); + match self.insert((name, factoid.idx), factoid) { + None => DbResponse::Success, + Some(_) => DbResponse::Failed("Factoid was overwritten"), + } + } + + fn get(&self, name: &str, idx: i32) -> Option<Factoid> { + self.get(&(String::from(name), idx)).map(|f| f.clone()) + } + + fn delete(&mut self, name: &str, idx: i32) -> DbResponse { + match self.remove(&(String::from(name), idx)) { + Some(_) => DbResponse::Success, + None => DbResponse::Failed("Factoid not found"), + } + } + + fn count(&self, name: &str) -> Result<i32, &'static str> { + Ok(self.iter() + .filter(|&(&(ref n, _), _)| n == name) + .count() as i32) + } +} + +// Diesel automatically define the factoids module as public. +// For now this is how we keep it private. +#[cfg(feature = "mysql")] +mod mysql { + table! { + factoids (name, idx) { + name -> Varchar, + idx -> Integer, + content -> Text, + author -> Varchar, + created -> Timestamp, + } + } +} + +#[cfg(feature = "mysql")] +impl Database for MysqlConnection { + fn insert(&mut self, factoid: &NewFactoid) -> DbResponse { + use diesel; + match diesel::insert_into(factoids::table) + .values(factoid) + .execute(self) { + Ok(_) => DbResponse::Success, + Err(e) => { + error!("DB Insertion Error: \"{}\"", e); + DbResponse::Failed("Failed to add factoid") + } + } + } + + fn get(&self, name: &str, idx: i32) -> Option<Factoid> { + factoids::table.find((name, idx)).first(self).ok() + } + + fn delete(&mut self, name: &str, idx: i32) -> DbResponse { + use diesel; + use self::factoids::columns; + match diesel::delete(factoids::table + .filter(columns::name.eq(name)) + .filter(columns::idx.eq(idx))) + .execute(self) { + Ok(v) => { + if v > 0 { + DbResponse::Success + } else { + DbResponse::Failed("Could not find any factoid with that name") + } + } + Err(e) => { + error!("DB Deletion Error: \"{}\"", e); + DbResponse::Failed("Failed to delete factoid") + } + } + } + + fn count(&self, name: &str) -> Result<i32, &'static str> { + let count: Result<i64, _> = factoids::table + .filter(factoids::columns::name.eq(name)) + .count() + .first(self); + + match count { + Ok(c) => Ok(c as i32), + Err(_) => Err("Database Error"), + } + } +} diff --git a/src/plugins/factoids/mod.rs b/src/plugins/factoids/mod.rs new file mode 100644 index 0000000..49ace10 --- /dev/null +++ b/src/plugins/factoids/mod.rs @@ -0,0 +1,330 @@ +extern crate rlua; + +use std::fmt; +use std::str::FromStr; +use std::sync::Mutex; +use self::rlua::prelude::*; +use irc::client::prelude::*; +use irc::error::IrcError; + +use time; +use chrono::NaiveDateTime; + +use plugin::*; +pub mod database; +use self::database::{Database, DbResponse}; + +mod utils; +use self::utils::*; + +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, &str> { + let count = try_lock!(self.factoids).count(&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, tm.nsec as u32), + }; + + match try_lock!(self.factoids).insert(&factoid) { + DbResponse::Success => Ok("Successfully added"), + DbResponse::Failed(e) => Err(e), + } + } + + fn add(&self, client: &IrcClient, command: &mut PluginCommand) -> Result<(), IrcError> { + if command.tokens.len() < 2 { + return self.invalid_command(client, command); + } + + let name = command.tokens.remove(0); + let content = command.tokens.join(" "); + + match self.create_factoid(&name, &content, &command.source) { + Ok(v) => client.send_notice(&command.source, v), + Err(e) => client.send_notice(&command.source, e), + } + } + + fn from_url(&self, client: &IrcClient, command: &mut PluginCommand) -> Result<(), IrcError> { + if command.tokens.len() < 2 { + return self.invalid_command(client, command); + } + + let name = command.tokens.remove(0); + let url = &command.tokens[0]; + if let Some(content) = ::utils::download(1024, url) { + match self.create_factoid(&name, &content, &command.source) { + Ok(v) => client.send_notice(&command.source, v), + Err(e) => client.send_notice(&command.source, e), + } + } else { + client.send_notice(&command.source, "Failed to download.") + } + } + + fn remove(&self, client: &IrcClient, command: &mut PluginCommand) -> Result<(), IrcError> { + if command.tokens.len() < 1 { + return self.invalid_command(client, command); + } + + let name = command.tokens.remove(0); + let count = match try_lock!(self.factoids).count(&name) { + Ok(c) => c, + Err(e) => return client.send_notice(&command.source, e), + }; + + match try_lock!(self.factoids).delete(&name, count - 1) { + DbResponse::Success => client.send_notice(&command.source, "Successfully removed"), + DbResponse::Failed(e) => client.send_notice(&command.source, &e), + } + } + + fn get(&self, client: &IrcClient, command: &PluginCommand) -> Result<(), IrcError> { + + let (name, idx) = match command.tokens.len() { + 0 => return self.invalid_command(client, command), + 1 => { + let name = &command.tokens[0]; + let count = match try_lock!(self.factoids).count(name) { + Ok(c) => c, + Err(e) => return client.send_notice(&command.source, e), + }; + + if count < 1 { + return client.send_notice(&command.source, &format!("{} does not exist", name)); + } + + (name, count - 1) + } + _ => { + let name = &command.tokens[0]; + let idx = match i32::from_str(&command.tokens[1]) { + Ok(i) => i, + Err(_) => return client.send_notice(&command.source, "Invalid index"), + }; + + (name, idx) + } + }; + + let factoid = match try_lock!(self.factoids).get(name, idx) { + Some(v) => v, + None => { + return client.send_notice(&command.source, + &format!("{}~{} does not exist", name, idx)) + } + }; + + let message = factoid.content.replace("\n", "|").replace("\r", ""); + + client.send_privmsg(&command.target, + &format!("{}: {}", factoid.name, message)) + } + + fn info(&self, client: &IrcClient, command: &PluginCommand) -> Result<(), IrcError> { + + match command.tokens.len() { + 0 => self.invalid_command(client, command), + 1 => { + let name = &command.tokens[0]; + let count = match try_lock!(self.factoids).count(name) { + Ok(c) => c, + Err(e) => return client.send_notice(&command.source, e), + }; + + match count { + 0 => client.send_notice(&command.source, &format!("{} does not exist", name)), + 1 => { + client.send_privmsg(&command.target, + &format!("There is 1 version of {}", name)) + } + _ => { + client.send_privmsg(&command.target, + &format!("There are {} versions of {}", count, name)) + } + } + } + _ => { + let name = &command.tokens[0]; + let idx = match i32::from_str(&command.tokens[1]) { + Ok(i) => i, + Err(_) => return client.send_notice(&command.source, "Invalid index"), + }; + + let factoid = match try_lock!(self.factoids).get(name, idx) { + Some(v) => v, + None => { + return client.send_notice(&command.source, + &format!("{}~{} does not exist", name, idx)) + } + }; + + client.send_privmsg(&command.target, + &format!("{}: Added by {} at {} UTC", + name, + factoid.author, + factoid.created)) + } + + } + } + + fn exec(&self, + client: &IrcClient, + mut command: PluginCommand, + error: bool) + -> Result<(), IrcError> { + if command.tokens.len() < 1 { + self.invalid_command(client, &command) + + } else { + let name = command.tokens.remove(0); + let count = match try_lock!(self.factoids).count(&name) { + Ok(c) => c, + Err(e) => return client.send_notice(&command.source, e), + }; + + let factoid = match try_lock!(self.factoids).get(&name, count - 1) { + Some(v) => v.content, + None if error => return self.invalid_command(client, &command), + None => return Ok(()), + }; + + let value = &if factoid.starts_with(">") { + let factoid = String::from(&factoid[1..]); + + if factoid.starts_with(">") { + factoid + } else { + match self.run_lua(&name, &factoid, &command) { + Ok(v) => v, + Err(e) => format!("\"{}\"", e), + } + } + } else { + factoid + }; + + client.send_privmsg(&command.target, &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("|")) + } + + fn invalid_command(&self, client: &IrcClient, command: &PluginCommand) -> Result<(), IrcError> { + client.send_notice(&command.source, "Invalid Command") + } +} + +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<(), IrcError> { + 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_string(), + target: message.response_target().unwrap().to_string(), + tokens: t, + }; + + self.exec(client, c, false) + + } else { + Ok(()) + } + } + + fn command(&self, client: &IrcClient, mut command: PluginCommand) -> Result<(), IrcError> { + if command.tokens.is_empty() { + return self.invalid_command(client, &command); + } + + let sub_command = command.tokens.remove(0); + match sub_command.as_ref() { + "add" => self.add(client, &mut command), + "fromurl" => self.from_url(client, &mut command), + "remove" => self.remove(client, &mut command), + "get" => self.get(client, &command), + "info" => self.info(client, &command), + "exec" => self.exec(client, command, true), + _ => self.invalid_command(client, &command), + } + } + + 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 {{ ... }}") + } +} diff --git a/src/plugins/factoids/sandbox.lua b/src/plugins/factoids/sandbox.lua new file mode 100644 index 0000000..3fde65e --- /dev/null +++ b/src/plugins/factoids/sandbox.lua @@ -0,0 +1,51 @@ +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 env = { print = send, + println = sendln, + args = args, + input = input, + user = user, + channel = channel, + request = download, + pairs = pairs, + table = table, + string = string, + tostring = tostring, + tonumber = tonumber, + math = math } + +local f, e = load(factoid, nil, nil, env) + +-- 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 + +-- 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..fc86fb3 --- /dev/null +++ b/src/plugins/factoids/utils.rs @@ -0,0 +1,21 @@ +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(1024, &url) { + Some(v) => Ok(v), + None => Err(RuntimeError(format!("Failed to download {}", url))), + } +} + +pub fn sleep(_: &Lua, dur: u64) -> Result<(), LuaError> { + thread::sleep(Duration::from_millis(dur)); + Ok(()) +} diff --git a/src/plugins/mod.rs b/src/plugins/mod.rs index 834e5b1..5b32efd 100644 --- a/src/plugins/mod.rs +++ b/src/plugins/mod.rs @@ -4,6 +4,7 @@ mod url; mod emoji; mod tell; mod currency; +mod factoids; mod keepnick; pub use self::help::Help; @@ -11,4 +12,6 @@ pub use self::url::Url; pub use self::emoji::Emoji; pub use self::tell::Tell; pub use self::currency::Currency; +pub use self::factoids::Factoids; +pub use self::factoids::database; pub use self::keepnick::KeepNick; diff --git a/src/plugins/url.rs b/src/plugins/url.rs index 455aa4e..df4fdf2 100644 --- a/src/plugins/url.rs +++ b/src/plugins/url.rs @@ -1,5 +1,4 @@ extern crate regex; -extern crate reqwest; extern crate select; use irc::client::prelude::*; @@ -7,15 +6,11 @@ use irc::error::IrcError; use self::regex::Regex; -use std::str; -use std::io::{self, Read}; -use self::reqwest::Client; -use self::reqwest::header::Connection; - use self::select::document::Document; use self::select::predicate::Name; use plugin::*; +use utils; lazy_static! { static ref RE: Regex = Regex::new(r"(^|\s)(https?://\S+)").unwrap(); @@ -43,64 +38,14 @@ impl Url { } } - fn download(&self, url: &str) -> Option<String> { - let response = Client::new().get(url).header(Connection::close()).send(); - - match response { - Ok(mut response) => { - let mut body = String::new(); - - // 500 kilobyte buffer - let mut buf = [0; 500 * 1000]; - let mut written = 0; - // 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) => { - debug!("Download from {:?} failed: {}", url, e); - return None; - } - }; - - let slice = match str::from_utf8(&buf[..len]) { - Ok(slice) => slice, - Err(e) => { - debug!("Failed to read bytes from {:?} as UTF8: {}", url, e); - return None; - } - }; - - body.push_str(slice); - written += len; - - // Check if the file is too large to download - if written > self.max_kib * 1024 { - debug!( - "Stopping download - File from {:?} is larger than {} KiB", - url, self.max_kib - ); - return None; - } - } - Some(body) // once told me - } - Err(e) => { - debug!("Bad response from {:?}: ({})", url, e); - None - } - } - } - fn url(&self, text: &str) -> Result<String, &str> { let url = match self.grep_url(text) { Some(url) => url, None => return Err("No Url was found."), }; - match self.download(&url) { + + match utils::download(self.max_kib, &url) { Some(body) => { let doc = Document::from(body.as_ref()); if let Some(title) = doc.find(Name("title")).next() { |