From ac2f9b937509c87937918c67ce763d54d1ab5d68 Mon Sep 17 00:00:00 2001
From: Awiteb <a@4rs.nl>
Date: Wed, 5 Jun 2024 18:32:13 +0300
Subject: [PATCH] feat: Support `cmd` and `args` querys

Signed-off-by: Awiteb <a@4rs.nl>
---
 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<String>,
+) -> Option<String> {
+    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::<String>("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::<Arc<AppState>>()
         .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<Mutex<Vec<crate::PingedBot
     }
 }
 
-pub(crate) async fn send_start(
+pub(crate) async fn send_command(
     client: &Client,
     bot_username: &str,
     bots: &Arc<Mutex<Vec<crate::PingedBot>>>,
+    command: &str,
+    args: &str,
 ) -> ServerResult<u64> {
     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)