inline 模式

Cargo.lock

@@ -1126,6 +1126,7 @@ dependencies = [
  "migration",
  "models",
  "rand",
+ "regex",
  "reqwest",
  "sea-orm",
  "shadow-rs",

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]

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<Self, DbErr> {
-        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<Message, RequestError> {
-        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<Message, RequestError> {
-        self.send_text_reply(bot, message, Commands::descriptions().to_string())
-            .await
-    }
-
-    pub async fn about_handler(
-        &self,
-        bot: &Bot,
-        message: &Message,
-    ) -> Result<Message, RequestError> {
-        self.send_text_reply(bot, message, BOT_ABOUT.to_string())
-            .await
-    }
-
-    pub async fn hangit_handler(
-        &self,
-        bot: &Bot,
-        message: &Message,
-    ) -> Result<Message, RequestError> {
-        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::<usize>() % BOT_TEXT_HANGED_SELF.len()
-                } else {
-                    OsRng.gen::<usize>() % 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<Message, 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 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<String, String> = 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<String, String> = 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
+}

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,
@@ -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<Vec<TopData>> {
-        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<Vec<StatsModel>> {
+        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
+            }
+        }
+    }
 }

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::<Vec<_>>(),
+
+        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(())
+}

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::<Commands>().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),
     }
 }

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] = [
@@ -38,3 +40,4 @@ 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 = "这是个幽灵……";
+pub const BOT_TEXT_INLINE_HANG: &'static str = "吊死";

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::<usize>() % BOT_TEXT_HANGED_SELF.len()
+    } else {
+        OsRng.gen::<usize>() % 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::<Vec<_>>();
+
+    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(())
+}