admin/
admin.rs

1use std::str::FromStr;
2
3use chrono::Duration;
4use teloxide::{prelude::*, types::ChatPermissions, utils::command::BotCommands};
5
6// Derive BotCommands to parse text with a command into this enumeration.
7//
8// 1. `rename_rule = "lowercase"` turns all the commands into lowercase letters.
9// 2. `description = "..."` specifies a text before all the commands.
10//
11// That is, you can just call Command::descriptions() to get a description of
12// your commands in this format:
13// %GENERAL-DESCRIPTION%
14// %PREFIX%%COMMAND% - %DESCRIPTION%
15
16/// Use commands in format /%command% %num% %unit%
17#[derive(BotCommands, Clone)]
18#[command(rename_rule = "lowercase", parse_with = "split")]
19enum Command {
20    /// Kick user from chat.
21    Kick,
22    /// Ban user in chat.
23    Ban {
24        time: u64,
25        unit: UnitOfTime,
26    },
27    /// Mute user in chat.
28    Mute {
29        time: u64,
30        unit: UnitOfTime,
31    },
32    Help,
33}
34
35#[derive(Clone)]
36enum UnitOfTime {
37    Seconds,
38    Minutes,
39    Hours,
40}
41
42impl FromStr for UnitOfTime {
43    type Err = &'static str;
44    fn from_str(s: &str) -> Result<Self, <Self as FromStr>::Err> {
45        match s {
46            "h" | "hours" => Ok(UnitOfTime::Hours),
47            "m" | "minutes" => Ok(UnitOfTime::Minutes),
48            "s" | "seconds" => Ok(UnitOfTime::Seconds),
49            _ => Err("Allowed units: h, m, s"),
50        }
51    }
52}
53
54#[tokio::main]
55async fn main() {
56    pretty_env_logger::init();
57    log::info!("Starting admin bot...");
58
59    let bot = teloxide::Bot::from_env();
60
61    Command::repl(bot, action).await;
62}
63
64async fn action(bot: Bot, msg: Message, cmd: Command) -> ResponseResult<()> {
65    match cmd {
66        Command::Help => {
67            bot.send_message(msg.chat.id, Command::descriptions().to_string()).await?;
68        }
69        Command::Kick => kick_user(bot, msg).await?,
70        Command::Ban { time, unit } => ban_user(bot, msg, calc_restrict_time(time, unit)).await?,
71        Command::Mute { time, unit } => mute_user(bot, msg, calc_restrict_time(time, unit)).await?,
72    };
73
74    Ok(())
75}
76
77// Kick a user with a replied message.
78async fn kick_user(bot: Bot, msg: Message) -> ResponseResult<()> {
79    match msg.reply_to_message() {
80        Some(replied) => {
81            // bot.unban_chat_member can also kicks a user from a group chat.
82            bot.unban_chat_member(msg.chat.id, replied.from.as_ref().unwrap().id).await?;
83        }
84        None => {
85            bot.send_message(msg.chat.id, "Use this command in reply to another message").await?;
86        }
87    }
88    Ok(())
89}
90
91// Ban a user with replied message.
92async fn ban_user(bot: Bot, msg: Message, time: Duration) -> ResponseResult<()> {
93    match msg.reply_to_message() {
94        Some(replied) => {
95            bot.kick_chat_member(
96                msg.chat.id,
97                replied.from.as_ref().expect("Must be MessageKind::Common").id,
98            )
99            .until_date(msg.date + time)
100            .await?;
101        }
102        None => {
103            bot.send_message(msg.chat.id, "Use this command in a reply to another message!")
104                .await?;
105        }
106    }
107    Ok(())
108}
109
110// Mute a user with a replied message.
111async fn mute_user(bot: Bot, msg: Message, time: Duration) -> ResponseResult<()> {
112    match msg.reply_to_message() {
113        Some(replied) => {
114            bot.restrict_chat_member(
115                msg.chat.id,
116                replied.from.as_ref().expect("Must be MessageKind::Common").id,
117                ChatPermissions::empty(),
118            )
119            .until_date(msg.date + time)
120            .await?;
121        }
122        None => {
123            bot.send_message(msg.chat.id, "Use this command in a reply to another message!")
124                .await?;
125        }
126    }
127    Ok(())
128}
129
130// Calculates time of user restriction.
131fn calc_restrict_time(time: u64, unit: UnitOfTime) -> Duration {
132    // FIXME: actually handle the case of too big integers correctly, instead of
133    // unwrapping
134    match unit {
135        UnitOfTime::Hours => Duration::try_hours(time as i64).unwrap(),
136        UnitOfTime::Minutes => Duration::try_minutes(time as i64).unwrap(),
137        UnitOfTime::Seconds => Duration::try_seconds(time as i64).unwrap(),
138    }
139}