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,39 +55,22 @@ 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) {
-            event!(Level::ERROR, "{}", e);
-        }
-    }
-}
-
-async fn download_fallback(url: &str, info: YtDlpInfo) -> Result<TmpFile, DownloadError> {
         let av = match info.best_av_format() {
             Some(av) => av,
             None => {
@@ -105,19 +90,16 @@ async fn download_fallback(url: &str, info: YtDlpInfo) -> Result<TmpFile, Downlo
         };
 
         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 info = YtDlp::load_info(url).await?;
         let vf = match info.best_video_format() {
             Some(vf) => vf,
-        None => return download_fallback(url, info).await,
+            None => return Err(DownloadError::NoFormatFound),
         };
         let af = match info.best_audio_format() {
             Some(af) => af,
-        None => return download_fallback(url, info).await,
+            None => return Err(DownloadError::NoFormatFound),
         };
 
         let video = YtDlp::download(url, &info, &vf).await?;
@@ -151,4 +133,56 @@ pub async fn download(url: &str) -> Result<TmpFile, DownloadError> {
             Ok(()) => Ok(output),
             Err(e) => Err(DownloadError::Message(e.to_string())),
         }
+    }
+
+    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)?;
+
+        Ok(YtDlp::download(url, info, original).await?)
+    }
+
+    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?;