Skip to content

Commit

Permalink
add TelegramBot
Browse files Browse the repository at this point in the history
  • Loading branch information
rustdesk committed Jun 26, 2024
1 parent dbbd917 commit faf363c
Showing 1 changed file with 73 additions and 0 deletions.
73 changes: 73 additions & 0 deletions src/auth_2fa.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use hbb_common::{
anyhow::anyhow,
bail,
config::Config,
get_time,
Expand Down Expand Up @@ -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
}

0 comments on commit faf363c

Please sign in to comment.