From faf363cfd27f58b5755fbd73aa1ce4b882180862 Mon Sep 17 00:00:00 2001
From: rustdesk <info@rustdesk.com>
Date: Wed, 26 Jun 2024 18:49:27 +0800
Subject: [PATCH] add TelegramBot

---
 src/auth_2fa.rs | 73 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 73 insertions(+)

diff --git a/src/auth_2fa.rs b/src/auth_2fa.rs
index 1ea773342..f46bc71e6 100644
--- a/src/auth_2fa.rs
+++ b/src/auth_2fa.rs
@@ -1,4 +1,5 @@
 use hbb_common::{
+    anyhow::anyhow,
     bail,
     config::Config,
     get_time,
@@ -109,3 +110,75 @@ pub fn get_2fa(raw: Option<String>) -> Option<TOTP> {
         .map(|x| Some(x))
         .unwrap_or_default()
 }
+
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct TelegramBot {
+    #[serde(skip)]
+    pub token_str: String,
+    pub token: Vec<u8>,
+    pub chat_id: String,
+}
+
+impl TelegramBot {
+    fn into_string(&self) -> ResultType<String> {
+        let token = encrypt_vec_or_original(self.token_str.as_bytes(), "00", 1024);
+        let bot = TelegramBot {
+            token,
+            ..self.clone()
+        };
+        let s = serde_json::to_string(&bot)?;
+        Ok(s)
+    }
+
+    fn save(&self) -> ResultType<()> {
+        let s = self.into_string()?;
+        #[cfg(not(any(target_os = "android", target_os = "ios")))]
+        crate::ipc::set_option("telegram_bot", &s);
+        #[cfg(any(target_os = "android", target_os = "ios"))]
+        Config::set_option("telegram_bot".to_owned(), s);
+        Ok(())
+    }
+
+    fn get() -> ResultType<TelegramBot> {
+        let data = Config::get_option("telegram_bot");
+        let mut bot = serde_json::from_str::<TelegramBot>(&data)?;
+        let (token, success, _) = decrypt_vec_or_original(&bot.token, "00");
+        if success {
+            bot.token_str = String::from_utf8(token)?;
+            return Ok(bot);
+        }
+        bail!("decrypt_vec_or_original telegram bot token failed")
+    }
+}
+
+// https://gist.github.com/dideler/85de4d64f66c1966788c1b2304b9caf1
+pub async fn send_2fa_code_to_telegram(code: &str) -> ResultType<()> {
+    let bot = TelegramBot::get()?;
+    let url = format!("https://api.telegram.org/bot{}/sendMessage", bot.token_str);
+    let params = serde_json::json!({"chat_id": bot.chat_id, "text": code});
+    crate::post_request(url, params.to_string(), "").await?;
+    Ok(())
+}
+
+pub async fn get_chatid_telegram(bot_token: &str) -> ResultType<Option<String>> {
+    // send a message to the bot first please, otherwise the chat_id will be empty
+    let url = format!("https://api.telegram.org/bot{}/getUpdates", bot_token);
+    let resp = crate::post_request(url, "".to_owned(), "")
+        .await
+        .map_err(|e| anyhow!(e))?;
+    let res = serde_json::from_str::<serde_json::Value>(&resp)
+        .map(|x| {
+            let chat_id = x["result"][0]["message"]["chat"]["id"].as_str();
+            chat_id.map(|x| x.to_owned())
+        })
+        .map_err(|e| anyhow!(e));
+    if let Ok(Some(chat_id)) = res.as_ref() {
+        let bot = TelegramBot {
+            token_str: bot_token.to_owned(),
+            chat_id: chat_id.to_owned(),
+            ..Default::default()
+        };
+        bot.save()?;
+    }
+    res
+}