diff --git a/Cargo.lock b/Cargo.lock index d3cc3c5..76ee9b7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2191,6 +2191,7 @@ dependencies = [ "mongodb", "rusty-tesseract", "serde", + "tokio", "tracing", "tracing-subscriber", ] diff --git a/swordfish-common/Cargo.toml b/swordfish-common/Cargo.toml index 6aabb4d..7ff7012 100644 --- a/swordfish-common/Cargo.toml +++ b/swordfish-common/Cargo.toml @@ -10,6 +10,7 @@ leptess = "0.14.0" log = "0.4.20" rusty-tesseract = "1.1.9" serde = "1.0.195" +tokio = "1.35.1" tracing = "0.1.40" tracing-subscriber = { version = "0.3.18", features = ["env-filter"] } diff --git a/swordfish-common/src/constants.rs b/swordfish-common/src/constants.rs index cfbfb18..8ba1f34 100644 --- a/swordfish-common/src/constants.rs +++ b/swordfish-common/src/constants.rs @@ -1,2 +1,3 @@ pub const KATANA_ID: u64 = 646937666251915264; pub const SOFA_ID: u64 = 853629533855809596; +pub const QINGQUE_ID: u64 = 772642704257187840; diff --git a/swordfish-common/src/database/katana.rs b/swordfish-common/src/database/katana.rs index 37b9bff..05b9ce0 100644 --- a/swordfish-common/src/database/katana.rs +++ b/swordfish-common/src/database/katana.rs @@ -3,8 +3,10 @@ use crate::structs::Card; use mongodb::Collection; use std::sync::OnceLock; use std::time::{SystemTime, UNIX_EPOCH}; +use tokio::task; +use tracing::trace; -static KATANA: OnceLock> = OnceLock::new(); +pub static KATANA: OnceLock> = OnceLock::new(); /// /// Initialize the "katana" collection in MongoDB @@ -39,8 +41,7 @@ pub async fn query_card(name: &str, series: &str) -> Option { .unwrap() } -pub async fn write_card(mut card: Card) { - // todo!("Write card to database"); +pub async fn write_card(mut card: Card) -> Result<(), String> { let old_card = KATANA .get() .unwrap() @@ -59,7 +60,7 @@ pub async fn write_card(mut card: Card) { .expect("Time went backwards"); card.last_update_ts = current_time_ts.as_secs() as i64; if old_card.is_some() { - KATANA + match KATANA .get() .unwrap() .replace_one( @@ -71,8 +72,96 @@ pub async fn write_card(mut card: Card) { None, ) .await - .unwrap(); + { + Ok(_) => { + return Ok(()); + } + Err(e) => { + return Err(format!("Failed to update card: {}", e)); + } + } } else { - KATANA.get().unwrap().insert_one(card, None).await.unwrap(); + match KATANA.get().unwrap().insert_one(card, None).await { + Ok(_) => { + return Ok(()); + } + Err(e) => { + return Err(format!("Failed to insert card: {}", e)); + } + } } } + +pub async fn write_cards(cards: Vec) -> Result<(), String> { + let mut new_cards: Vec = Vec::new(); + let mut handles: Vec, String>>> = Vec::new(); + for mut card in cards { + trace!("Writing card: {:?}", card); + handles.push(task::spawn(async { + let old_card = KATANA + .get() + .unwrap() + .find_one( + mongodb::bson::doc! { + "name": card.name.clone(), + "series": card.series.clone() + }, + None, + ) + .await + .unwrap(); + let start = SystemTime::now(); + let current_time_ts = start + .duration_since(UNIX_EPOCH) + .expect("Time went backwards"); + card.last_update_ts = current_time_ts.as_secs() as i64; + if old_card.is_some() { + match KATANA + .get() + .unwrap() + .replace_one( + mongodb::bson::doc! { + "name": card.name.clone(), + "series": card.series.clone() + }, + card, + None, + ) + .await + { + Ok(_) => { + return Ok(None); + } + Err(e) => { + return Err(format!("Failed to update card: {}", e)); + } + } + } else { + return Ok(Some(card)); + }; + })); + } + for handle in handles { + match handle.await.unwrap() { + Ok(card) => { + if card.is_some() { + new_cards.push(card.unwrap()); + } + } + Err(e) => { + return Err(format!("Failed to update card: {}", e)); + } + } + } + if new_cards.len() > 0 { + match KATANA.get().unwrap().insert_many(new_cards, None).await { + Ok(_) => { + return Ok(()); + } + Err(e) => { + return Err(format!("Failed to insert card: {}", e)); + } + } + } + Ok(()) +} diff --git a/swordfish-common/src/lib.rs b/swordfish-common/src/lib.rs index cea07e8..2ffa8c0 100644 --- a/swordfish-common/src/lib.rs +++ b/swordfish-common/src/lib.rs @@ -1,4 +1,5 @@ #![feature(lazy_cell)] +#![feature(string_remove_matches)] pub use log; pub use tracing::{debug, error, info, trace, warn}; use tracing_subscriber::{self, fmt, EnvFilter}; @@ -6,6 +7,7 @@ pub mod constants; pub mod database; pub mod structs; pub mod tesseract; +pub mod utils; pub fn setup_logger(level: &str) -> Result<(), ()> { let formatter = fmt::format() diff --git a/swordfish-common/src/structs.rs b/swordfish-common/src/structs.rs index 9b48ccf..4ec6650 100644 --- a/swordfish-common/src/structs.rs +++ b/swordfish-common/src/structs.rs @@ -1,8 +1,8 @@ use serde::{Deserialize, Serialize}; -#[derive(Debug, Deserialize, Serialize)] +#[derive(Debug, Deserialize, Serialize, Clone)] pub struct Card { - pub wishlist: Option, + pub wishlist: Option, pub name: String, pub series: String, pub print: i32, diff --git a/swordfish-common/src/utils/katana.rs b/swordfish-common/src/utils/katana.rs index fee54aa..4352133 100644 --- a/swordfish-common/src/utils/katana.rs +++ b/swordfish-common/src/utils/katana.rs @@ -1,3 +1,61 @@ -fn parse_card() { - -} \ No newline at end of file +use crate::structs::Card; +use log::{error, trace}; + +pub fn parse_cards_from_qingque_atopwl(content: &String) -> Vec { + let mut cards: Vec = Vec::new(); + for line in content.split("\n") { + trace!("Parsing line: {}", line); + let mut line_split = line.split(" · "); + let wishlist = match line_split.nth(1) { + Some(wishlist_str) => { + let mut wl_string = wishlist_str.to_string(); + // Remove ` + wl_string.remove(0); + // Remove ❤ (Double because heart is 2 bytes) + wl_string.remove(0); + wl_string.remove(0); + // Remove last `` + wl_string.pop(); + // Remove "," in the number + wl_string.remove_matches(","); + // Remove whitespace + wl_string = wl_string + .split_whitespace() + .collect::() + .trim() + .to_string(); + trace!("Formatted wishlist number:{}", wl_string); + match wl_string.parse::() { + Ok(wishlist) => wishlist, + Err(_) => { + error!("Failed to parse wishlist number: {}", wishlist_str); + continue; + } + } + } + None => continue, + }; + let series = match line_split.next() { + Some(series) => series.to_string(), + None => continue, + }; + let name = match line_split.next() { + Some(name) => { + let mut name_string = name.to_string(); + name_string.remove_matches("**"); + name_string + } + None => continue, + }; + let card = Card { + wishlist: Some(wishlist), + name, + series, + print: 0, + last_update_ts: 0, + }; + trace!("Parsed card: {:?}", card); + cards.push(card); + } + cards +} diff --git a/swordfish-common/src/utils/mod.rs b/swordfish-common/src/utils/mod.rs new file mode 100644 index 0000000..2fd6688 --- /dev/null +++ b/swordfish-common/src/utils/mod.rs @@ -0,0 +1 @@ +pub mod katana; diff --git a/swordfish/src/config.rs b/swordfish/src/config.rs index c6b4039..f57069d 100644 --- a/swordfish/src/config.rs +++ b/swordfish/src/config.rs @@ -17,10 +17,37 @@ pub struct Tesseract { pub backend: String, } +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct Debug { + pub allowed_users: Vec, +} + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct List { + pub enabled: bool, + pub servers: Vec, + pub channels: Vec, +} + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct DropAnalyzer { + pub enabled: bool, + pub blacklist: List, + pub whitelist: List, +} + +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct Features { + pub katana_drop_analysis: DropAnalyzer, + pub sofa_drop_analysis: DropAnalyzer, +} + #[derive(Serialize, Deserialize, Debug, Clone)] pub struct Config { pub log: Log, pub tesseract: Tesseract, + pub debug: Debug, + pub features: Features, } impl Config { @@ -36,6 +63,37 @@ impl Config { tesseract: Tesseract { backend: "libtesseract".to_string(), }, + debug: Debug { + allowed_users: vec![], + }, + features: Features { + katana_drop_analysis: DropAnalyzer { + enabled: false, + blacklist: List { + enabled: false, + servers: vec![], + channels: vec![], + }, + whitelist: List { + enabled: false, + servers: vec![], + channels: vec![], + }, + }, + sofa_drop_analysis: DropAnalyzer { + enabled: false, + blacklist: List { + enabled: false, + servers: vec![], + channels: vec![], + }, + whitelist: List { + enabled: false, + servers: vec![], + channels: vec![], + }, + }, + }, } } pub fn save(&self, path: &str) { diff --git a/swordfish/src/helper.rs b/swordfish/src/helper.rs index fda92e8..1bc7ddc 100644 --- a/swordfish/src/helper.rs +++ b/swordfish/src/helper.rs @@ -1,9 +1,44 @@ +use crate::config::List; use crate::template::message; use serenity::builder::CreateMessage; use serenity::client::Context; use serenity::model::channel::Message; use swordfish_common::error; +pub fn message_in_blacklist(msg: &Message, blacklist: &List) -> bool { + if !blacklist.enabled { + return false; + } + let guild_id = match msg.guild_id { + Some(id) => id, + None => return false, + }; + if blacklist.servers.contains(&guild_id.get()) { + return true; + } + if blacklist.channels.contains(&msg.channel_id.get()) { + return true; + } + return false; +} + +pub fn message_in_whitelist(msg: &Message, whitelist: &List) -> bool { + if !whitelist.enabled { + return true; + } + let guild_id = match msg.guild_id { + Some(id) => id, + None => return false, + }; + if whitelist.servers.contains(&guild_id.get()) { + return true; + } + if whitelist.channels.contains(&msg.channel_id.get()) { + return true; + } + return false; +} + pub async fn error_message(ctx: &Context, msg: &Message, content: String, title: Option) { match msg .channel_id diff --git a/swordfish/src/main.rs b/swordfish/src/main.rs index eda6912..58b67c7 100644 --- a/swordfish/src/main.rs +++ b/swordfish/src/main.rs @@ -1,5 +1,6 @@ #![feature(lazy_cell)] use dotenvy::dotenv; +use serenity::all::MessageUpdateEvent; use serenity::async_trait; use serenity::framework::standard::macros::{command, group}; use serenity::framework::standard::{CommandResult, Configuration, StandardFramework}; @@ -25,7 +26,7 @@ const GITHUB_URL: &str = "https://github.com/teppyboy/swordfish"; static CONFIG: OnceLock = OnceLock::new(); #[group] -#[commands(ping, kdropanalyze, info)] +#[commands(ping, debug, info)] struct General; struct Handler; #[async_trait] @@ -45,31 +46,145 @@ impl EventHandler for Handler { } } } + async fn message_update( + &self, + ctx: Context, + old_if_available: Option, + new: Option, + event: MessageUpdateEvent, + ) { + let author = match event.author { + Some(ref v) => v, + None => { + return; + } + }; + if author.id == ctx.cache.current_user().id { + return; + } + let content = match event.content { + Some(ref v) => v, + None => { + return; + } + }; + trace!("Message update: {}, sender: {}", content, author.id); + if author.id.get() == constants::QINGQUE_ID { + parse_qingque(&ctx, event).await.unwrap(); + } + } } -async fn parse_katana(_ctx: &Context, msg: &Message) -> Result<(), String> { +async fn parse_qingque(ctx: &Context, event: MessageUpdateEvent) -> Result<(), String> { + if event.embeds.is_none() || event.embeds.clone().unwrap().len() == 0 { + return Ok(()); + } + let embed = &event.embeds.unwrap()[0]; + let embed_title = match embed.title { + Some(ref title) => title, + None => { + return Ok(()); + } + }; + match embed_title.as_str() { + "Top Wishlist" => { + let cards = utils::katana::parse_cards_from_qingque_atopwl( + &embed.description.as_ref().unwrap(), + ); + trace!("Begin importing cards"); + match database::katana::write_cards(cards).await { + Ok(_) => { + trace!("Imported successully"); + } + Err(why) => { + error!("Failed to import card: {:?}", why); + } + } + } + _ => { + return Ok(()); + } + } + Ok(()) +} + +async fn parse_katana(ctx: &Context, msg: &Message) -> Result<(), String> { if msg.content.contains("is dropping 3 cards!") || msg .content .contains("I'm dropping 3 cards since this server is currently active!") { - // trace!("Card drop detected, executing drop analyzer..."); - // match katana::analyze_drop_message(&LEPTESS_ARC, msg).await { - // Ok(_) => { - // // msg.reply(ctx, "Drop analysis complete").await?; - // } - // Err(why) => { - // trace!("Failed to analyze drop: `{:?}`", why); - // // helper::error_message(ctx, msg, format!("Failed to analyze drop: `{:?}`", why)).await; - // } - // }; + let config = CONFIG.get().unwrap(); + if !config.features.katana_drop_analysis.enabled { + return Ok(()); + } + if helper::message_in_blacklist(msg, &config.features.katana_drop_analysis.blacklist) { + return Ok(()); + } + if !helper::message_in_whitelist(msg, &config.features.katana_drop_analysis.whitelist) { + return Ok(()); + } + let start = Instant::now(); + match katana::analyze_drop_message(msg).await { + Ok(cards) => { + let duration = start.elapsed(); + let mut reply_str = String::new(); + for card in cards { + // reply_str.push_str(&format!("{:?}\n", card)); + let wishlist_str: String = match card.wishlist { + Some(wishlist) => { + let mut out_str = wishlist.to_string(); + while out_str.len() < 5 { + out_str.push(' '); + } + out_str + } + None => "None ".to_string(), + }; + let last_update_ts_str = match card.last_update_ts { + 0 => "`Never`".to_string(), + ts => { + format!("", ts.to_string()) + } + }; + reply_str.push_str( + format!( + ":heart: `{}` • `{}` • **{}** • {} • {}\n", + wishlist_str, card.print, card.name, card.series, last_update_ts_str + ) + .as_str(), + ) + } + reply_str.push_str(&format!("Time taken (to analyze): `{:?}`", duration)); + match msg.reply(ctx, reply_str).await { + Ok(_) => {} + Err(why) => { + error!("Failed to reply to message: {:?}", why); + } + }; + } + Err(why) => { + helper::error_message( + ctx, + msg, + format!("Failed to analyze drop: `{:?}`", why), + None, + ) + .await; + } + }; } Ok(()) } #[tokio::main] async fn main() { - dotenv().unwrap(); + match dotenv() { + Ok(_) => {} + Err(why) => { + eprintln!("Failed to load .env: {:?}", why); + } + } let token = env::var("DISCORD_TOKEN").expect("Token not found"); let config: Config; if Path::new("./config.toml").exists() { @@ -114,9 +229,298 @@ async fn ping(ctx: &Context, msg: &Message) -> CommandResult { } #[command] -async fn kdropanalyze(ctx: &Context, msg: &Message) -> CommandResult { +async fn debug(ctx: &Context, msg: &Message) -> CommandResult { + let config = CONFIG.get().unwrap(); + if !config.debug.allowed_users.contains(&msg.author.id.get()) { + return Ok(()); + } let mut args = msg.content.split(" "); - let target_channel_id = match args.nth(1) { + let subcommand = match args.nth(1) { + Some(content) => content, + None => { + helper::error_message( + ctx, + msg, + "Usage: `debug [args...]`".to_string(), + None, + ) + .await; + return Ok(()); + } + }; + match subcommand { + "kdropanalyze" => dbg_kdropanalyze(ctx, msg).await?, + "kda" => dbg_kdropanalyze(ctx, msg).await?, + "embed" => dbg_embed(ctx, msg).await?, + "parse-qingque-atopwl" => dbg_parse_qingque_atopwl(ctx, msg).await?, + _ => { + helper::error_message( + ctx, + msg, + format!("Unknown subcommand: `{}`", subcommand), + None, + ) + .await; + return Ok(()); + } + } + Ok(()) +} + +#[command] +async fn info(ctx: &Context, msg: &Message) -> CommandResult { + let reply_str = format!( + "Swordfish v{} - {}\n\ + Log level: `{}`\n\ + Build type: `{}`\n\n\ + Like my work? Consider supporting me at my [Ko-fi](https://ko-fi.com/tretrauit) or [Patreon](https://patreon.com/tretrauit)!\n\n\ + *Debug information*\n\ + Tesseract backend: `{}`\n\ + ", + env!("CARGO_PKG_VERSION"), + GITHUB_URL, + CONFIG.get().unwrap().log.level.clone().as_str(), + env!("BUILD_PROFILE"), + CONFIG.get().unwrap().tesseract.backend.clone().as_str(), + ); + helper::info_message(ctx, msg, reply_str, Some("Information".to_string())).await; + Ok(()) +} + +async fn dbg_parse_qingque_atopwl(ctx: &Context, msg: &Message) -> CommandResult { + let mut args = msg.content.split(" "); + let target_channel_id = match args.nth(2) { + Some(content) => match content.parse::() { + Ok(id) => id, + Err(why) => { + helper::error_message( + ctx, + msg, + format!("Failed to parse channel ID: `{:?}`", why), + None, + ) + .await; + return Ok(()); + } + }, + None => { + helper::error_message( + ctx, + msg, + "Usage: `parse-qingque-atopwl `".to_string(), + None, + ) + .await; + return Ok(()); + } + }; + let target_msg_id = match args.nth(0) { + Some(content) => match content.parse::() { + Ok(id) => id, + Err(why) => { + helper::error_message( + ctx, + msg, + format!("Failed to parse message ID: `{:?}`", why), + None, + ) + .await; + return Ok(()); + } + }, + None => { + helper::error_message( + ctx, + msg, + "Usage: `parse-qingque-atopwl `".to_string(), + None, + ) + .await; + return Ok(()); + } + }; + let target_msg = match ctx + .http() + .get_message( + ChannelId::new(target_channel_id), + MessageId::new(target_msg_id), + ) + .await + { + Ok(msg) => msg, + Err(why) => { + helper::error_message( + ctx, + msg, + format!("Failed to get message: `{:?}`", why), + None, + ) + .await; + return Ok(()); + } + }; + if target_msg.embeds.len() == 0 { + helper::error_message( + ctx, + msg, + "Message does not contain any embeds".to_string(), + None, + ) + .await; + return Ok(()); + } + let embed = &target_msg.embeds[0]; + let embed_description = match embed.description { + Some(ref description) => description, + None => { + helper::error_message( + ctx, + msg, + "Embed does not contain a description".to_string(), + None, + ) + .await; + return Ok(()); + } + }; + let cards = utils::katana::parse_cards_from_qingque_atopwl(embed_description); + helper::info_message( + ctx, + msg, + format!("Parsed cards: ```\n{:?}\n```", cards), + None, + ) + .await; + Ok(()) +} + +async fn dbg_embed(ctx: &Context, msg: &Message) -> CommandResult { + let mut args = msg.content.split(" "); + let target_channel_id = match args.nth(2) { + Some(content) => match content.parse::() { + Ok(id) => id, + Err(why) => { + helper::error_message( + ctx, + msg, + format!("Failed to parse channel ID: `{:?}`", why), + None, + ) + .await; + return Ok(()); + } + }, + None => { + helper::error_message( + ctx, + msg, + "Usage: `embed `".to_string(), + None, + ) + .await; + return Ok(()); + } + }; + let target_msg_id = match args.nth(0) { + Some(content) => match content.parse::() { + Ok(id) => id, + Err(why) => { + helper::error_message( + ctx, + msg, + format!("Failed to parse message ID: `{:?}`", why), + None, + ) + .await; + return Ok(()); + } + }, + None => { + helper::error_message( + ctx, + msg, + "Usage: `embed `".to_string(), + None, + ) + .await; + return Ok(()); + } + }; + let target_msg = match ctx + .http() + .get_message( + ChannelId::new(target_channel_id), + MessageId::new(target_msg_id), + ) + .await + { + Ok(msg) => msg, + Err(why) => { + helper::error_message( + ctx, + msg, + format!("Failed to get message: `{:?}`", why), + None, + ) + .await; + return Ok(()); + } + }; + if target_msg.embeds.len() == 0 { + helper::error_message( + ctx, + msg, + "Message does not contain any embeds".to_string(), + None, + ) + .await; + return Ok(()); + } + let embed = &target_msg.embeds[0]; + let embed_title = match embed.title { + Some(ref title) => title, + None => { + helper::error_message(ctx, msg, "Embed does not contain a title".to_string(), None) + .await; + return Ok(()); + } + }; + let embed_description = match embed.description { + Some(ref description) => description, + None => { + helper::error_message( + ctx, + msg, + "Embed does not contain a description".to_string(), + None, + ) + .await; + return Ok(()); + } + }; + helper::info_message( + ctx, + msg, + format!( + "Title: \n\ + ```\ + {}\n\ + ```\n\ + Description: \n\ + ```\n\ + {}\n\ + ```", + embed_title, embed_description + ), + Some("Embed information".to_string()), + ) + .await; + Ok(()) +} + +async fn dbg_kdropanalyze(ctx: &Context, msg: &Message) -> CommandResult { + let mut args = msg.content.split(" "); + let target_channel_id = match args.nth(2) { Some(content) => match content.parse::() { Ok(id) => id, Err(why) => { @@ -232,23 +636,3 @@ async fn kdropanalyze(ctx: &Context, msg: &Message) -> CommandResult { }; Ok(()) } - -#[command] -async fn info(ctx: &Context, msg: &Message) -> CommandResult { - let reply_str = format!( - "Swordfish v{} - {}\n\ - Log level: `{}`\n\ - Build type: `{}`\n\n\ - Like my work? Consider supporting me at my [Ko-fi](https://ko-fi.com/tretrauit) or [Patreon](https://patreon.com/tretrauit)!\n\n\ - *Debug information*\n\ - Tesseract backend: `{}`\n\ - ", - env!("CARGO_PKG_VERSION"), - GITHUB_URL, - CONFIG.get().unwrap().log.level.clone().as_str(), - env!("BUILD_PROFILE"), - CONFIG.get().unwrap().tesseract.backend.clone().as_str(), - ); - helper::info_message(ctx, msg, reply_str, Some("Information".to_string())).await; - Ok(()) -}