telegram_bot2/
bot_builder.rs

1use crate::__private::{GenericHandler, HandlerError};
2use crate::models::{GetUpdates, Update, UpdateType};
3use crate::{Bot, BotState};
4use std::env::var;
5use std::fs::File;
6use std::io::Read;
7
8#[cfg(feature = "commands")]
9use {crate::__private::command::CommandHandler, crate::models::MessageEntityType, lazy_static::lazy_static, regex::Regex, std::collections::BTreeMap};
10#[cfg(feature = "daemons")]
11use {crate::__private::daemon::DaemonStruct, tokio::sync::Mutex};
12
13use log::{error, info, trace};
14use std::sync::Arc;
15use std::time::Duration;
16
17/// Builds a bot.
18/// This should be used in combination with the [bot][macro@crate::bot] macro
19pub struct BotBuilder<T>
20where
21    T: Sync + Send + 'static,
22{
23    bot: Arc<Bot>,
24    state: Arc<Option<BotState<T>>>,
25
26    #[cfg(feature = "daemons")]
27    daemons: Vec<DaemonStruct<T>>,
28
29    #[cfg(feature = "commands")]
30    commands: Option<BTreeMap<String, Vec<CommandHandler<T>>>>,
31
32    handler: Vec<GenericHandler<T>>,
33
34    update_limit: Option<i64>,
35    allowed_updates: Option<Vec<UpdateType>>,
36    offset: Option<i64>,
37    timeout: Option<i64>,
38    interval: Option<Duration>,
39}
40
41impl BotBuilder<()> {
42    /// Builds a bot and uses the environment variable `BOT_TOKEN` as token or extract it from the `BOT_TOKEN_FILE` file. If both are defined, `BOT_TOKEN` is used.
43    pub fn new() -> BotBuilder<()> {
44        BotBuilder {
45            bot: Arc::new(Bot::login(
46                var("BOT_TOKEN")
47                    .ok()
48                    .or(var("BOT_TOKEN_FILE").ok().and_then(|f| {
49                        let mut s = String::new();
50                        File::open(f).expect("Could not open token file").read_to_string(&mut s).expect("Could not read from token file");
51                        Some(s)
52                    }))
53                    .expect("No token provided (define BOT_TOKEN or BOT_TOKEN_FILE variable)"),
54            )),
55            state: Arc::new(None),
56            #[cfg(feature = "commands")]
57            commands: None,
58            #[cfg(feature = "daemons")]
59            daemons: vec![],
60            handler: vec![],
61            update_limit: None,
62            allowed_updates: None,
63            offset: None,
64            timeout: None,
65            interval: None,
66        }
67    }
68
69    /// Builds a bot with the given token
70    pub fn new_with_token(token: String) -> BotBuilder<()> {
71        BotBuilder {
72            bot: Arc::new(Bot::login(token)),
73            state: Arc::new(None),
74            #[cfg(feature = "commands")]
75            commands: None,
76            #[cfg(feature = "daemons")]
77            daemons: vec![],
78            handler: vec![],
79            update_limit: None,
80            allowed_updates: None,
81            offset: None,
82            timeout: None,
83            interval: None,
84        }
85    }
86}
87
88impl Default for BotBuilder<()> {
89    fn default() -> Self {
90        BotBuilder::new()
91    }
92}
93
94impl<T: Sync + Send + 'static> BotBuilder<T> {
95    /// Add a shared state to the bot
96    /// The state can be used by any update or command handler
97    pub fn with_state<S>(self, state: S) -> BotBuilder<S>
98    where
99        S: Sync + Send + 'static,
100    {
101        if self.state.is_some() {
102            panic!("State already set")
103        }
104        if !self.handler.is_empty() {
105            panic!("Cannot set state after handler")
106        }
107        #[cfg(feature = "daemons")]
108        if !self.daemons.is_empty() {
109            panic!("Cannot set state after daemons")
110        }
111        #[cfg(feature = "commands")]
112        if self.commands.is_some() {
113            panic!("Cannot set state after commands")
114        }
115
116        BotBuilder {
117            bot: self.bot,
118            state: Arc::new(Some(BotState { state })),
119            #[cfg(feature = "commands")]
120            commands: None,
121            #[cfg(feature = "daemons")]
122            daemons: vec![],
123            handler: vec![],
124            update_limit: self.update_limit,
125            allowed_updates: self.allowed_updates,
126            offset: self.offset,
127            timeout: self.timeout,
128            interval: self.interval,
129        }
130    }
131
132    /// Set the update handler of the bot
133    pub fn handlers(mut self, handlers: Vec<GenericHandler<T>>) -> Self {
134        if !self.handler.is_empty() {
135            panic!("Handlers already set")
136        }
137
138        handlers.iter().for_each(|h| info!("Registered handler {} with priority {}", h.name, h.rank));
139
140        self.handler = handlers;
141        self.handler.sort_by_key(|f| f.rank);
142        self
143    }
144
145    /// Set the maximum fetched updates per poll
146    /// Default is 100, and max is 100
147    pub fn update_limit(mut self, limit: i64) -> BotBuilder<T> {
148        if self.update_limit.is_some() {
149            panic!("Update limit already set")
150        }
151
152        self.update_limit = Some(limit);
153        self
154    }
155
156    /// Specify the time (in seconds) before the polling exits with timeout
157    /// Should be a stricly positive integer
158    pub fn timeout(mut self, timeout: i64) -> BotBuilder<T> {
159        if self.timeout.is_some() {
160            panic!("Timeout already set")
161        }
162        // TODO warn for short timeouts
163
164        self.timeout = Some(timeout);
165        self
166    }
167
168    /// Minimum waiting time between two pollings
169    pub fn interval(mut self, interval: Duration) -> BotBuilder<T> {
170        if self.interval.is_some() {
171            panic!("Interval already set")
172        }
173
174        self.interval = Some(interval);
175        self
176    }
177
178    #[cfg(feature = "commands")]
179    /// Set the command handlers (see [`commands`][crate::commands])
180    pub fn commands(mut self, commands: BTreeMap<String, Vec<CommandHandler<T>>>) -> BotBuilder<T> {
181        if self.commands.is_some() {
182            panic!("Commands already set")
183        }
184
185        for cmd in &commands {
186            info!("Registering handler for command {}: ", cmd.0);
187            for handler in cmd.1 {
188                info!(r#" - "{}" with priority {}"#, handler.syntax, handler.rank);
189            }
190        }
191
192        self.commands = Some(commands);
193        self
194    }
195
196    #[cfg(feature = "daemons")]
197    /// Set the daemons of the bot (see [`daemon`][crate::daemon])
198    pub fn daemons(mut self, daemons: Vec<DaemonStruct<T>>) -> BotBuilder<T> {
199        if !self.daemons.is_empty() {
200            panic!("Daemons already set")
201        }
202
203        daemons.iter().for_each(|d| info!("Registering daemon {}", d.syntax));
204
205        self.daemons = daemons;
206        self
207    }
208
209    /// Start the bot (should not be used, see [`bot`][macro@crate::bot])
210    pub async fn launch(mut self) {
211        let allowed_updates = self.allowed_updates.unwrap_or_default();
212        let timeout = self.timeout.or(Some(5));
213
214        info!("Launching bot with parameters: [interval: {:?}, timeout: {:?}]", self.interval, timeout);
215
216        #[cfg(feature = "daemons")]
217        Self::launch_daemons(self.daemons, self.bot.clone(), self.state.clone()).await;
218
219        loop {
220            if let Some(interval) = self.interval {
221                tokio::time::sleep(interval).await;
222            }
223
224            // Wait for an update to come
225            let mut updates = vec![];
226            while updates.is_empty() {
227                updates = self
228                    .bot
229                    .get_updates(GetUpdates {
230                        offset: self.offset,
231                        limit: self.update_limit,
232                        timeout,
233                        allowed_updates: {
234                            #[cfg(feature = "commands")]
235                            if self.commands.is_some() {
236                                [allowed_updates.clone(), vec![UpdateType::NewMessage]].concat()
237                            } else {
238                                allowed_updates.clone()
239                            }
240
241                            #[cfg(not(feature = "commands"))]
242                            allowed_updates.clone()
243                        },
244                    })
245                    .await
246                    .unwrap()
247            }
248
249            // Sort all updates by their id such that they are processed correctly
250            updates.sort_by_key(|u| u.update_id);
251
252            // Set the offset to the id after the last one recieved
253            // We may unwrap it, as the loop above ensure that the vec is not empty
254            self.offset = Some(updates.last().unwrap().update_id + 1);
255
256            for update in updates {
257                // Check whether the update is a command and there is a valid registered handleer
258                if let Ok(update) = Update::try_from(update.clone()) {
259                    trace!("Received update {} ({:?})", update.get_id(), update.get_type());
260
261                    #[cfg(feature = "commands")]
262                    if Self::try_commands(&self.commands, &self.bot, self.state.as_ref().as_ref(), &update).await {
263                        continue;
264                    }
265
266                    // Else pass the update to the default handler
267                    for h in &self.handler {
268                        if h.updates.is_empty() || h.updates.iter().any(|u| *u == update.get_type()) {
269                            match h.handler.handle(&self.bot, &update, self.state.as_ref().as_ref()).await {
270                                Ok(_) => break,
271                                Err(HandlerError::Parse) => continue,
272                                Err(HandlerError::Runtime) => {
273                                    error!("Handler failure for {:?}", update);
274                                    break;
275                                }
276                            }
277                        }
278                    }
279                } else {
280                    error!("Invalid update received: {:?}", update)
281                }
282            }
283        }
284    }
285
286    #[cfg(feature = "commands")]
287    /// Tries to parse the update for as a command. Returns whether it was successful
288    async fn try_commands(commands: &Option<BTreeMap<String, Vec<CommandHandler<T>>>>, bot: &Bot, state: Option<&BotState<T>>, update: &Update) -> bool {
289        if let (Update::NewMessage(_, m), Some(commands)) = (&update, commands) {
290            if let Some(e) = m.entities.iter().find(|e| e.offset == 0 && matches!(e._type, MessageEntityType::BotCommand)) {
291                trace!("Command \"{}\" issued", m.text.as_ref().unwrap_or(&"".to_string()));
292
293                lazy_static! {
294                    static ref PARAMETER_DELIMITER: Regex = Regex::new(r"\s+").unwrap();
295                }
296
297                if let Some(handlers) = commands.get(&m.text.as_ref().unwrap().as_str()[(e.offset + 1) as usize..(e.offset + e.length) as usize]) {
298                    let args: Vec<String> = PARAMETER_DELIMITER
299                        .split(&m.text.as_ref().unwrap().as_str()[(e.offset + e.length) as usize..])
300                        .filter(|s| !s.is_empty())
301                        .map(str::to_string)
302                        .collect();
303
304                    for h in handlers {
305                        match h.handler.handle(&args, update, bot, state).await {
306                            Ok(_) => break,
307                            Err(HandlerError::Parse) => continue,
308                            Err(HandlerError::Runtime) => panic!(),
309                        }
310                    }
311
312                    trace!("Could not delegate to any handler");
313
314                    return true;
315                } else {
316                    trace!("No corresponding handler found")
317                }
318            }
319        }
320
321        false
322    }
323
324    #[cfg(feature = "daemons")]
325    /// Launch all the bot daemons and returns their lock
326    async fn launch_daemons(daemons: Vec<DaemonStruct<T>>, bot: Arc<Bot>, state: Arc<Option<BotState<T>>>) -> Vec<Arc<Mutex<()>>> {
327        daemons
328            .into_iter()
329            .map(|d| {
330                let mutex = Arc::new(Mutex::new(()));
331
332                let bot = bot.clone();
333                let state = state.clone();
334                let lock = mutex.clone();
335
336                tokio::task::spawn(async move { d.daemon.launch(bot, state, lock).await });
337
338                mutex
339            })
340            .collect()
341    }
342}
telegram_bot2/bot_builder.rs

telegram_bot2/
bot_builder.rs
