Add download function & fromurl command

8 files changed, 125 insertions, 79 deletions

diff --git a/bin/main.rs b/bin/main.rs
index 9373afd..fd0b4ab 100644
--- a/bin/main.rs
+++ b/bin/main.rs
@@ -144,7 +144,7 @@ fn main() {
         if let Some(disabled_plugins) = disabled_plugins {
             for name in disabled_plugins {
                 if let None = bot.remove_plugin(name) {
-                    error!("{:?} was not found - could not disable", name);
+                    error!("\"{}\" was not found - could not disable", name);
                 }
             }
         }
diff --git a/src/lib.rs b/src/lib.rs
index 8b55dab..186d35c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -51,6 +51,7 @@ extern crate time;
 
 pub mod plugin;
 pub mod plugins;
+pub mod utils;
 
 use std::collections::HashMap;
 use std::fmt;
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<T: Database> 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, server: &IrcServer, command: &mut PluginCommand) -> Result<(), IrcError> {
         if command.tokens.len() < 2 {
             return self.invalid_command(server, command);
@@ -42,24 +63,27 @@ impl<T: Database> Factoids<T> {
 
         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<T: Database> Factoids<T> {
 
         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<T: Database> Factoids<T> {
                 } 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<T: Database> Factoids<T> {
         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<T: Database> Plugin for Factoids<T> {
         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<String, LuaError> {
+    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<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);
-                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());
diff --git a/src/utils.rs b/src/utils.rs
new file mode 100644
index 0000000..59a5a54
--- /dev/null
+++ b/src/utils.rs
@@ -0,0 +1,59 @@
+extern crate reqwest;
+
+use std::str;
+use std::io::{self, Read};
+
+use self::reqwest::Client;
+use self::reqwest::header::Connection;
+
+pub fn download(max_kib: usize, 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 > max_kib * 1024 {
+                    debug!("Stopping download - File from {:?} is larger than {} KiB", url, max_kib);
+                    return None;
+                }
+
+            }
+            Some(body)
+        }
+        Err(e) => {
+            debug!("Bad response from {:?}: ({})", url, e);
+            return None;
+        }
+    }
+}