tretrauit/swordfish
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: implement write_cards & support Qingque

Browse Source
This commit is contained in:
tretrauit 2024-01-07 01:05:47 +07:00
parent bd1d54e202
commit b16c0dadab
11 changed files with 676 additions and 46 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
Cargo.lock generated
View File

@ -2191,6 +2191,7 @@ dependencies = [
"mongodb",
"rusty-tesseract",
"serde",
"tokio",
"tracing",
"tracing-subscriber",
]

1
swordfish-common/Cargo.toml
View File

@ -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"] }

1
swordfish-common/src/constants.rs
View File

@ -1,2 +1,3 @@
pub const KATANA_ID: u64 = 646937666251915264;
pub const SOFA_ID: u64 = 853629533855809596;
pub const QINGQUE_ID: u64 = 772642704257187840;

101
swordfish-common/src/database/katana.rs
View File

@ -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<Collection<Card>> = OnceLock::new();
pub static KATANA: OnceLock<Collection<Card>> = OnceLock::new();
///
/// Initialize the "katana" collection in MongoDB
@ -39,8 +41,7 @@ pub async fn query_card(name: &str, series: &str) -> Option<Card> {
.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<Card>) -> Result<(), String> {
let mut new_cards: Vec<Card> = Vec::new();
let mut handles: Vec<task::JoinHandle<Result<Option<Card>, 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(())
}

2
swordfish-common/src/lib.rs
View File

@ -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()

4
swordfish-common/src/structs.rs
View File

@ -1,8 +1,8 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Deserialize, Serialize)]
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct Card {
pub wishlist: Option<i32>,
pub wishlist: Option<u32>,
pub name: String,
pub series: String,
pub print: i32,

60
swordfish-common/src/utils/katana.rs
View File

@ -1,3 +1,61 @@
fn parse_card() {
use crate::structs::Card;
use log::{error, trace};
pub fn parse_cards_from_qingque_atopwl(content: &String) -> Vec<Card> {
let mut cards: Vec<Card> = 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::<String>()
.trim()
.to_string();
trace!("Formatted wishlist number:{}", wl_string);
match wl_string.parse::<u32>() {
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
}

1
swordfish-common/src/utils/mod.rs Normal file
View File

@ -0,0 +1 @@
pub mod katana;

58
swordfish/src/config.rs
View File

@ -17,10 +17,37 @@ pub struct Tesseract {
pub backend: String,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Debug {
pub allowed_users: Vec<u64>,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct List {
pub enabled: bool,
pub servers: Vec<u64>,
pub channels: Vec<u64>,
}
#[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) {

35
swordfish/src/helper.rs
View File

@ -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<String>) {
match msg
.channel_id

454
swordfish/src/main.rs
View File

@ -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<Config> = 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<Message>,
new: Option<Message>,
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!("<t:{}:R>", 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 <subcommand> [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::<u64>() {
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 <channel ID> <message ID>`".to_string(),
None,
)
.await;
return Ok(());
}
};
let target_msg_id = match args.nth(0) {
Some(content) => match content.parse::<u64>() {
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 <channel ID> <message ID>`".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::<u64>() {
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 <channel ID> <message ID>`".to_string(),
None,
)
.await;
return Ok(());
}
};
let target_msg_id = match args.nth(0) {
Some(content) => match content.parse::<u64>() {
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 <channel ID> <message ID>`".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::<u64>() {
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(())
}