From ac2f9b937509c87937918c67ce763d54d1ab5d68 Mon Sep 17 00:00:00 2001 From: Awiteb Date: Wed, 5 Jun 2024 18:32:13 +0300 Subject: [PATCH] feat: Support `cmd` and `args` querys Signed-off-by: Awiteb --- src/api.rs | 45 +++++++++++++++++++++++++++++++++++++++++++-- src/main.rs | 1 + src/superbot.rs | 8 ++++++-- 3 files changed, 50 insertions(+), 4 deletions(-) diff --git a/src/api.rs b/src/api.rs index 04c046a..4771aeb 100644 --- a/src/api.rs +++ b/src/api.rs @@ -16,6 +16,7 @@ use std::sync::{Arc, Mutex}; +use base64::Engine; use salvo::{catcher::Catcher, http::HeaderValue, hyper::header, logging::Logger, prelude::*}; use crate::{superbot, PingList}; @@ -82,12 +83,46 @@ fn write_json_body(res: &mut Response, json_body: impl serde::Serialize) { .ok(); } +fn decode_base64url_or( + or: impl FnOnce() -> String, + res: &mut Response, + query: Option, +) -> Option { + log::debug!("Decode {query:?}"); + query.map_or_else( + || Some(or()), + |query| { + crate::Base64Url.decode(query).map_or_else( + |_| { + log::error!("Is invalid base64Url"); + res.status_code(StatusCode::BAD_REQUEST); + write_json_body( + res, + MessageSchema::new("Invalid base64Url args").code(StatusCode::BAD_REQUEST), + ); + None + }, + |decoded_query| Some(String::from_utf8_lossy(&decoded_query).into_owned()), + ) + }, + ) +} + #[handler] async fn ping(req: &Request, res: &mut Response, depot: &mut Depot) { let bot_username = req .param::("bot_username") .expect("The path param is required") .to_lowercase(); + + let Some(bot_cmd) = decode_base64url_or(|| "start".to_owned(), res, req.query("cmd")) else { + return; + }; + let Some(cmd_args) = decode_base64url_or(String::new, res, req.query("args")) else { + return; + }; + + let app_state = depot .obtain::>() .expect("The app state is injected"); @@ -99,8 +134,14 @@ async fn ping(req: &Request, res: &mut Response, depot: &mut Depot) { let msg = if !app_state.bots.contains(&bot_username) { MessageSchema::new("Is not authorized to check the status of this bot") .code(StatusCode::BAD_REQUEST) - } else if let Ok(telegram_id) = - superbot::send_start(&app_state.tg_client, &bot_username, bots).await + } else if let Ok(telegram_id) = superbot::send_command( + &app_state.tg_client, + &bot_username, + bots, + &bot_cmd, + &cmd_args, + ) + .await { if bots.check(telegram_id) { MessageSchema::new("Alive") diff --git a/src/main.rs b/src/main.rs index f092458..29af004 100644 --- a/src/main.rs +++ b/src/main.rs @@ -30,6 +30,7 @@ mod errors; mod superbot; mod traits; +pub(crate) use base64::engine::general_purpose::URL_SAFE_NO_PAD as Base64Url; pub(crate) use errors::{Error as ServerError, Result as ServerResult}; use tokio::signal; pub(crate) use traits::PingList; diff --git a/src/superbot.rs b/src/superbot.rs index e1a371d..a4253ef 100644 --- a/src/superbot.rs +++ b/src/superbot.rs @@ -101,16 +101,20 @@ pub(crate) async fn handler(client: Client, bots: Arc>>, + command: &str, + args: &str, ) -> ServerResult { if let Some(chat) = client.resolve_username(bot_username).await? { log::debug!("Bots: {bots:?}"); let telegram_id = chat.id() as u64; bots.add_new(telegram_id); - client.send_message(chat, "/start").await?; + client + .send_message(chat, format!("/{command} {args}")) + .await?; // Sleep, wating the response time::sleep(time::Duration::from_secs(2)).await; Ok(telegram_id)