diff --git a/Cargo.lock b/Cargo.lock index e6acd5c..31bcebe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1126,6 +1126,7 @@ dependencies = [ "migration", "models", "rand", + "regex", "reqwest", "sea-orm", "shadow-rs", diff --git a/Cargo.toml b/Cargo.toml index 0a3b793..4ae6250 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,7 @@ futures = "0.3.29" strfmt = "^0.2.4" reqwest= "^0.11" rand="^0.8.5" +regex = "^1.10.2" #lazy_static="^1.4.0" [dependencies.clap] diff --git a/src/commands.rs b/src/commands.rs index b66e1ac..34adfc1 100644 --- a/src/commands.rs +++ b/src/commands.rs @@ -1,7 +1,5 @@ use std::collections::HashMap; -use rand::{rngs::OsRng, Rng}; -use sea_orm::DbErr; use strfmt::Format; use teloxide::{ payloads::SendMessageSetters, @@ -13,13 +11,13 @@ use teloxide::{ }; use crate::{ - config::Args, db_controller::Controller, messages::{ - BOT_ABOUT, BOT_TEXT_HANGED, BOT_TEXT_HANGED_SELF, BOT_TEXT_IS_CHANNEL, BOT_TEXT_NO_TARGET, - BOT_TEXT_TOP_GLOBAL, BOT_TEXT_TOP_GROUP, BOT_TEXT_TOP_NONE, BOT_TEXT_TOP_TEMPLATE, - BOT_TEXT_TOP_TITLE, BOT_TEXT_HANG_BOT, BOT_TEXT_HANG_ANONYMOUS, BOT_TEXT_HANG_CHANNEL, + BOT_ABOUT, BOT_TEXT_HANG_ANONYMOUS, BOT_TEXT_HANG_BOT, BOT_TEXT_HANG_CHANNEL, + BOT_TEXT_IS_CHANNEL, BOT_TEXT_NO_TARGET, BOT_TEXT_TOP_GLOBAL, BOT_TEXT_TOP_GROUP, + BOT_TEXT_TOP_NONE, BOT_TEXT_TOP_TEMPLATE, BOT_TEXT_TOP_TITLE, }, + utils::hangit_text, }; #[derive(BotCommands, PartialEq, Debug, Clone)] @@ -44,146 +42,96 @@ impl Default for Commands { } } -#[derive(Debug, Clone)] -pub struct CommandHandler { - pub controller: Controller, +async fn send_text_reply(bot: &Bot, message: &Message, text: String) -> Result<(), RequestError> { + bot.send_message(message.chat.id, text) + .reply_to_message_id(message.id) + .parse_mode(ParseMode::MarkdownV2) + .await?; + Ok(()) } -impl CommandHandler { - pub async fn new(config: &Args) -> Result { - Ok(Self { - controller: Controller::new(config.database_uri.to_owned()).await?, - }) - } +pub async fn help_handler(bot: &Bot, message: &Message) -> Result<(), RequestError> { + send_text_reply(bot, message, Commands::descriptions().to_string()).await +} - pub async fn init(&self) -> Result<(), DbErr> { - self.controller.migrate().await - } +pub async fn about_handler(bot: &Bot, message: &Message) -> Result<(), RequestError> { + send_text_reply(bot, message, BOT_ABOUT.to_string()).await +} - async fn send_text_reply( - &self, - bot: &Bot, - message: &Message, - text: String, - ) -> Result { - bot.send_message(message.chat.id, text) - .reply_to_message_id(message.id) - .parse_mode(ParseMode::MarkdownV2) - .await - } - - pub async fn help_handler( - &self, - bot: &Bot, - message: &Message, - ) -> Result { - self.send_text_reply(bot, message, Commands::descriptions().to_string()) - .await - } - - pub async fn about_handler( - &self, - bot: &Bot, - message: &Message, - ) -> Result { - self.send_text_reply(bot, message, BOT_ABOUT.to_string()) - .await - } - - pub async fn hangit_handler( - &self, - bot: &Bot, - message: &Message, - ) -> Result { - let reply = match message.reply_to_message() { - Some(reply) => reply, - None => { - return self - .send_text_reply(bot, message, BOT_TEXT_NO_TARGET.to_string()) - .await - } - }; - - match reply.from() { - Some(user) => { - if user.is_bot { - return self.send_text_reply(bot, reply, BOT_TEXT_HANG_BOT.to_string()).await - } - - if user.is_anonymous() { - return self.send_text_reply(bot, reply, BOT_TEXT_HANG_ANONYMOUS.to_string()).await - } - - if user.is_channel() { - return self.send_text_reply(bot, reply, BOT_TEXT_HANG_CHANNEL.to_string()).await - } - - let is_self = match message.from() { - Some(f) => f.first_name == user.first_name, - None => false, - }; - - let mut vars = HashMap::new(); - - let index = if is_self { - OsRng.gen::() % BOT_TEXT_HANGED_SELF.len() - } else { - OsRng.gen::() % BOT_TEXT_HANGED.len() - }; - - let text = if is_self { - BOT_TEXT_HANGED_SELF[index] - } else { - BOT_TEXT_HANGED[index] - }; - - let name = escape(user.first_name.as_str()); - vars.insert("name".to_string(), name.as_str()); - - let _ = self - .controller - .hangit(&user.full_name(), message.chat.id) - .await; - self.send_text_reply(bot, reply, text.format(&vars).unwrap()) - .await - } - None => { - self.send_text_reply(bot, message, BOT_TEXT_IS_CHANNEL.to_string()) - .await - } +pub async fn hangit_handler( + db: &Controller, + bot: &Bot, + message: &Message, +) -> Result<(), RequestError> { + let reply = match message.reply_to_message() { + Some(reply) => reply, + None => { + return send_text_reply(bot, message, BOT_TEXT_NO_TARGET.to_string()).await; } - } + }; - pub async fn top_handler(&self, bot: &Bot, message: &Message) -> Result { - let chat = &message.chat; - let scope = match chat.is_group() || chat.is_supergroup() { - true => BOT_TEXT_TOP_GROUP, - false => BOT_TEXT_TOP_GLOBAL, - }; - - let mut index = 1; - let mut text = format!("{}\\-{}\n\n", BOT_TEXT_TOP_TITLE, scope); - let results = match self.controller.top(chat).await { - Some(result) => result, - None => { - return self - .send_text_reply(bot, message, BOT_TEXT_TOP_NONE.to_string()) - .await + match reply.from() { + Some(user) => { + if user.is_bot { + return send_text_reply(bot, message, BOT_TEXT_HANG_BOT.to_string()).await; } - }; - for result in results { - let mut vars: HashMap = HashMap::new(); + if user.is_anonymous() { + return send_text_reply(bot, message, BOT_TEXT_HANG_ANONYMOUS.to_string()).await; + } - vars.insert("name".to_string(), escape(result.name.as_str())); - vars.insert("count".to_string(), result.counts.to_string()); + if user.is_channel() { + return send_text_reply(bot, message, BOT_TEXT_HANG_CHANNEL.to_string()).await; + } - let record = BOT_TEXT_TOP_TEMPLATE.format(&vars).unwrap(); + let is_self = match message.from() { + Some(f) => f.first_name == user.first_name, + None => false, + }; - text = format!("{}{}\\. {}\n", text, index, record); - index += 1; + let _ = db.hangit(&user.full_name(), message.chat.id).await; + send_text_reply( + bot, + reply, + hangit_text(user.first_name.to_string(), is_self, true), + ) + .await } - - self.send_text_reply(bot, message, text).await + None => send_text_reply(bot, message, BOT_TEXT_IS_CHANNEL.to_string()).await, } } + +pub async fn top_handler( + db: &Controller, + bot: &Bot, + message: &Message, +) -> Result<(), RequestError> { + let chat = &message.chat; + let scope = match chat.is_group() || chat.is_supergroup() { + true => BOT_TEXT_TOP_GROUP, + false => BOT_TEXT_TOP_GLOBAL, + }; + + let mut index = 1; + let mut text = format!("{}\\-{}\n\n", BOT_TEXT_TOP_TITLE, scope); + let results = match db.top(chat).await { + Some(result) => result, + None => { + return send_text_reply(bot, message, BOT_TEXT_TOP_NONE.to_string()).await; + } + }; + + for result in results { + let mut vars: HashMap = HashMap::new(); + + vars.insert("name".to_string(), escape(result.name.as_str())); + vars.insert("count".to_string(), result.counts.to_string()); + + let record = BOT_TEXT_TOP_TEMPLATE.format(&vars).unwrap(); + + text = format!("{}{}\\. {}\n", text, index, record); + index += 1; + } + + send_text_reply(bot, message, text).await +} diff --git a/src/db_controller.rs b/src/db_controller.rs index 72924a7..6391fa7 100644 --- a/src/db_controller.rs +++ b/src/db_controller.rs @@ -8,6 +8,8 @@ use sea_orm::{ use teloxide::types::{Chat, ChatId}; use wd_log::{log_debug_ln, log_error_ln, log_info_ln}; +const LIMIT: u64 = 10; + #[derive(Debug, FromQueryResult)] pub struct TopData { pub name: String, @@ -32,7 +34,7 @@ impl Controller { if let Err(err) = Migrator::install(&self.db).await { return Err(err); } - + if let Err(err) = Migrator::up(&self.db, None).await { Err(err) } else { @@ -70,9 +72,23 @@ impl Controller { transcation.commit().await } + pub async fn update_group(&self, name: String, group_id: ChatId) -> Result<(), DbErr> { + log_debug_ln!("name={:?}", name); + + let transcation = self.db.begin().await?; + match Stats::find().filter(StatsColumn::GroupId.eq(0)).filter(StatsColumn::Name.eq(name)).one(&transcation).await? { + Some(one) => { + let mut one: StatsActiveModel = one.into(); + one.group_id = Set(group_id.0); + one.save(&transcation).await?; + }, + None => {}, + } + transcation.commit().await + } + /// stats pub async fn top(&self, chat: &Chat) -> Option> { - const LIMIT: u64 = 10; let transcation = match self.db.begin().await { Ok(t) => t, Err(error) => { @@ -89,9 +105,10 @@ impl Controller { .order_by_desc(StatsColumn::Counts.sum()) .limit(LIMIT); - let query = match chat.is_group() || chat.is_supergroup() { - true => query.filter(StatsColumn::GroupId.eq(chat.id.0)), - false => query, + let query = if chat.is_group() || chat.is_supergroup() { + query.filter(StatsColumn::GroupId.eq(chat.id.0)) + } else { + query }; log_debug_ln!( @@ -109,4 +126,26 @@ impl Controller { } } } + + pub async fn find_by_name(&self, name: &String) -> Option> { + let transcation = match self.db.begin().await { + Ok(t) => t, + Err(error) => { + log_error_ln!("{}", error); + return None; + } + }; + + let result = Stats::find() + .filter(StatsColumn::Name.contains(name)) + .limit(LIMIT).all(&transcation).await; + + match result { + Ok(result) => Some(result), + Err(error) => { + log_error_ln!("{}", error); + None + } + } + } } diff --git a/src/inline_query.rs b/src/inline_query.rs new file mode 100644 index 0000000..543c228 --- /dev/null +++ b/src/inline_query.rs @@ -0,0 +1,62 @@ +use teloxide::{ + prelude::Bot, + requests::{Request, Requester}, + types::{ + ChatId, ChosenInlineResult, InlineQuery, InlineQueryResult, InlineQueryResultArticle, + InputMessageContent, InputMessageContentText, + }, + RequestError, +}; +use wd_log::{log_error_ln, log_debug_ln}; + +use crate::{ + db_controller::Controller, + messages::BOT_TEXT_INLINE_HANG, + utils::{hangit_text, IS_SELF, NEED_ESCAPE}, +}; + +pub async fn inline_menu(db: &Controller, bot: &Bot, q: InlineQuery) -> Result<(), RequestError> { + let name = q.query; + + let mut results = match db.find_by_name(&name).await { + Some(list) => list + .iter() + .map(|n| { + InlineQueryResult::Article(InlineQueryResultArticle::new( + n.name.clone(), + format!("{} {}", BOT_TEXT_INLINE_HANG, n.name), + InputMessageContent::Text(InputMessageContentText::new(hangit_text( + n.name.clone(), + q.from.first_name == n.name, + !NEED_ESCAPE, + ))), + )) + }) + .collect::>(), + + None => vec![], + }; + + results.push(InlineQueryResult::Article(InlineQueryResultArticle::new( + name.clone(), + format!("{} {}", BOT_TEXT_INLINE_HANG, name), + InputMessageContent::Text(InputMessageContentText::new(hangit_text( + name, + !IS_SELF, + !NEED_ESCAPE, + ))), + ))); + + bot.answer_inline_query(&q.id, results).send().await?; + Ok(()) +} + +pub async fn inline_anwser(db: &Controller, a: ChosenInlineResult) -> Result<(), RequestError> { + log_debug_ln!("{:#?}", a); + + if let Err(err) = db.hangit(&a.result_id, ChatId(0)).await { + log_error_ln!("{:?}", err); + } + + Ok(()) +} diff --git a/src/main.rs b/src/main.rs index 1e92faa..92a0003 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,17 +1,23 @@ mod commands; mod config; mod db_controller; +mod inline_query; mod messages; +mod utils; use clap::Parser; -use commands::{CommandHandler, Commands}; +use commands::{about_handler, hangit_handler, help_handler, top_handler, Commands}; use config::Args; +use db_controller::Controller; +use inline_query::{inline_anwser, inline_menu}; use teloxide::{ prelude::*, requests::{Request, Requester}, + types::{Me, Update}, utils::command::BotCommands, }; +use utils::message_handler; use wd_log::{ log_debug_ln, log_error_ln, log_info_ln, log_panic, set_level, set_prefix, DEBUG, INFO, }; @@ -28,43 +34,77 @@ async fn main() { } else { set_level(INFO); } - let command_handler = match CommandHandler::new(&args).await { - Err(err) => log_panic!("{}", err), - Ok(c) => c, + + let db_controller = match db_controller::Controller::new(args.database_uri.to_owned()).await { + Ok(db) => db, + Err(err) => { + log_panic!("{:?}", err); + } }; - command_handler.init().await.unwrap(); + if let Err(err) = db_controller.migrate().await { + log_panic!("{:?}", err); + } let bot = Bot::new(args.tgbot_token.to_owned()) .set_api_url(reqwest::Url::parse(&args.api_url.as_str()).unwrap()); - get_me(&bot).await; + let me = get_me(&bot).await; register_commands(&bot).await; - Commands::repl(bot, move |bot: Bot, message: Message, cmd: Commands| { - let command_handler = command_handler.clone(); + let handler = dptree::entry() + .branch( + Update::filter_message() + .branch(dptree::entry().filter_command::().endpoint( + |db: Controller, bot: Bot, message: Message, cmd: Commands| async move { + let r = match cmd { + Commands::Help => help_handler(&bot, &message).await, + Commands::About => about_handler(&bot, &message).await, + Commands::Top => top_handler(&db, &bot, &message).await, + Commands::HangIt => hangit_handler(&db, &bot, &message).await, + }; - async move { - let r = match cmd { - Commands::Help => command_handler.help_handler(&bot, &message).await, - Commands::About => command_handler.about_handler(&bot, &message).await, - Commands::Top => command_handler.top_handler(&bot, &message).await, - Commands::HangIt => command_handler.hangit_handler(&bot, &message).await, - }; + match r { + Ok(_) => Ok(()), + Err(err) => { + log_error_ln!("{:?}", err); + Err(err) + } + } + }, + )) + .branch( + dptree::filter(|msg: Message| msg.chat.is_group() || msg.chat.is_supergroup()) + .endpoint(|db: Controller, msg: Message, me: Me| async move { + let r = message_handler(&db, msg, &me).await; + match r { + Ok(_) => Ok(()), + Err(err) => { + log_error_ln!("{:?}", err); + Err(err) + } + } + }), + ), + ) + .branch( + Update::filter_inline_query().endpoint( + |db: Controller, bot: Bot, q: InlineQuery| async move { + inline_menu(&db, &bot, q).await + }, + ), + ) + .branch(Update::filter_chosen_inline_result().endpoint( + |db: Controller, a: ChosenInlineResult| async move { inline_anwser(&db, a).await }, + )); - match r { - Ok(_r) => { - log_debug_ln!("will send: {:?}", _r.text()); - Ok(()) - } - Err(err) => { - log_error_ln!("{:?}", err); - Err(err) - } - } - } - }) - .await; + Dispatcher::builder(bot, handler) + .dependencies(dptree::deps![db_controller, me]) + .default_handler(|upd| async move { log_debug_ln!("unhandled update: {:?}", upd) }) + .enable_ctrlc_handler() + .build() + .dispatch() + .await; } async fn register_commands(bot: &Bot) { @@ -75,13 +115,16 @@ async fn register_commands(bot: &Bot) { } } -async fn get_me(bot: &Bot) { +async fn get_me(bot: &Bot) -> Me { match bot.get_me().send().await { - Ok(result) => log_info_ln!( - "connect succeed: id={}, botname=\"{}\"", - result.id, - result.username() - ), + Ok(result) => { + log_info_ln!( + "connect succeed: id={}, botname=\"{}\"", + result.id, + result.username() + ); + result + } Err(error) => log_panic!("{}", error), } } diff --git a/src/messages.rs b/src/messages.rs index 25bebb4..8a4d25c 100644 --- a/src/messages.rs +++ b/src/messages.rs @@ -14,7 +14,8 @@ const BOT_TEXT_HANGED_2: &'static str = "因为 {name} 太过逆天，我们把 const BOT_TEXT_HANGED_3: &'static str = "{name} 吊在了路灯上，TA 兴风作浪的时代结束了……"; const BOT_TEXT_HANGED_4: &'static str = "吊在路灯上的 {name} 正在接受大家的鄙视……"; const BOT_TEXT_HANGED_5: &'static str = "对 {name} 来说，绳命来得快去得也快，只有路灯是永恒的……"; -const BOT_TEXT_HANGED_6: &'static str = "被套上麻袋的 {name} 在经历了一顿胖揍之后，最后还是成了路灯的挂件……"; +const BOT_TEXT_HANGED_6: &'static str = + "被套上麻袋的 {name} 在经历了一顿胖揍之后，最后还是成了路灯的挂件……"; pub const BOT_TEXT_HANGED: [&str; 6] = [ BOT_TEXT_HANGED_1, @@ -26,7 +27,8 @@ pub const BOT_TEXT_HANGED: [&str; 6] = [ ]; const BOT_TEXT_HANGED_SELF_1: &'static str = "{name} 承受不了自己所做的一切，选择了自行了断……"; -const BOT_TEXT_HANGED_SELF_2: &'static str = "对于 {name} 来说，把自己吊在路灯上可能是最好的选择了……"; +const BOT_TEXT_HANGED_SELF_2: &'static str = + "对于 {name} 来说，把自己吊在路灯上可能是最好的选择了……"; const BOT_TEXT_HANGED_SELF_3: &'static str = "{name} 最终还是选择了逃避……"; pub const BOT_TEXT_HANGED_SELF: [&str; 3] = [ @@ -37,4 +39,5 @@ pub const BOT_TEXT_HANGED_SELF: [&str; 3] = [ pub const BOT_TEXT_HANG_BOT: &'static str = "机器人是无法被吊死的……"; pub const BOT_TEXT_HANG_CHANNEL: &'static str = "这是个频道……"; -pub const BOT_TEXT_HANG_ANONYMOUS: &'static str = "这是个幽灵……"; \ No newline at end of file +pub const BOT_TEXT_HANG_ANONYMOUS: &'static str = "这是个幽灵……"; +pub const BOT_TEXT_INLINE_HANG: &'static str = "吊死"; diff --git a/src/utils.rs b/src/utils.rs new file mode 100644 index 0000000..d00258f --- /dev/null +++ b/src/utils.rs @@ -0,0 +1,84 @@ +use std::collections::HashMap; + +use rand::{rngs::OsRng, Rng}; +use regex::Regex; +use strfmt::Format; +use teloxide::{ + types::{Me, Message}, + utils::markdown::escape, + RequestError, +}; +use wd_log::{log_debug_ln, log_error_ln}; + +use crate::{ + db_controller::Controller, + messages::{BOT_TEXT_HANGED, BOT_TEXT_HANGED_SELF}, +}; + +pub const IS_SELF: bool = true; +pub const NEED_ESCAPE: bool = true; + +pub fn hangit_text(name: String, is_self: bool, need_escape: bool) -> String { + let mut vars = HashMap::new(); + let index = if is_self { + OsRng.gen::() % BOT_TEXT_HANGED_SELF.len() + } else { + OsRng.gen::() % BOT_TEXT_HANGED.len() + }; + + let text = if is_self { + BOT_TEXT_HANGED_SELF[index] + } else { + BOT_TEXT_HANGED[index] + }; + + let name = if need_escape { + escape(name.as_str()) + } else { + name + }; + vars.insert("name".to_string(), name.as_str()); + + text.format(&vars).unwrap() +} + +pub async fn message_handler(db: &Controller, msg: Message, me: &Me) -> Result<(), RequestError> { + let text = match msg.text() { + Some(t) => t.to_owned(), + None => { + log_debug_ln!("{:?}", msg); + return Ok(()); + } + }; + + let formats = vec![BOT_TEXT_HANGED.to_vec(), BOT_TEXT_HANGED_SELF.to_vec()] + .concat() + .iter() + .map(|i| Regex::new(&format!("^{}$", i.replace("{name}", "(.+)"))).unwrap()) + .collect::>(); + + if let Some(via_bot) = msg.via_bot { + if via_bot.is_bot && via_bot.id == me.id { + for f in formats { + log_debug_ln!("Regexp {:?}", f); + if !f.is_match(text.as_str()) { + continue; + } + + if let Some(cap) = f.captures(text.as_str()) { + if let Some(name) = cap.get(1) { + log_debug_ln!("got username: {:?}", name.as_str()); + if let Err(error) = db + .update_group(name.as_str().to_string(), msg.chat.id) + .await + { + log_error_ln!("{:?}", error); + } + break; + } + } + } + } + } + Ok(()) +}