From b0afa21511530974ec9269fc4de66b22054499f8 Mon Sep 17 00:00:00 2001 From: mykola2312 Date: Sat, 30 Mar 2024 09:03:03 +0200 Subject: [PATCH] implement separate download logic for youtube and tiktok, allowing for better quality decisions. DI via enums --- src/dl.rs | 196 +++++++++++++++++++++++++++-------------------- src/dl/yt_dlp.rs | 2 - 2 files changed, 115 insertions(+), 83 deletions(-) diff --git a/src/dl.rs b/src/dl.rs index 6be367f..4a61c42 100644 --- a/src/dl.rs +++ b/src/dl.rs @@ -1,12 +1,12 @@ use std::fmt; -use std::fs; use tracing::{event, Level}; +use crate::bot::sanitize::{extract_url, parse_url}; use crate::dl::ffmpeg::FFMpeg; use self::spawn::SpawnError; use self::tmpfile::{TmpFile, TmpFileError}; -use self::yt_dlp::{YtDlp, YtDlpError, YtDlpFormat, YtDlpInfo}; +use self::yt_dlp::{YtDlp, YtDlpError, YtDlpInfo}; pub mod ffmpeg; mod spawn; @@ -15,6 +15,7 @@ pub mod yt_dlp; pub enum DownloadError { Message(String), + NotAnURL, NoFormatFound, MakePathError, } @@ -44,6 +45,7 @@ impl fmt::Display for DownloadError { use DownloadError as DE; match &self { DE::Message(msg) => write!(f, "{}", msg), + DE::NotAnURL => write!(f, "no url or malformed url were provided"), DE::NoFormatFound => write!( f, "no best format found. you may want to specify one yourself" @@ -53,102 +55,134 @@ impl fmt::Display for DownloadError { } } -fn make_download_path( - info: &YtDlpInfo, - suffix: Option<&str>, - format: &YtDlpFormat, -) -> Result { - std::env::temp_dir() - .join(format!( - "{}_{}.{}", - info.id, - suffix.unwrap_or(""), - format.ext - )) - .into_os_string() - .into_string() - .map_err(|e| DownloadError::MakePathError) +enum Downloader { + Default, + YouTube, + TikTok, } -fn file_exists(path: &str) -> bool { - match fs::metadata(path) { - Ok(_) => true, - Err(_) => false, - } -} +const DEFAULT_DOWNLOADER: (&'static str, Downloader) = ("", Downloader::Default); +const DOWNLOADERS: [(&'static str, Downloader); 4] = [ + ("www.youtube.com", Downloader::YouTube), + ("youtu.be", Downloader::YouTube), + ("www.tiktok.com", Downloader::TikTok), + ("vm.tiktok.com", Downloader::TikTok), +]; -pub fn delete_if_exists(path: &str) { - if file_exists(path) { - if let Err(e) = fs::remove_file(path) { - event!(Level::ERROR, "{}", e); - } - } -} - -async fn download_fallback(url: &str, info: YtDlpInfo) -> Result { - let av = match info.best_av_format() { - Some(av) => av, - None => { - event!( - Level::WARN, - "no best format found for {}, reverting to default", - url - ); - match info.default_format() { - Some(format) => format, - None => { - event!(Level::ERROR, "no formats found for {}", url); - return Err(DownloadError::NoFormatFound); +impl Downloader { + async fn default_download(url: &str, info: &YtDlpInfo) -> Result { + let av = match info.best_av_format() { + Some(av) => av, + None => { + event!( + Level::WARN, + "no best format found for {}, reverting to default", + url + ); + match info.default_format() { + Some(format) => format, + None => { + event!(Level::ERROR, "no formats found for {}", url); + return Err(DownloadError::NoFormatFound); + } } } - } - }; + }; - Ok(YtDlp::download(url, &info, &av).await?) -} + Ok(YtDlp::download(url, &info, &av).await?) + } -pub async fn download(url: &str) -> Result { - event!(Level::INFO, "url {}", url); + async fn youtube_download(url: &str, info: &YtDlpInfo) -> Result { + let vf = match info.best_video_format() { + Some(vf) => vf, + None => return Err(DownloadError::NoFormatFound), + }; + let af = match info.best_audio_format() { + Some(af) => af, + None => return Err(DownloadError::NoFormatFound), + }; - let info = YtDlp::load_info(url).await?; - let vf = match info.best_video_format() { - Some(vf) => vf, - None => return download_fallback(url, info).await, - }; - let af = match info.best_audio_format() { - Some(af) => af, - None => return download_fallback(url, info).await, - }; + let video = YtDlp::download(url, &info, &vf).await?; + let audio = YtDlp::download(url, &info, &af).await?; - let video = YtDlp::download(url, &info, &vf).await?; - let audio = YtDlp::download(url, &info, &af).await?; + let abr = if let Some(abr) = af.abr { + FFMpeg::round_mp3_bitrate(abr) + } else { + event!( + Level::ERROR, + "somehow url {} audio format {} doesnt have abr", + url, + af.format_id + ); - let abr = if let Some(abr) = af.abr { - FFMpeg::round_mp3_bitrate(abr) - } else { + 192 + }; + + let output = TmpFile::new(format!("{}.{}", &info.id, &vf.ext).as_str())?; event!( - Level::ERROR, - "somehow url {} audio format {} doesnt have abr", + Level::INFO, + "for {} we joining video {} and audio {}", url, + vf.format_id, af.format_id ); - 192 - }; + let res = FFMpeg::join_video_audio(&video.path, &audio.path, abr, &output.path).await; - let output = TmpFile::new(format!("{}.{}", &info.id, &vf.ext).as_str())?; - event!( - Level::INFO, - "for {} we joining video {} and audio {}", - url, - vf.format_id, - af.format_id - ); + match res { + Ok(()) => Ok(output), + Err(e) => Err(DownloadError::Message(e.to_string())), + } + } - let res = FFMpeg::join_video_audio(&video.path, &audio.path, abr, &output.path).await; + async fn tiktok_download(url: &str, info: &YtDlpInfo) -> Result { + let original = info.formats + .iter() + .find(|f| f.format_id == "0") + .ok_or(DownloadError::NoFormatFound)?; - match res { - Ok(()) => Ok(output), - Err(e) => Err(DownloadError::Message(e.to_string())), + Ok(YtDlp::download(url, info, original).await?) + } + + pub async fn download(&self, url: &str, info: &YtDlpInfo) -> Result { + match self { + Downloader::Default => Self::default_download(url, info).await, + Downloader::YouTube => Self::youtube_download(url, info).await, + Downloader::TikTok => Self::tiktok_download(url, info).await + } } } + +impl fmt::Display for Downloader { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Downloader::Default => write!(f, "Default"), + Downloader::YouTube => write!(f, "YouTube"), + Downloader::TikTok => write!(f, "TikTok") + } + } +} + +pub async fn download(url: &str) -> Result { + let url = parse_url(extract_url(url).ok_or(DownloadError::NotAnURL)?) + .ok_or(DownloadError::NotAnURL)?; + let host_url = url.host_str().ok_or(DownloadError::NotAnURL)?; + + let downloader = &DOWNLOADERS + .iter() + .find(|f| f.0 == host_url) + .unwrap_or(&DEFAULT_DOWNLOADER).1; + event!(Level::INFO, "using {} downloader for {}", downloader, url); + + let info = YtDlp::load_info(url.as_str()).await?; + let output = match downloader.download(url.as_str(), &info).await { + Ok(output) => output, + Err(e) => { + event!(Level::ERROR, "downloader {} failed: {}. falling back to default downloader", downloader, e); + + DEFAULT_DOWNLOADER.1.download(url.as_str(), &info).await? + } + }; + + Ok(output) +} diff --git a/src/dl/yt_dlp.rs b/src/dl/yt_dlp.rs index 8f23ed1..b14fb83 100644 --- a/src/dl/yt_dlp.rs +++ b/src/dl/yt_dlp.rs @@ -241,8 +241,6 @@ impl fmt::Display for YtDlpError { } pub struct YtDlp {} - -// BUG: REAL ARGUMENT INJECTION! FIX ASAP impl YtDlp { pub async fn load_info(url: &str) -> Result { let output = spawn("python", &["-m", "yt_dlp", url, "-j"]).await?;