diff --git a/Cargo.lock b/Cargo.lock index 72cd045..4022763 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1739,7 +1739,6 @@ dependencies = [ "infer", "rand", "regex", - "serde", "teloxide", "tempfile", "thiserror 2.0.16", diff --git a/Cargo.toml b/Cargo.toml index 69973a4..6c2535f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,7 +14,6 @@ futures = "0.3" infer = "0.19" rand = "0.9" regex = "1.11" -serde = { version = "1.0", features = ["derive"] } teloxide = { version = "0.17", features = ["macros"] } tempfile = "3" thiserror = "2.0" diff --git a/src/commands.rs b/src/commands.rs index aeea625..d5bebf9 100644 --- a/src/commands.rs +++ b/src/commands.rs @@ -1,7 +1,6 @@ +use crate::comments::global_comments; use teloxide::{prelude::*, utils::command::BotCommands}; -use crate::comments::{Comments, global_comments}; - #[derive(BotCommands, Clone)] #[command(rename_rule = "lowercase")] pub enum Command { @@ -25,9 +24,8 @@ pub async fn answer(bot: Bot, msg: Message, cmd: Command) -> ResponseResult<()> .await? } Command::Curse => { - let comment = global_comments().map(Comments::build_caption); - bot.send_message(msg.chat.id, comment.unwrap_or_else(|| "To comment".into())) - .await? + let comment = global_comments().build_caption(); + bot.send_message(msg.chat.id, comment).await? } }; diff --git a/src/comments.rs b/src/comments.rs index 44e5605..c6bca6b 100644 --- a/src/comments.rs +++ b/src/comments.rs @@ -7,6 +7,8 @@ use std::{ }; use tokio::fs::read_to_string; +static GLOBAL_COMMENTS: OnceLock = OnceLock::new(); + const DISCLAIMER: &str = "(Roleplay — fictional messages for entertainment.)"; pub const TELEGRAM_CAPTION_LIMIT: usize = 4096; const FALLBACK_COMMENTS: &[&str] = &[ @@ -92,26 +94,28 @@ impl Comments { pub fn lines(&self) -> &[String] { &self.lines } + + /// Initialize the global comments (call once at startup). + /// + /// # Errors + /// + /// Returns `Error::Other` when the global is already initialized. + pub fn init(self) -> Result<()> { + GLOBAL_COMMENTS + .set(self) + .map_err(|_| Error::other("comments already initialized")) + } } -static GLOBAL_COMMENTS: OnceLock = OnceLock::new(); - -/// Initialize the global comments (call once at startup). +/// Get global comments (initialized by `Comments::init(self)`). /// -/// # Errors +/// # Panics /// -/// - Returns `Error::Other` when the global is already initialized. -pub fn init_global_comments(comments: Comments) -> Result<()> { - GLOBAL_COMMENTS - .set(comments) - .map_err(|_| Error::other("comments already initialized")) -} - -/// Get global comments (if initialized). +/// Panics if comments have not been initialized. #[inline] #[must_use] -pub fn global_comments() -> Option<&'static Comments> { - GLOBAL_COMMENTS.get() +pub fn global_comments() -> &'static Comments { + GLOBAL_COMMENTS.get().expect("comments not initialized") } impl Display for Comments { diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..11a32ac --- /dev/null +++ b/src/config.rs @@ -0,0 +1,111 @@ +use crate::error::{Error, Result}; +use std::{env, fmt::Debug, path::PathBuf, sync::OnceLock}; + +static GLOBAL_CONFIG: OnceLock = OnceLock::new(); + +#[derive(Debug, Clone, Default)] +pub struct Config { + pub youtube: YoutubeConfig, + pub instagram: InstagramConfig, + pub tiktok: TiktokConfig, + pub twitter: TwitterConfig, +} + +#[derive(Debug, Clone)] +pub struct YoutubeConfig { + pub cookies_path: Option, + pub postprocessor_args: String, +} + +#[derive(Debug, Clone, Default)] +pub struct InstagramConfig { + pub cookies_path: Option, +} + +#[derive(Debug, Clone, Default)] +pub struct TiktokConfig { + pub cookies_path: Option, +} + +#[derive(Debug, Clone, Default)] +pub struct TwitterConfig { + pub cookies_path: Option, +} + +impl Config { + /// Load configuration from environment variables. + #[must_use] + pub fn from_env() -> Self { + Self { + youtube: YoutubeConfig::from_env(), + instagram: InstagramConfig::from_env(), + tiktok: TiktokConfig::from_env(), + twitter: TwitterConfig::from_env(), + } + } + + /// Initialize the global config (call once at startup). + /// + /// # Errors + /// + /// Returns error if config is already initialized. + pub fn init(self) -> Result<()> { + GLOBAL_CONFIG + .set(self) + .map_err(|_| Error::other("config already initialized")) + } +} +/// Get global config (initialized by `Config::init(self)`). +#[must_use] +pub fn global_config() -> Config { + GLOBAL_CONFIG.get().cloned().unwrap_or_default() +} + +impl YoutubeConfig { + const DEFAULT_POSTPROCESSOR_ARGS: &'static str = "ffmpeg:-vf setsar=1 -c:v libx264 -crf 20 -preset ultrafast -c:a aac -b:a 128k -movflags +faststart"; + + fn from_env() -> Self { + Self { + cookies_path: get_path_from_env("YOUTUBE_SESSION_COOKIE_PATH"), + postprocessor_args: env::var("YOUTUBE_POSTPROCESSOR_ARGS") + .unwrap_or_else(|_| Self::DEFAULT_POSTPROCESSOR_ARGS.to_string()), + } + } +} + +impl InstagramConfig { + fn from_env() -> Self { + Self { + cookies_path: get_path_from_env("IG_SESSION_COOKIE_PATH"), + } + } +} + +impl TiktokConfig { + fn from_env() -> Self { + Self { + cookies_path: get_path_from_env("TIKTOK_SESSION_COOKIE_PATH"), + } + } +} + +impl TwitterConfig { + fn from_env() -> Self { + Self { + cookies_path: get_path_from_env("TWITTER_SESSION_COOKIE_PATH"), + } + } +} + +fn get_path_from_env(key: &str) -> Option { + env::var(key).ok().map(PathBuf::from) +} + +impl Default for YoutubeConfig { + fn default() -> Self { + Self { + cookies_path: None, + postprocessor_args: Self::DEFAULT_POSTPROCESSOR_ARGS.into(), + } + } +} diff --git a/src/download.rs b/src/download.rs index df41670..64a56bb 100644 --- a/src/download.rs +++ b/src/download.rs @@ -1,3 +1,4 @@ +use crate::config::global_config; use crate::{ error::{Error, Result}, utils::{ @@ -8,7 +9,6 @@ use crate::{ use futures::{StreamExt, stream}; use std::{ cmp::min, - env, ffi::OsStr, fs::{self, metadata}, path::{Path, PathBuf}, @@ -105,21 +105,12 @@ async fn run_command_in_tempdir(cmd: &str, args: &[&str]) -> Result) -> Result { - let base_args = ["-t", "mp4", "--extractor-args", "instagram:"]; - let mut args = base_args + let config = global_config(); + let args = ["-t", "mp4", "--extractor-args", "instagram:"] .iter() .map(ToString::to_string) - .collect::>(); - - if let Ok(cookies_path) = env::var("IG_SESSION_COOKIE_PATH") { - args.extend(["--cookies".into(), cookies_path]); - } - - args.push(url.into()); - - let args_ref = args.iter().map(String::as_ref).collect::>(); - - run_command_in_tempdir("yt-dlp", &args_ref).await + .collect(); + run_yt_dlp(args, config.instagram.cookies_path.as_ref(), &url.into()).await } /// Download a Tiktok URL with yt-dlp. @@ -129,21 +120,12 @@ pub async fn download_instagram(url: impl Into) -> Result) -> Result { - let base_args = ["-t", "mp4", "--extractor-args", "tiktok:"]; - let mut args = base_args + let config = global_config(); + let args = ["-t", "mp4", "--extractor-args", "tiktok:"] .iter() .map(ToString::to_string) - .collect::>(); - - if let Ok(cookies_path) = env::var("TIKTOK_SESSION_COOKIE_PATH") { - args.extend(["--cookies".into(), cookies_path]); - } - - args.push(url.into()); - - let args_ref = args.iter().map(String::as_ref).collect::>(); - - run_command_in_tempdir("yt-dlp", &args_ref).await + .collect(); + run_yt_dlp(args, config.tiktok.cookies_path.as_ref(), &url.into()).await } /// Download a Twitter URL with yt-dlp. @@ -153,8 +135,12 @@ pub async fn download_tiktok(url: impl Into) -> Result { /// - Propagates `run_command_in_tempdir` errors. #[cfg(feature = "twitter")] pub async fn download_twitter(url: impl Into) -> Result { - let args = ["-t", "mp4", "--extractor-args", "twitter:", &url.into()]; - run_command_in_tempdir("yt-dlp", &args).await + let config = global_config(); + let args = ["-t", "mp4", "--extractor-args", "twitter:"] + .iter() + .map(ToString::to_string) + .collect(); + run_yt_dlp(args, config.twitter.cookies_path.as_ref(), &url.into()).await } /// Download a URL with yt-dlp. @@ -164,7 +150,8 @@ pub async fn download_twitter(url: impl Into) -> Result /// - Propagates `run_command_in_tempdir` errors. #[cfg(feature = "youtube")] pub async fn download_youtube(url: impl Into) -> Result { - let base_args = [ + let config = global_config(); + let args = [ "--no-playlist", "-t", "mp4", @@ -172,21 +159,12 @@ pub async fn download_youtube(url: impl Into) -> Result "-o", "%(title)s.%(ext)s", "--postprocessor-args", - "ffmpeg:-vf setsar=1 -c:v libx264 -crf 20 -preset veryfast -c:a aac -b:a 128k -movflags +faststart", - ]; - let mut args = base_args - .iter() - .map(ToString::to_string) - .collect::>(); - - if let Ok(cookies_path) = env::var("YOUTUBE_SESSION_COOKIE_PATH") { - args.extend(["--cookies".into(), cookies_path]); - } - args.push(url.into()); - - let args_ref = args.iter().map(String::as_ref).collect::>(); - - run_command_in_tempdir("yt-dlp", &args_ref).await + &config.youtube.postprocessor_args, + ] + .iter() + .map(ToString::to_string) + .collect(); + run_yt_dlp(args, config.youtube.cookies_path.as_ref(), &url.into()).await } /// Post-process a `DownloadResult`. @@ -278,6 +256,20 @@ fn is_potential_media_file(path: &Path) -> bool { .any(|allowed| allowed.eq_ignore_ascii_case(&ext)) } +async fn run_yt_dlp( + mut args: Vec, + cookies_path: Option<&PathBuf>, + url: &str, +) -> Result { + if let Some(path) = cookies_path { + args.extend(["--cookies".to_string(), path.to_string_lossy().to_string()]); + } + args.push(url.to_string()); + + let args_ref = args.iter().map(String::as_ref).collect::>(); + run_command_in_tempdir("yt-dlp", &args_ref).await +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/lib.rs b/src/lib.rs index 5395148..4800029 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,6 @@ pub mod commands; pub mod comments; +pub mod config; pub mod download; pub mod error; pub mod handler; diff --git a/src/main.rs b/src/main.rs index ddb9970..2b69d56 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,7 +1,7 @@ use dotenv::dotenv; use teloxide::{prelude::*, respond}; use tg_relay_rs::{ - comments::{Comments, init_global_comments}, + comments::Comments, handler::{Handler, create_handlers}, telemetry::setup_logger, }; @@ -13,15 +13,15 @@ async fn main() -> color_eyre::Result<()> { color_eyre::install().expect("color-eyre install"); setup_logger(); - let comments = Comments::load_from_file("comments.txt") + Comments::load_from_file("comments.txt") .await .map_err(|e| { warn!("failed to laod comments.txt: {}; using dummy comments", e); e }) - .unwrap_or_else(|_| Comments::dummy()); - - init_global_comments(comments).expect("failed to initialize global comments"); + .unwrap_or_else(|_| Comments::dummy()) + .init() + .expect("failed to initialize comments"); let bot = Bot::from_env(); info!("bot starting"); diff --git a/src/utils.rs b/src/utils.rs index 1a5bd6d..b49728d 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,5 +1,5 @@ use crate::{ - comments::{Comments, global_comments}, + comments::global_comments, error::{Error, Result}, }; use capitalize::Capitalize; @@ -109,18 +109,13 @@ pub async fn send_media_from_path( path: PathBuf, kind: MediaKind, ) -> Result<()> { - let caption_opt = global_comments() - .map(Comments::build_caption) - .filter(|caption| !caption.is_empty()); - + let caption = global_comments().build_caption(); let input = InputFile::file(path); macro_rules! send_msg { ($request_expr:expr) => {{ let mut request = $request_expr; - if let Some(cap) = caption_opt { - request = request.caption(cap); - } + request = request.caption(caption); match request.await { Ok(message) => info!(message_id = message.id.to_string(), "{} sent", kind), Err(e) => {