From 83148f848ac65c3b8ed4a28003b35b3fc034c6b4 Mon Sep 17 00:00:00 2001 From: Awiteb Date: Fri, 31 May 2024 18:49:27 +0300 Subject: [PATCH] feat: Make `config.toml` hold all configrations Signed-off-by: Awiteb --- src/config.rs | 122 ++++++++++++++++++++++++++++++++++++++++++++++++++ src/errors.rs | 2 + src/main.rs | 75 +++++++++++-------------------- 3 files changed, 149 insertions(+), 50 deletions(-) create mode 100644 src/config.rs diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..30b05a9 --- /dev/null +++ b/src/config.rs @@ -0,0 +1,122 @@ +// A simple API to ping telegram bots and returns if it's online or not. +// Copyright (C) 2023-2024 Awiteb +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published +// by the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::{fs, path::Path}; + +use serde::{de::Error as DeError, Deserialize}; + +use crate::{ServerError, ServerResult}; + +/// The config of telegram client +#[derive(Deserialize, Debug)] +pub(crate) struct TelegramClient { + pub api_hash: String, + pub api_id: i32, +} + +/// The config of the api +#[derive(Deserialize, Debug)] +pub(crate) struct ApiConfig { + #[serde(deserialize_with = "host_deserialize")] + #[serde(default = "default_host")] + pub host: String, + #[serde(default = "default_port")] + pub port: u16, +} + +/// The config struct +#[derive(Deserialize, Debug)] +pub(crate) struct Config { + pub client: TelegramClient, + pub api: ApiConfig, + #[serde(deserialize_with = "bots_deserialize")] + pub bots: Vec, + #[serde(deserialize_with = "one_or_more_string")] + pub tokens: Vec, +} + +impl Config { + /// Initialize the config from toml file + pub fn from_toml_file(file_path: impl AsRef) -> ServerResult { + toml::from_str( + &fs::read_to_string(file_path).map_err(|err| ServerError::Config(err.to_string()))?, + ) + .map_err(|err| ServerError::Config(err.to_string())) + } +} + +/// A function to deserialize the host, make sure it's a valid host to band +fn host_deserialize<'de, D: serde::Deserializer<'de>>(d: D) -> Result { + let host = String::deserialize(d)?; + if host == "localhost" { + return Ok(host); + } + let octets = host + .split(".") + .map(|octet| octet.parse::()) + .collect::, _>>(); + if let Ok(octets) = octets { + if octets.len() != 4 { + return Err(DeError::custom("There is more then 4 octets")); + } + } else { + return Err(DeError::custom("Contain invalid number")); + } + Ok(host) +} + +fn bots_deserialize<'de, D: serde::Deserializer<'de>>(d: D) -> Result, D::Error> { + let bots = Vec::::deserialize(d)?; + if bots.is_empty() { + return Err(DeError::custom("There must be one or more bots")); + } + + for bot in &bots { + if !bot.starts_with('@') { + return Err(DeError::custom(format!( + "Invalid bot username `{bot}`: must starts with `@`" + ))); + } else if !bot.to_lowercase().ends_with("bot") { + return Err(DeError::custom(format!( + "Invalid bot username `{bot}`: must end with `bot`" + ))); + } + } + + Ok(bots.iter().map(|s| s.trim().to_owned()).collect()) +} + +fn one_or_more_string<'de, D>(d: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + let str_vec = Vec::::deserialize(d)?; + if str_vec.is_empty() { + return Err(DeError::custom("There is must be one at least")); + } + Ok(str_vec.iter().map(|s| s.trim().to_owned()).collect()) +} + + +/// The default host `0.0.0.0` +fn default_host() -> String { + "0.0.0.0".to_owned() +} + +/// The default port `3939` +const fn default_port() -> u16 { + 3939 +} diff --git a/src/errors.rs b/src/errors.rs index 7debbbf..c4bc270 100644 --- a/src/errors.rs +++ b/src/errors.rs @@ -28,6 +28,8 @@ pub(crate) type Result = std::result::Result; pub(crate) enum Error { #[error("Cli Error: {0}")] CliParse(String), + #[error("Config Error: {0}")] + Config(String), #[error("IO Error: {0}")] Io(#[from] IoError), #[error("Thread Error: {0}")] diff --git a/src/main.rs b/src/main.rs index e31f34d..63d2b5e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,18 +16,20 @@ #![doc = include_str!("../README.md")] -use std::{env, fs, process::ExitCode, sync::Mutex}; +use std::{process::ExitCode, sync::Mutex}; use lazy_static::lazy_static; use salvo::{conn::TcpListener, Listener}; mod api; mod cli_parser; +mod config; mod errors; mod superbot; mod traits; pub(crate) use errors::{Error as ServerError, Result as ServerResult}; +use tokio::signal; pub(crate) use traits::PingList; #[derive(Default, Clone)] @@ -67,59 +69,32 @@ async fn try_main() -> ServerResult<()> { return Ok(()); } + let config = config::Config::from_toml_file(&cli_args.config_file)?; - let bots: Vec = fs::read_to_string("bots.txt")? - .lines() - .map(|b| b.trim().to_owned()) - .collect(); - let tokens: Vec = fs::read_to_string("tokens.txt")? - .lines() - .map(|b| b.trim().to_owned()) - .collect(); + let (client, sign_out) = superbot::login(config.client.api_hash, config.client.api_id).await?; + let app_state = api::AppState::new(config.bots, config.tokens, client.clone()); - if bots - .iter() - .any(|b| !b.starts_with('@') || !b.to_lowercase().ends_with("bot")) - { - bots.iter().for_each(|b| { - if !b.starts_with('@') { - eprintln!("Invalid bot username `{b}`: must starts with `@`"); - } else if !b.to_lowercase().ends_with("bot") { - eprintln!("Invalid bot username `{b}`: must end with `bot`"); - } - }) - } else { - let (client, sign_out) = superbot::login( - env::var("TELEPINGBOT_API_HASH") - .expect("`TELEPINGBOT_API_HASH` environment variable is required"), - env::var("TELEPINGBOT_API_ID") - .expect("`TELEPINGBOT_API_ID` environment variable is required") - .parse() - .expect("Invalid value for `TELEPINGBOT_API_ID` must be a number"), - ) - .await?; - let host = env::var("TELEOINGBOT_HOST") - .expect("`TELEOINGBOT_HOST` environment variable must be set"); - let port = env::var("TELEOINGBOT_PORT") - .expect("`TELEOINGBOT_PORT` environment variable must be set"); - let app_state = api::AppState::new(bots, tokens, client.clone()); + let handler_client = client.clone(); + let acceptor = TcpListener::new(format!("{}:{}", config.api.host, config.api.port)) + .bind() + .await; + let client_handler = tokio::spawn(async move { superbot::handler(handler_client).await }); + let server_handler = tokio::spawn(async move { + salvo::Server::new(acceptor) + .serve(api::service(app_state)) + .await + }); + log::info!("Bind the API to {}:{}", config.api.host, config.api.port); - let handler_client = client.clone(); - let acceptor = TcpListener::new(format!("{host}:{port}")).bind().await; - let client_handler = tokio::spawn(async move { superbot::handler(handler_client).await }); - let server_handler = tokio::spawn(async move { - salvo::Server::new(acceptor) - .serve(api::service(app_state)) - .await - }); + tokio::select! { + _ = client_handler => {}, + _ = server_handler=> {}, + _ = signal::ctrl_c() => {}, + } - client_handler.await?; - server_handler.await?; - - log::debug!("Close the API, telegram sign out status: {sign_out}"); - if sign_out { - client.sign_out_disconnect().await?; - } + log::debug!("Close the API, telegram sign out status: {sign_out}"); + if sign_out { + client.sign_out_disconnect().await?; } Ok(()) }