diff --git a/src/commands/misc/sauce.rs b/src/commands/misc/sauce.rs index 116295f..d1058ee 100644 --- a/src/commands/misc/sauce.rs +++ b/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 = - source_msg.attachments.into_iter().map(|a| a.url).collect(); + log::debug!("Getting attachments..."); + let mut attachment_urls: Vec = 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::>(); 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(()); } diff --git a/src/messages/sauce.rs b/src/messages/sauce.rs index 678327b..78ffa44 100644 --- a/src/messages/sauce.rs +++ b/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")) diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 36959b1..6fe164e 100644 --- a/src/utils/mod.rs +++ b/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 { - 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 { + 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; +}