Fix sauce command image extraction

Signed-off-by: trivernis <trivernis@protonmail.com>

src/commands/misc/sauce.rs

@@ -1,5 +1,5 @@
 use crate::messages::sauce::show_sauce_menu;
-use crate::utils::get_previous_message_or_reply;
+use crate::utils::{get_previous_message_or_reply, is_image, is_video};
 
 use sauce_api::Sauce;
 
@@ -12,6 +12,7 @@ use serenity::model::channel::Message;
 #[command]
 #[description("Searches for the source of a previously posted image or an image replied to.")]
 #[usage("")]
+#[aliases("source")]
 async fn sauce(ctx: &Context, msg: &Message) -> CommandResult {
     log::debug!("Got sauce command");
     let source_msg = get_previous_message_or_reply(ctx, msg).await?;
@@ -23,14 +24,20 @@ async fn sauce(ctx: &Context, msg: &Message) -> CommandResult {
     }
     let source_msg = source_msg.unwrap();
     log::trace!("Source message is {:?}", source_msg);
-    let mut attachment_urls: Vec<String> =
-        source_msg.attachments.into_iter().map(|a| a.url).collect();
+    log::debug!("Getting attachments...");
+    let mut attachment_urls: Vec<String> = source_msg
+        .attachments
+        .into_iter()
+        .map(|a| a.url)
+        .filter(|url| is_image(url) || is_video(url))
+        .collect();
 
+    log::debug!("Getting embedded images...");
     let mut embed_images = source_msg
         .embeds
         .into_iter()
-        .filter_map(|e| e.thumbnail)
-        .map(|t| t.url)
+        .filter_map(|e| e.thumbnail.map(|t| t.url).or(e.image.map(|i| i.url)))
+        .filter(|url| is_image(url) || is_video(url))
         .collect::<Vec<String>>();
 
     attachment_urls.append(&mut embed_images);
@@ -38,7 +45,9 @@ async fn sauce(ctx: &Context, msg: &Message) -> CommandResult {
 
     if attachment_urls.is_empty() {
         log::debug!("No images in source image");
-        msg.channel_id.say(ctx, "Images in message found.").await?;
+        msg.channel_id
+            .say(ctx, "I could not find any images in the message.")
+            .await?;
         return Ok(());
     }
 

src/messages/sauce.rs

@@ -1,9 +1,13 @@
 use crate::utils::error::BotResult;
 use crate::utils::get_domain_for_url;
-use sauce_api::SauceResult;
+use sauce_api::{SauceItem, SauceResult};
 use serenity::builder::CreateMessage;
 use serenity::{model::channel::Message, prelude::*};
 use serenity_utils::prelude::*;
+use std::cmp::Ordering;
+
+static MAX_RESULTS: usize = 6;
+static MIN_SIMILARITY: f32 = 50.0;
 
 /// Builds a new sauce menu
 pub async fn show_sauce_menu(
@@ -24,7 +28,15 @@ pub async fn show_sauce_menu(
             },
         )
     } else {
-        Menu::new(ctx, msg, &pages, MenuOptions::default())
+        Menu::new(
+            ctx,
+            msg,
+            &pages,
+            MenuOptions {
+                timeout: 600.,
+                ..Default::default()
+            },
+        )
     };
     menu.run().await?;
 
@@ -32,23 +44,47 @@ pub async fn show_sauce_menu(
 }
 
 /// Creates a single sauce page
-fn create_sauce_page<'a>(result: SauceResult) -> CreateMessage<'a> {
+fn create_sauce_page<'a>(mut result: SauceResult) -> CreateMessage<'a> {
     let mut message = CreateMessage::default();
     let mut description_lines = Vec::new();
     let original = result.original_url;
     description_lines.push(format!("[Original]({})", original));
     description_lines.push(String::new());
+    // sort by similarity
+    result.items.sort_by(|a, b| {
+        if a.similarity > b.similarity {
+            Ordering::Greater
+        } else if a.similarity < b.similarity {
+            Ordering::Less
+        } else {
+            Ordering::Equal
+        }
+    });
+    // display with descending similarity
+    result.items.reverse();
+    let items: Vec<(usize, SauceItem)> = result
+        .items
+        .into_iter()
+        .filter(|i| i.similarity >= MIN_SIMILARITY)
+        .enumerate()
+        .collect();
 
-    for item in result.items {
-        if item.similarity > 70. {
-            description_lines.push(format!(
-                "{}% Similarity: [{}]({})",
-                item.similarity,
-                get_domain_for_url(&item.link).unwrap_or("Source".to_string()),
-                item.link
-            ));
+    if items.is_empty() {
+        description_lines.push("*No Sources found*".to_string());
+    }
+
+    for (i, item) in items {
+        if i >= MAX_RESULTS {
+            break;
         }
+        description_lines.push(format!(
+            "{}% Similarity: [{}]({})",
+            item.similarity,
+            get_domain_for_url(&item.link).unwrap_or("Source".to_string()),
+            item.link
+        ));
     }
+
     message.embed(|e| {
         e.title("Sources")
             .description(description_lines.join("\n"))

src/utils/mod.rs

@@ -41,8 +41,46 @@ pub async fn get_previous_message_or_reply(
 
 /// Returns the domain for a given url
 pub fn get_domain_for_url(url: &str) -> Option<String> {
-    let domain_regex: Regex = Regex::new(r"^(https?://)?(www\.)?((\w+\.)+\w+).*$").unwrap();
-    let captures = domain_regex.captures(url)?;
+    lazy_static::lazy_static! {
+        static ref DOMAIN_REGEX: Regex = Regex::new(r"^(https?://)?(www\.)?((\w+\.)+\w+).*$").unwrap();
+    }
+
+    let captures = DOMAIN_REGEX.captures(url)?;
+
+    captures.get(3).map(|c| c.as_str().to_string())
+}
+
+/// Returns the file for a given domain
+pub fn get_file_name_for_domain(url: &str) -> Option<String> {
+    lazy_static::lazy_static! {
+        static ref FILE_REGEX: Regex = Regex::new(r"^(https?://)?(www\.)?(\w+\.)+\w+([^/]*/)*([^/]+)$").unwrap();
+    }
+
+    let captures = FILE_REGEX.captures(url)?;
 
     captures.get(3).map(|c| c.as_str().to_string())
 }
+
+/// Returns if the given file is an image
+#[inline]
+pub fn is_image(url: &str) -> bool {
+    static IMAGE_EXTENSIONS: &[&str] = &["png", "webp", "jpg", "jpeg", "ico", "gif"];
+    for ext in IMAGE_EXTENSIONS {
+        if url.ends_with(ext) {
+            return true;
+        }
+    }
+    return false;
+}
+
+/// Returns if the given file is an image
+#[inline]
+pub fn is_video(url: &str) -> bool {
+    static IMAGE_EXTENSIONS: &[&str] = &["mp4", "flv", "mkv", "webm"];
+    for ext in IMAGE_EXTENSIONS {
+        if url.ends_with(ext) {
+            return true;
+        }
+    }
+    return false;
+}