From 32e59aa3cc6567b909ab995a4e69beb12bdf0322 Mon Sep 17 00:00:00 2001 From: Jokler Date: Tue, 13 Feb 2018 21:37:58 +0100 Subject: Add download function & fromurl command --- src/plugins/factoids/database.rs | 4 +-- src/plugins/factoids/mod.rs | 63 +++++++++++++++++++++++++++++----------- src/plugins/factoids/sandbox.lua | 1 + src/plugins/factoids/utils.rs | 13 +++++++++ src/plugins/url.rs | 61 ++------------------------------------ 5 files changed, 64 insertions(+), 78 deletions(-) create mode 100644 src/plugins/factoids/utils.rs (limited to 'src/plugins') diff --git a/src/plugins/factoids/database.rs b/src/plugins/factoids/database.rs index cb6f422..b612d6f 100644 --- a/src/plugins/factoids/database.rs +++ b/src/plugins/factoids/database.rs @@ -106,7 +106,7 @@ impl Database for MysqlConnection { .execute(self) { Ok(_) => DbResponse::Success, Err(e) => { - error!("DB Insertion Error: {:?}", e); + error!("DB Insertion Error: \"{}\"", e); DbResponse::Failed("Failed to add factoid") } } @@ -131,7 +131,7 @@ impl Database for MysqlConnection { } } Err(e) => { - error!("DB Deletion Error: {:?}", e); + error!("DB Deletion Error: \"{}\"", e); DbResponse::Failed("Failed to delete factoid") } } diff --git a/src/plugins/factoids/mod.rs b/src/plugins/factoids/mod.rs index d67416f..f857f85 100644 --- a/src/plugins/factoids/mod.rs +++ b/src/plugins/factoids/mod.rs @@ -14,6 +14,9 @@ use plugin::*; pub mod database; use self::database::{Database, DbResponse}; +mod utils; +use self::utils::download; + static LUA_SANDBOX: &'static str = include_str!("sandbox.lua"); #[derive(PluginName)] @@ -35,6 +38,24 @@ impl Factoids { 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, server: &IrcServer, command: &mut PluginCommand) -> Result<(), IrcError> { if command.tokens.len() < 2 { return self.invalid_command(server, command); @@ -42,24 +63,27 @@ impl Factoids { let name = command.tokens.remove(0); let content = command.tokens.join(" "); - let count = match try_lock!(self.factoids).count(&name) { - Ok(c) => c, - Err(e) => return server.send_notice(&command.source, e), - }; - let tm = time::now().to_timespec(); + match self.create_factoid(&name, &content, &command.source) { + Ok(v) => server.send_notice(&command.source, v), + Err(e) => server.send_notice(&command.source, e), + } + } - let factoid = database::NewFactoid { - name: &name, - idx: count, - content: &content, - author: &command.source, - created: NaiveDateTime::from_timestamp(tm.sec, tm.nsec as u32), - }; + fn from_url(&self, server: &IrcServer, command: &mut PluginCommand) -> Result<(), IrcError> { + if command.tokens.len() < 2 { + return self.invalid_command(server, command); + } - match try_lock!(self.factoids).insert(&factoid) { - DbResponse::Success => server.send_notice(&command.source, "Successfully added"), - DbResponse::Failed(e) => server.send_notice(&command.source, &e), + 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) => server.send_notice(&command.source, v), + Err(e) => server.send_notice(&command.source, e), + } + } else { + server.send_notice(&command.source, "Failed to download.") } } @@ -110,7 +134,10 @@ impl Factoids { let factoid = match try_lock!(self.factoids).get(name, idx) { Some(v) => v, - None => return server.send_notice(&command.source, &format!("{}~{} does not exist", name, idx)), + None => { + return server.send_notice(&command.source, + &format!("{}~{} does not exist", name, idx)) + } }; server.send_privmsg(&command.target, @@ -194,7 +221,7 @@ impl Factoids { } else { match self.run_lua(&name, &factoid, &command) { Ok(v) => v, - Err(e) => format!("{}", e), + Err(e) => format!("\"{}\"", e), } } } else { @@ -222,6 +249,7 @@ impl Factoids { let globals = lua.globals(); globals.set("factoid", code)?; + globals.set("download", lua.create_function(download))?; globals.set("args", args)?; globals.set("input", command.tokens.join(" "))?; globals.set("user", command.source.clone())?; @@ -274,6 +302,7 @@ impl Plugin for Factoids { let sub_command = command.tokens.remove(0); match sub_command.as_ref() { "add" => self.add(server, &mut command), + "fromurl" => self.from_url(server, &mut command), "remove" => self.remove(server, &mut command), "get" => self.get(server, &command), "info" => self.info(server, &command), diff --git a/src/plugins/factoids/sandbox.lua b/src/plugins/factoids/sandbox.lua index 2fff150..282a903 100644 --- a/src/plugins/factoids/sandbox.lua +++ b/src/plugins/factoids/sandbox.lua @@ -18,6 +18,7 @@ local env = { print = send, input = input, user = user, channel = channel, + request = download, pairs = pairs, table = table, string = string, diff --git a/src/plugins/factoids/utils.rs b/src/plugins/factoids/utils.rs new file mode 100644 index 0000000..ad25eef --- /dev/null +++ b/src/plugins/factoids/utils.rs @@ -0,0 +1,13 @@ +extern crate reqwest; + +use utils; +use super::rlua::prelude::*; + +use self::LuaError::RuntimeError; + +pub fn download(_: &Lua, url: String) -> Result { + match utils::download(1024, &url) { + Some(v) => Ok(v), + None => Err(RuntimeError(format!("Failed to download {}", url))), + } +} diff --git a/src/plugins/url.rs b/src/plugins/url.rs index f6e06f3..6f4a68f 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::Error as 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,58 +38,6 @@ impl Url { } } - fn download(&self, url: &str) -> Option { - 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); - return None; - } - } - } - fn url(&self, server: &IrcServer, message: &str, target: &str) -> Result<(), IrcError> { let url = match self.grep_url(message) { Some(url) => url, @@ -104,7 +47,7 @@ impl Url { }; - match self.download(&url) { + match utils::download(self.max_kib, &url) { Some(body) => { let doc = Document::from(body.as_ref()); -- cgit v1.2.3-70-g09d2