From c9fcae05db24e09dabd3cf0be5e45786a6588554 Mon Sep 17 00:00:00 2001 From: trivernis Date: Sat, 10 Apr 2021 18:50:41 +0200 Subject: [PATCH] Add sauce command Signed-off-by: trivernis --- Cargo.lock | 241 +++++++++++++++++++++++++++++++++++- Cargo.toml | 6 +- README.md | 2 + src/commands/misc/mod.rs | 4 +- src/commands/misc/pekofy.rs | 20 +-- src/commands/misc/sauce.rs | 60 +++++++++ src/commands/mod.rs | 2 +- src/messages/mod.rs | 1 + src/messages/sauce.rs | 60 +++++++++ src/utils/context_data.rs | 8 ++ src/utils/error.rs | 3 + src/utils/logging.rs | 1 + src/utils/mod.rs | 30 +++++ 13 files changed, 419 insertions(+), 19 deletions(-) create mode 100644 src/commands/misc/sauce.rs create mode 100644 src/messages/sauce.rs diff --git a/Cargo.lock b/Cargo.lock index bbfbc0b..ee74dd0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -142,6 +142,21 @@ version = "0.13.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "904dfeac50f3cdaba28fc6f57fdcddb75f49ed61346676a78c4ffe55877802fd" +[[package]] +name = "bit-set" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e11e16035ea35e4e5997b393eacbf6f63983188f7a2ad25bfb13465f5ad59de" +dependencies = [ + "bit-vec", +] + +[[package]] +name = "bit-vec" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb" + [[package]] name = "bitflags" version = "0.5.0" @@ -533,6 +548,16 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "futf" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c9c1ce3fa9336301af935ab852c437817d14cd33690446569392e65170aac3b" +dependencies = [ + "mac", + "new_debug_unreachable", +] + [[package]] name = "futures" version = "0.3.13" @@ -712,6 +737,20 @@ dependencies = [ "libc", ] +[[package]] +name = "html5ever" +version = "0.25.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aafcf38a1a36118242d29b92e1b08ef84e67e4a5ed06e0a80be20e6a32bfed6b" +dependencies = [ + "log 0.4.14", + "mac", + "markup5ever", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "http" version = "0.2.3" @@ -981,6 +1020,41 @@ dependencies = [ "serde_json", ] +[[package]] +name = "mac" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c41e0c4fef86961ac6d6f8a82609f55f31b05e4fce149ac5710e439df7619ba4" + +[[package]] +name = "markup5ever" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aae38d669396ca9b707bfc3db254bc382ddb94f57cc5c235f34623a669a01dab" +dependencies = [ + "log 0.4.14", + "phf", + "phf_codegen", + "serde", + "serde_derive", + "serde_json", + "string_cache", + "string_cache_codegen", + "tendril", +] + +[[package]] +name = "markup5ever_rcdom" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f015da43bcd8d4f144559a3423f4591d69b8ce0652c905374da7205df336ae2b" +dependencies = [ + "html5ever", + "markup5ever", + "tendril", + "xml5ever", +] + [[package]] name = "matches" version = "0.1.8" @@ -1113,6 +1187,12 @@ dependencies = [ "tempfile", ] +[[package]] +name = "new_debug_unreachable" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4a24736216ec316047a1fc4252e27dabb04218aa4a3f37c6e7ddbf1f9782b54" + [[package]] name = "ntapi" version = "0.3.6" @@ -1245,6 +1325,44 @@ version = "2.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d4fd5641d01c8f18a23da7b6fe29298ff4b55afcccdf78973b24cf3175fee32e" +[[package]] +name = "phf" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3dfb61232e34fcb633f43d12c58f83c1df82962dcdfa565a4e866ffc17dafe12" +dependencies = [ + "phf_shared", +] + +[[package]] +name = "phf_codegen" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cbffee61585b0411840d3ece935cce9cb6321f01c45477d30066498cd5e1a815" +dependencies = [ + "phf_generator", + "phf_shared", +] + +[[package]] +name = "phf_generator" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "17367f0cc86f2d25802b2c26ee58a7b23faeccf78a396094c13dced0d0182526" +dependencies = [ + "phf_shared", + "rand 0.7.3", +] + +[[package]] +name = "phf_shared" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c00cf8b9eafe68dde5e9eaa2cef8ee84a9336a47d566ec55ca16589633b65af7" +dependencies = [ + "siphasher", +] + [[package]] name = "pin-project" version = "1.0.6" @@ -1334,6 +1452,12 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "precomputed-hash" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "925383efa346730478fb4838dbe9137d2a47675ad789c546d150a6e1dd4ab31c" + [[package]] name = "proc-macro-hack" version = "0.5.19" @@ -1386,6 +1510,7 @@ dependencies = [ "rand_chacha 0.2.2", "rand_core 0.5.1", "rand_hc 0.2.0", + "rand_pcg", ] [[package]] @@ -1456,6 +1581,15 @@ dependencies = [ "rand_core 0.6.2", ] +[[package]] +name = "rand_pcg" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16abd0c1b639e9eb4d7c50c0b8100b0d0f849be2349829c740fe8e6eb4816429" +dependencies = [ + "rand_core 0.5.1", +] + [[package]] name = "rayon" version = "1.5.0" @@ -1612,6 +1746,22 @@ dependencies = [ "zeroize", ] +[[package]] +name = "sauce-api" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca3a8391f67eb2e99020a7f0ee9b8dd53bbcf0944329ab41b3363b0ded373876" +dependencies = [ + "async-trait", + "reqwest", + "select", + "serde", + "serde_json", + "strfmt", + "thiserror", + "urlencoding", +] + [[package]] name = "schannel" version = "0.1.19" @@ -1676,6 +1826,17 @@ dependencies = [ "libc", ] +[[package]] +name = "select" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ee061f90afcc8678bef7a78d0d121683f0ba753f740ff7005f833ec445876b7" +dependencies = [ + "bit-set", + "html5ever", + "markup5ever_rcdom", +] + [[package]] name = "serde" version = "1.0.125" @@ -1779,6 +1940,16 @@ dependencies = [ "serde_repr", ] +[[package]] +name = "serenity_utils" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19368e88b3d6183b52704a2fbc1e5788ad4100876eddaa7217e1c8e1d2d2535c" +dependencies = [ + "serenity", + "tokio", +] + [[package]] name = "sha-1" version = "0.9.4" @@ -1801,6 +1972,12 @@ dependencies = [ "libc", ] +[[package]] +name = "siphasher" +version = "0.3.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cbce6d4507c7e4a3962091436e56e95290cb71fa302d0d270e32130b75fbff27" + [[package]] name = "slab" version = "0.4.2" @@ -1896,6 +2073,37 @@ dependencies = [ "loom", ] +[[package]] +name = "strfmt" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b278b244ef7aa5852b277f52dd0c6cac3a109919e1f6d699adde63251227a30f" + +[[package]] +name = "string_cache" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ddb1139b5353f96e429e1a5e19fbaf663bddedaa06d1dbd49f82e352601209a" +dependencies = [ + "lazy_static", + "new_debug_unreachable", + "phf_shared", + "precomputed-hash", + "serde", +] + +[[package]] +name = "string_cache_codegen" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f24c8e5e19d22a726626f1a5e16fe15b132dcf21d10177fa5a45ce7962996b97" +dependencies = [ + "phf_generator", + "phf_shared", + "proc-macro2", + "quote", +] + [[package]] name = "subtle" version = "2.4.0" @@ -1992,6 +2200,17 @@ dependencies = [ "winapi 0.3.9", ] +[[package]] +name = "tendril" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9ef557cb397a4f0a5a3a628f06515f78563f2209e64d47055d9dc6052bf5e33" +dependencies = [ + "futf", + "mac", + "utf-8", +] + [[package]] name = "term" version = "0.4.6" @@ -2050,7 +2269,7 @@ checksum = "cda74da7e1a664f795bb1f8a87ec406fb89a02522cf6e50620d016add6dbbf5c" [[package]] name = "tobi-rs" -version = "0.1.3" +version = "0.2.0" dependencies = [ "aspotify", "chrono", @@ -2066,10 +2285,12 @@ dependencies = [ "rand 0.8.3", "regex", "reqwest", + "sauce-api", "serde", "serde_derive", "serde_json", "serenity", + "serenity_utils", "songbird", "sysinfo", "thiserror", @@ -2329,6 +2550,12 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "urlencoding" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c9232eb53352b4442e40d7900465dfc534e8cb2dc8f18656fcb2ac16112b5593" + [[package]] name = "utf-8" version = "0.7.5" @@ -2533,6 +2760,18 @@ dependencies = [ "winapi 0.3.9", ] +[[package]] +name = "xml5ever" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b1b52e6e8614d4a58b8e70cf51ec0cc21b256ad8206708bcff8139b5bbd6a59" +dependencies = [ + "log 0.4.14", + "mac", + "markup5ever", + "time", +] + [[package]] name = "xsalsa20poly1305" version = "0.6.0" diff --git a/Cargo.toml b/Cargo.toml index 59f533d..6cd07e4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "tobi-rs" -version = "0.1.3" +version = "0.2.0" authors = ["trivernis "] edition = "2018" @@ -28,4 +28,6 @@ colored = "2.0.0" sysinfo = "0.16.5" database = {path="./database"} reqwest = "0.11.2" -chrono-tz = "0.5.3" \ No newline at end of file +chrono-tz = "0.5.3" +sauce-api = "0.7.1" +serenity_utils = "0.6.1" \ No newline at end of file diff --git a/README.md b/README.md index af57864..d9232b3 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,7 @@ The bot depends on the following APIs - [Discord](https://discord.com/developers/applications): It's a discord bot... - [Spotify](https://developer.spotify.com/documentation/web-api/): To fetch song names to be searched on youtube for music playback - [lyrics.ohv](https://lyricsovh.docs.apiary.io): To fetch lyrics for playing songs +- [SauceNAO](https://saucenao.com): To fetch source information for images ## Dev Dependencies @@ -55,6 +56,7 @@ The required values are: - `DATABASE_URL` (required): Connection uri to the postgres database in the schema `postgres://myuser:mypassword@localhost:5432/database` - `SPOTIFY_CLIENT_ID` (required): Spotify API Client ID - `SPOTIFY_CLIENT_SECRET` (required): Spotify API Client Secret +- `SAUCENAO_API_KEY` (required): SauceNAO API Key - `BOT_PREFIX` (optional): The prefix of the bot. Defaults to `~` if not set. - `LOG_DIR` (optional): Directory to store log files in. Defaults to `logs` in the cwd. diff --git a/src/commands/misc/mod.rs b/src/commands/misc/mod.rs index 03b48eb..0a524d3 100644 --- a/src/commands/misc/mod.rs +++ b/src/commands/misc/mod.rs @@ -3,6 +3,7 @@ use serenity::framework::standard::macros::group; use pekofy::PEKOFY_COMMAND; use ping::PING_COMMAND; use qalc::QALC_COMMAND; +use sauce::SAUCE_COMMAND; use shutdown::SHUTDOWN_COMMAND; use stats::STATS_COMMAND; use time::TIME_COMMAND; @@ -12,11 +13,12 @@ pub(crate) mod help; mod pekofy; mod ping; mod qalc; +mod sauce; mod shutdown; mod stats; mod time; mod timezones; #[group] -#[commands(ping, stats, shutdown, pekofy, time, timezones, qalc)] +#[commands(ping, stats, shutdown, pekofy, time, timezones, qalc, sauce)] pub struct Misc; diff --git a/src/commands/misc/pekofy.rs b/src/commands/misc/pekofy.rs index 4000eb0..d34492c 100644 --- a/src/commands/misc/pekofy.rs +++ b/src/commands/misc/pekofy.rs @@ -1,3 +1,4 @@ +use crate::utils::get_previous_message_or_reply; use rand::prelude::*; use regex::Regex; use serenity::framework::standard::{Args, CommandError, CommandResult}; @@ -24,21 +25,12 @@ async fn pekofy(ctx: &Context, msg: &Message, args: Args) -> CommandResult { let mut content = args.message().to_string(); if args.is_empty() { - if let Some(reference) = &msg.referenced_message { - reference_message = reference.id; - content = reference.content.clone(); - } else { - let messages = msg - .channel_id - .messages(ctx, |ret| ret.before(&msg.id).limit(1)) - .await?; - let reference = messages - .first() - .ok_or(CommandError::from("No message to pekofy"))?; + let reference = get_previous_message_or_reply(ctx, msg) + .await? + .ok_or(CommandError::from("No message to pekofy"))?; + reference_message = reference.id; + content = reference.content; - reference_message = reference.id; - content = reference.content.clone(); - }; let _ = msg.delete(ctx).await; } if content.is_empty() { diff --git a/src/commands/misc/sauce.rs b/src/commands/misc/sauce.rs new file mode 100644 index 0000000..116295f --- /dev/null +++ b/src/commands/misc/sauce.rs @@ -0,0 +1,60 @@ +use crate::messages::sauce::show_sauce_menu; +use crate::utils::get_previous_message_or_reply; + +use sauce_api::Sauce; + +use crate::utils::context_data::Store; +use serenity::client::Context; +use serenity::framework::standard::macros::command; +use serenity::framework::standard::CommandResult; +use serenity::model::channel::Message; + +#[command] +#[description("Searches for the source of a previously posted image or an image replied to.")] +#[usage("")] +async fn sauce(ctx: &Context, msg: &Message) -> CommandResult { + log::debug!("Got sauce command"); + let source_msg = get_previous_message_or_reply(ctx, msg).await?; + + if source_msg.is_none() { + log::debug!("No source message provided"); + msg.channel_id.say(ctx, "No source message found.").await?; + return Ok(()); + } + 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(); + + let mut embed_images = source_msg + .embeds + .into_iter() + .filter_map(|e| e.thumbnail) + .map(|t| t.url) + .collect::>(); + + attachment_urls.append(&mut embed_images); + log::trace!("Image urls {:?}", attachment_urls); + + if attachment_urls.is_empty() { + log::debug!("No images in source image"); + msg.channel_id.say(ctx, "Images in message found.").await?; + return Ok(()); + } + + log::debug!( + "Checking SauceNao for {} attachments", + attachment_urls.len() + ); + let data = ctx.data.read().await; + let store_data = data.get::().unwrap(); + let sources = store_data.sauce_nao.check_sauces(attachment_urls).await?; + log::trace!("Sources are {:?}", sources); + + log::debug!("Creating menu..."); + + show_sauce_menu(ctx, msg, sources).await?; + log::debug!("Menu created"); + + Ok(()) +} diff --git a/src/commands/mod.rs b/src/commands/mod.rs index 84ae7ed..6bafd62 100644 --- a/src/commands/mod.rs +++ b/src/commands/mod.rs @@ -4,8 +4,8 @@ pub use misc::MISC_GROUP; pub use music::MUSIC_GROUP; pub use settings::SETTINGS_GROUP; +mod common; pub(crate) mod minecraft; pub(crate) mod misc; pub(crate) mod music; pub(crate) mod settings; -mod common; diff --git a/src/messages/mod.rs b/src/messages/mod.rs index 3bad0b3..6eeb354 100644 --- a/src/messages/mod.rs +++ b/src/messages/mod.rs @@ -1 +1,2 @@ pub mod music; +pub mod sauce; diff --git a/src/messages/sauce.rs b/src/messages/sauce.rs new file mode 100644 index 0000000..678327b --- /dev/null +++ b/src/messages/sauce.rs @@ -0,0 +1,60 @@ +use crate::utils::error::BotResult; +use crate::utils::get_domain_for_url; +use sauce_api::SauceResult; +use serenity::builder::CreateMessage; +use serenity::{model::channel::Message, prelude::*}; +use serenity_utils::prelude::*; + +/// Builds a new sauce menu +pub async fn show_sauce_menu( + ctx: &Context, + msg: &Message, + sources: Vec, +) -> BotResult<()> { + let pages: Vec = sources.into_iter().map(create_sauce_page).collect(); + + let menu = if pages.len() == 1 { + Menu::new( + ctx, + msg, + &pages, + MenuOptions { + controls: vec![], + ..Default::default() + }, + ) + } else { + Menu::new(ctx, msg, &pages, MenuOptions::default()) + }; + menu.run().await?; + + Ok(()) +} + +/// Creates a single sauce page +fn create_sauce_page<'a>(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()); + + 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 + )); + } + } + message.embed(|e| { + e.title("Sources") + .description(description_lines.join("\n")) + .thumbnail(original) + .footer(|f| f.text("Powered by SauceNAO")) + }); + + message +} diff --git a/src/utils/context_data.rs b/src/utils/context_data.rs index 46bba6d..054e770 100644 --- a/src/utils/context_data.rs +++ b/src/utils/context_data.rs @@ -9,7 +9,9 @@ use crate::providers::music::queue::MusicQueue; use crate::providers::music::spotify::SpotifyApi; use crate::utils::messages::EventDrivenMessage; use database::Database; +use sauce_api::prelude::SauceNao; use serenity::client::Context; +use std::env; pub struct Store; @@ -17,16 +19,22 @@ pub struct StoreData { pub minecraft_data_api: minecraft_data_rs::api::Api, pub music_queues: HashMap>>, pub spotify_api: SpotifyApi, + pub sauce_nao: SauceNao, } impl StoreData { pub fn new() -> StoreData { + let mut sauce_nao = SauceNao::new(); + sauce_nao.set_api_key( + env::var("SAUCENAO_API_KEY").expect("No SAUCENAO_API_KEY key in environment."), + ); Self { minecraft_data_api: minecraft_data_rs::api::Api::new( minecraft_data_rs::api::versions::latest_stable().unwrap(), ), music_queues: HashMap::new(), spotify_api: SpotifyApi::new(), + sauce_nao, } } } diff --git a/src/utils/error.rs b/src/utils/error.rs index 1666014..33d7c8e 100644 --- a/src/utils/error.rs +++ b/src/utils/error.rs @@ -31,6 +31,9 @@ pub enum BotError { #[error("Detected CLI injection attempt")] CliInject, + #[error("Serenity Utils Error: {0}")] + SerenityUtils(#[from] serenity_utils::Error), + #[error("{0}")] Msg(String), } diff --git a/src/utils/logging.rs b/src/utils/logging.rs index 4dda7cf..c2a323b 100644 --- a/src/utils/logging.rs +++ b/src/utils/logging.rs @@ -59,6 +59,7 @@ pub fn init_logger() { .level_for("want", log::LevelFilter::Warn) .level_for("mio", log::LevelFilter::Warn) .level_for("songbird", log::LevelFilter::Warn) + .level_for("html5ever", log::LevelFilter::Warn) .chain(std::io::stdout()) .chain( fern::log_file(log_dir.join(PathBuf::from(format!( diff --git a/src/utils/mod.rs b/src/utils/mod.rs index 89ef7cb..36959b1 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,4 +1,8 @@ +use crate::utils::error::BotResult; use rand::Rng; +use regex::Regex; +use serenity::client::Context; +use serenity::model::channel::Message; use std::collections::VecDeque; pub(crate) mod context_data; @@ -16,3 +20,29 @@ pub fn shuffle_vec_deque(deque: &mut VecDeque) { deque.swap(i, rng.gen_range(0..i + 1)) } } + +/// Returns the message the given message is a reply to or the message sent before that +pub async fn get_previous_message_or_reply( + ctx: &Context, + msg: &Message, +) -> BotResult> { + let referenced = if let Some(reference) = &msg.referenced_message { + Some(*reference.clone()) + } else { + let messages = msg + .channel_id + .messages(ctx, |ret| ret.before(&msg.id).limit(1)) + .await?; + messages.first().cloned() + }; + + Ok(referenced) +} + +/// 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)?; + + captures.get(3).map(|c| c.as_str().to_string()) +}