简易的tgbot入群验证队列

发布于 2024-01-21  8 次阅读


use chrono::{DateTime, Duration, Utc};
use log::{debug, error, info};
use std::collections::VecDeque;
use std::sync::Arc;
use teloxide::{
    types::{ChatId, UserId},
    Bot,
};
use tokio::sync::{Mutex, Notify};

use crate::handler;

#[derive(Debug, Clone)]
pub struct QueueElement {
    pub id: String,
    pub members: Vec<Members>,
    pub join_time: DateTime<Utc>,
}

#[derive(Debug, Clone)]
pub struct Members {
    pub chat_id: ChatId,
    pub user_id: UserId,
}

#[derive(Debug)]
pub struct WaitingQueueManager {
    queue: Arc<Mutex<VecDeque<QueueElement>>>,
    notify: Notify,
}

impl Default for WaitingQueueManager {
    fn default() -> Self {
        Self::new()
    }
}

impl WaitingQueueManager {
    pub fn new() -> Self {
        WaitingQueueManager {
            queue: Arc::new(Mutex::new(VecDeque::new())),
            notify: Notify::new(),
        }
    }

    pub async fn add_to_waiting_queue(&self, element: QueueElement) {
        let mut waiting_queue = self.queue.lock().await;
        waiting_queue.push_back(element);
        self.notify.notify_waiters();
        debug!("Element added to waiting queue: {:?}", waiting_queue);
    }

    async fn remove_from_waiting_queue_element(&self) {
        let mut waiting_queue = self.queue.lock().await;
        waiting_queue.pop_front();
    }

    pub async fn remove_from_waiting_queue(&self, chat_id: ChatId, user_id: UserId) {
        let mut waiting_queue = self.queue.lock().await;

        waiting_queue.iter_mut().for_each(|m| {
            m.members
                .retain(|f| !(f.chat_id == chat_id && f.user_id == user_id));
        });

        waiting_queue.retain(|m| !m.members.is_empty());
    }

    pub async fn process_waiting_queue(&self, bot: Bot) {
        loop {
            let front_value = self.queue.lock().await.front().cloned();

            if let Some(element) = front_value {
                let next_wait_time =
                    Duration::seconds(300) - Utc::now().signed_duration_since(element.join_time);
                tokio::time::sleep(next_wait_time.to_std().unwrap()).await;

                let waiting_queue = self.queue.lock().await;

                debug!(" waiting_queue: {:?}", waiting_queue);

                if let Some(updated_element) = waiting_queue.front().cloned() {
                    if updated_element.id == element.id {
                        for member in &element.members {
                            let chat_id = member.chat_id;
                            let user_id = member.user_id;

                            let bot_clone = bot.clone();
                            tokio::spawn(async move {
                                info!(
                                    "User {} has been in the queue for more than 5 minutes. Kicking...",
                                    user_id
                                );

                                // Kick the user
                                if let Err(err) =
                                    handler::kick_user(bot_clone, chat_id, user_id).await
                                {
                                    error!("Failed to kick user: {}", err);
                                }
                            });
                        }

                        drop(waiting_queue);
                        self.remove_from_waiting_queue_element().await;
                    } else {
                        continue;
                    }
                } else {
                    continue;
                }
            } else {
                info!("Waiting job...");
                self.notify.notified().await;
                info!("having...");
            }
        }
    }
}
僕と契約して、魔法少女になってよ!
最后更新于 2024-03-29