implement separate download logic for youtube and tiktok, allowing for better quality decisions. DI via enums

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(
+enum Downloader {
-    info: &YtDlpInfo,
+    Default,
-    suffix: Option<&str>,
+    YouTube,
-    format: &YtDlpFormat,
+    TikTok,
-) -> Result<String, DownloadError> {
-    std::env::temp_dir()
-        .join(format!(
-            "{}_{}.{}",
-            info.id,
-            suffix.unwrap_or(""),
-            format.ext
-        ))
-        .into_os_string()
-        .into_string()
-        .map_err(|e| DownloadError::MakePathError)
 }
 
-fn file_exists(path: &str) -> bool {
+const DEFAULT_DOWNLOADER: (&'static str, Downloader) = ("", Downloader::Default);
-    match fs::metadata(path) {
+const DOWNLOADERS: [(&'static str, Downloader); 4] = [
-        Ok(_) => true,
+    ("www.youtube.com", Downloader::YouTube),
-        Err(_) => false,
+    ("youtu.be", Downloader::YouTube),
-    }
+    ("www.tiktok.com", Downloader::TikTok),
-}
+    ("vm.tiktok.com", Downloader::TikTok),
+];
 
-pub fn delete_if_exists(path: &str) {
+impl Downloader {
-    if file_exists(path) {
+    async fn default_download(url: &str, info: &YtDlpInfo) -> Result<TmpFile, DownloadError> {
-        if let Err(e) = fs::remove_file(path) {
+        let av = match info.best_av_format() {
-            event!(Level::ERROR, "{}", e);
+            Some(av) => av,
-        }
+            None => {
-    }
+                event!(
-}
+                    Level::WARN,
-
+                    "no best format found for {}, reverting to default",
-async fn download_fallback(url: &str, info: YtDlpInfo) -> Result<TmpFile, DownloadError> {
+                    url
-    let av = match info.best_av_format() {
+                );
-        Some(av) => av,
+                match info.default_format() {
-        None => {
+                    Some(format) => format,
-            event!(
+                    None => {
-                Level::WARN,
+                        event!(Level::ERROR, "no formats found for {}", url);
-                "no best format found for {}, reverting to default",
+                        return Err(DownloadError::NoFormatFound);
-                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<TmpFile, DownloadError> {
+    async fn youtube_download(url: &str, info: &YtDlpInfo) -> Result<TmpFile, DownloadError> {
-    event!(Level::INFO, "url {}", url);
+        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 video = YtDlp::download(url, &info, &vf).await?;
-    let vf = match info.best_video_format() {
+        let audio = YtDlp::download(url, &info, &af).await?;
-        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 abr = if let Some(abr) = af.abr {
-    let audio = YtDlp::download(url, &info, &af).await?;
+            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 {
+            192
-        FFMpeg::round_mp3_bitrate(abr)
+        };
-    } else {
+
+        let output = TmpFile::new(format!("{}.{}", &info.id, &vf.ext).as_str())?;
         event!(
-            Level::ERROR,
+            Level::INFO,
-            "somehow url {} audio format {} doesnt have abr",
+            "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())?;
+        match res {
-    event!(
+            Ok(()) => Ok(output),
-        Level::INFO,
+            Err(e) => Err(DownloadError::Message(e.to_string())),
-        "for {} we joining video {} and audio {}",
+        }
-        url,
+    }
-        vf.format_id,
-        af.format_id
-    );
 
-    let res = FFMpeg::join_video_audio(&video.path, &audio.path, abr, &output.path).await;
+    async fn tiktok_download(url: &str, info: &YtDlpInfo) -> Result<TmpFile, DownloadError> {
+        let original = info.formats
+            .iter()
+            .find(|f| f.format_id == "0")
+            .ok_or(DownloadError::NoFormatFound)?;
 
-    match res {
+        Ok(YtDlp::download(url, info, original).await?)
-        Ok(()) => Ok(output),
+    }
-        Err(e) => Err(DownloadError::Message(e.to_string())),
+
+    pub async fn download(&self, url: &str, info: &YtDlpInfo) -> Result<TmpFile, DownloadError> {
+        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<TmpFile, DownloadError> {
+    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)
+}

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<YtDlpInfo, YtDlpError> {
         let output = spawn("python", &["-m", "yt_dlp", url, "-j"]).await?;