use hbb_common::{ anyhow::anyhow, bail, config::Config, get_time, password_security::{decrypt_vec_or_original, encrypt_vec_or_original}, ResultType, }; use serde_derive::{Deserialize, Serialize}; use std::sync::Mutex; use totp_rs::{Algorithm, Secret, TOTP}; lazy_static::lazy_static! { static ref CURRENT_2FA: Mutex> = Mutex::new(None); } const ISSUER: &str = "RustDesk"; const TAG_LOGIN: &str = "Connection"; #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct TOTPInfo { pub name: String, pub secret: Vec, pub digits: usize, pub created_at: i64, } impl TOTPInfo { fn new_totp(&self) -> ResultType { let totp = TOTP::new( Algorithm::SHA1, self.digits, 1, 30, self.secret.clone(), Some(format!("{} {}", ISSUER, TAG_LOGIN)), self.name.clone(), )?; Ok(totp) } fn gen_totp_info(name: String, digits: usize) -> ResultType { let secret = Secret::generate_secret(); let totp = TOTPInfo { secret: secret.to_bytes()?, name, digits, created_at: get_time(), ..Default::default() }; Ok(totp) } pub fn into_string(&self) -> ResultType { let secret = encrypt_vec_or_original(self.secret.as_slice(), "00", 1024); let totp_info = TOTPInfo { secret, ..self.clone() }; let s = serde_json::to_string(&totp_info)?; Ok(s) } pub fn from_str(data: &str) -> ResultType { let mut totp_info = serde_json::from_str::(data)?; let (secret, success, _) = decrypt_vec_or_original(&totp_info.secret, "00"); if success { totp_info.secret = secret; return Ok(totp_info.new_totp()?); } else { bail!("decrypt_vec_or_original 2fa secret failed") } } } pub fn generate2fa() -> String { #[cfg(not(any(target_os = "android", target_os = "ios")))] let id = crate::ipc::get_id(); #[cfg(any(target_os = "android", target_os = "ios"))] let id = Config::get_id(); if let Ok(info) = TOTPInfo::gen_totp_info(id, 6) { if let Ok(totp) = info.new_totp() { let code = totp.get_url(); *CURRENT_2FA.lock().unwrap() = Some((info, totp)); return code; } } "".to_owned() } pub fn verify2fa(code: String) -> bool { if let Some((info, totp)) = CURRENT_2FA.lock().unwrap().as_ref() { if let Ok(res) = totp.check_current(&code) { if res { if let Ok(v) = info.into_string() { #[cfg(not(any(target_os = "android", target_os = "ios")))] crate::ipc::set_option("2fa", &v); #[cfg(any(target_os = "android", target_os = "ios"))] Config::set_option("2fa".to_owned(), v); return res; } } } } false } pub fn get_2fa(raw: Option) -> Option { TOTPInfo::from_str(&raw.unwrap_or(Config::get_option("2fa"))) .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, pub chat_id: String, } impl TelegramBot { fn into_string(&self) -> ResultType { 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 { let data = Config::get_option("telegram_bot"); let mut bot = serde_json::from_str::(&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> { // 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::(&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 }