Cargo.toml

@@ -1,22 +1,22 @@
 [package]
 name = "mk-dl-bot"
-version = "0.2.0"
+version = "1.0.0"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-anyhow = "1.0.75"
+anyhow = "1.0.88"
 dotenv = "0.15.0"
-tokio = { version = "1.32.0", features = ["rt-multi-thread", "macros", "process"] }
-teloxide = { version = "0.12.2", git ="https://github.com/teloxide/teloxide", features = ["macros"] }
-sqlx = { version = "0.7.3", features = [ "runtime-tokio", "tls-native-tls", "postgres", "sqlx-postgres" ] }
-serde = { version = "1.0.196", features = ["derive"] }
-serde_json = "1.0.113"
-ordered-float = "4.2.0"
-regex = "1.10.3"
-url = "2.5.0"
+tokio = { version = "1.40.0", features = ["rt-multi-thread", "macros", "process"] }
+teloxide = { version = "0.13.0", features = ["macros"] }
+sqlx = { version = "0.7.4", features = [ "runtime-tokio", "tls-native-tls", "postgres", "sqlx-postgres" ] }
+serde = { version = "1.0.210", features = ["derive"] }
+serde_json = "1.0.128"
+ordered-float = "4.2.2"
+regex = "1.10.6"
+url = "2.5.2"
 tracing = { version = "0.1.40", features = ["async-await"] }
 tracing-appender = "0.2.3"
 tracing-subscriber = "0.3.18"
-rust-i18n = "3.0.1"
+rust-i18n = "3.1.2"

Jenkinsfile

@@ -13,7 +13,7 @@ node {
 
     stage('Push') {
         docker.withRegistry('https://registry.hub.docker.com', 'a2aa5264-dce1-4054-8828-8db95e3c6c3c') {
-            app.push('v0.1.1')
+            app.push('v1.0.1')
         }
     }
 

k8s/02_deployment.yml

@@ -19,7 +19,7 @@ spec:
         service: bot
     spec:
       containers:
-        - image: mykola2312/mk-dl-bot:v0.1.1
+        - image: mykola2312/mk-dl-bot:v1.0.1
           name: bot
           envFrom:
             - secretRef:

src/bot.rs

@@ -4,7 +4,6 @@ pub mod notify;
 pub mod op;
 pub mod request;
 pub mod request_chat;
-pub mod sanitize;
 pub mod start;
 pub mod types;
 pub mod version;

src/bot/bot.rs

@@ -1,12 +1,14 @@
 use anyhow;
 use rust_i18n::t;
-use url::Url;
 use std::str::{self, FromStr};
 use std::time::Duration;
 use teloxide::dispatching::{dialogue, dialogue::InMemStorage, UpdateHandler};
-use teloxide::types::{InputFile, InputMediaVideo, Me, MessageKind, MessageNewChatMembers, UpdateKind};
+use teloxide::types::{
+    InputFile, InputMediaVideo, Me, MessageKind, MessageNewChatMembers, UpdateKind,
+};
 use teloxide::{prelude::*, update_listeners::Polling, utils::command::BotCommands};
 use tracing::{event, Level};
+use url::Url;
 
 use super::start::handle_new_chat_member;
 use super::types::*;
@@ -25,8 +27,12 @@ use super::start::{cmd_start, handle_my_chat_member};
 pub async fn bot_main(db: DbPool) -> anyhow::Result<()> {
     event!(Level::INFO, "start");
 
-    let bot = Bot::new(unwrap_env("BOT_TOKEN"))
-        .set_api_url(Url::from_str(&unwrap_env("BOT_API_URL"))?);
+    let bot = if let Ok(api_url) = std::env::var("BOT_API_URL") {
+        Bot::new(unwrap_env("BOT_TOKEN"))
+            .set_api_url(Url::from_str(&api_url).unwrap())
+    } else {
+        Bot::new(unwrap_env("BOT_TOKEN"))
+    };
 
     let listener = Polling::builder(bot.clone())
         .timeout(Duration::from_secs(parse_env("POLLING_TIMEOUT")))
@@ -131,8 +137,18 @@ enum Command {
     DeclineChat(String),
 }
 
+use crate::dl::ffprobe::FFProbe;
+
 async fn cmd_test(bot: Bot, msg: Message, _db: DbPool) -> HandlerResult {
-    //bot.send_message(msg.chat.id, t!("test_response")).await?;
+    if cfg!(debug_assertions) {
+        if let Ok(probe) = FFProbe::probe("/home/mykola/Videos/test-video").await {
+            if let Some(vs) = probe.get_video_stream() {
+                dbg!(vs.get_video_resolution());
+            }
+        } else {
+            dbg!("failed");
+        }
+    }
 
     Ok(())
 }

src/bot/dl.rs

@@ -3,12 +3,13 @@ use teloxide::types::InputFile;
 use tracing::{event, Level};
 
 use super::types::HandlerResult;
-use crate::dl::delete_if_exists;
 use crate::dl::download;
 
+use crate::dl::ffprobe::FFProbe;
+
 async fn bot_download(bot: Bot, msg: Message, url: String) -> HandlerResult {
-    let output_path = match download(url.as_str()).await {
-        Ok(path) => path,
+    let output = match download(url.as_str()).await {
+        Ok(file) => file,
         Err(e) => {
             event!(Level::ERROR, "{}", e.to_string());
             bot.send_message(msg.chat.id, e.to_string()).await?;
@@ -16,14 +17,22 @@ async fn bot_download(bot: Bot, msg: Message, url: String) -> HandlerResult {
         }
     };
 
-    if let Err(e) = bot
-        .send_video(msg.chat.id, InputFile::file(&output_path))
-        .await
-    {
-        delete_if_exists(&output_path);
-        return Err(Box::new(e));
+    let mut video = bot.send_video(msg.chat.id, InputFile::file(&output.path));
+    // try getting video resolution
+    if let Ok(probe) = FFProbe::probe(&output.path).await {
+        if let Some(vs) = probe.get_video_stream() {
+            if let Some((width, height)) = vs.get_video_resolution() {
+                video.width = Some(width);
+                video.height = Some(height);
             }
 
+            // set video duration
+            video.duration = Some(vs.duration as u32);
+        }
+    }
+
+    video.await?;
+
     Ok(())
 }
 

src/dl.rs

@@ -1,18 +1,22 @@
 use std::fmt;
-use std::fs;
 use tracing::{event, Level};
 
 use crate::dl::ffmpeg::FFMpeg;
+use crate::security::sanitize::{extract_url, parse_url};
 
 use self::spawn::SpawnError;
-use self::yt_dlp::{YtDlp, YtDlpError, YtDlpFormat, YtDlpInfo};
+use self::tmpfile::{TmpFile, TmpFileError};
+use self::yt_dlp::{YtDlp, YtDlpError, YtDlpInfo};
 
 pub mod ffmpeg;
-mod spawn;
+pub mod ffprobe;
+pub mod spawn;
+mod tmpfile;
 pub mod yt_dlp;
 
 pub enum DownloadError {
     Message(String),
+    NotAnURL,
     NoFormatFound,
     MakePathError,
 }
@@ -29,11 +33,20 @@ impl From<YtDlpError> for DownloadError {
     }
 }
 
+impl From<TmpFileError> for DownloadError {
+    fn from(value: TmpFileError) -> Self {
+        match value {
+            TmpFileError::MakePathError => DownloadError::MakePathError,
+        }
+    }
+}
+
 impl fmt::Display for DownloadError {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         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"
@@ -43,92 +56,37 @@ impl fmt::Display for DownloadError {
     }
 }
 
-fn make_download_path(
-    info: &YtDlpInfo,
-    suffix: Option<&str>,
-    format: &YtDlpFormat,
-) -> 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)
+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),
+];
+
+impl Downloader {
+    async fn default_download(url: &str, info: &YtDlpInfo) -> Result<TmpFile, DownloadError> {
+        Ok(YtDlp::download(url, &info).await?)
     }
 
-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<String, DownloadError> {
-    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);
-                }
-            }
-        }
-    };
-
-    let output_path = make_download_path(&info, None, &av)?;
-    if let Err(e) = YtDlp::download(url, &av.format_id, output_path.as_str()).await {
-        delete_if_exists(&output_path);
-        return Err(DownloadError::Message(e.to_string()));
-    }
-
-    Ok(output_path)
-}
-
-pub async fn download(url: &str) -> Result<String, DownloadError> {
-    event!(Level::INFO, "url {}", url);
-
-    let info = YtDlp::load_info(url).await?;
+    async fn youtube_download(url: &str, info: &YtDlpInfo) -> Result<TmpFile, DownloadError> {
         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),
         };
 
-    // TODO: I should wrap those temp files in a impl Drop for defer deletion
-    let video_path = make_download_path(&info, Some("video"), &vf)?;
-    if let Err(e) = YtDlp::download(url, &vf.format_id, video_path.as_str()).await {
-        delete_if_exists(&video_path);
-        return Err(DownloadError::Message(e.to_string()));
-    }
-
-    let audio_path = make_download_path(&info, Some("audio"), &af)?;
-    if let Err(e) = YtDlp::download(url, &af.format_id, audio_path.as_str()).await {
-        delete_if_exists(&video_path);
-        delete_if_exists(&audio_path);
-        return Err(DownloadError::Message(e.to_string()));
-    }
+        let video = YtDlp::download_format(url, &info, &vf).await?;
+        let audio = YtDlp::download_format(url, &info, &af).await?;
 
         let abr = if let Some(abr) = af.abr {
             FFMpeg::round_mp3_bitrate(abr)
@@ -143,8 +101,7 @@ pub async fn download(url: &str) -> Result<String, DownloadError> {
             192
         };
 
-    let output_path = make_download_path(&info, None, &vf)?;
-
+        let output = TmpFile::new(format!("{}.{}", &info.id, &vf.ext).as_str())?;
         event!(
             Level::INFO,
             "for {} we joining video {} and audio {}",
@@ -153,18 +110,69 @@ pub async fn download(url: &str) -> Result<String, DownloadError> {
             af.format_id
         );
 
-    let res = FFMpeg::join_video_audio(
-        video_path.as_str(),
-        audio_path.as_str(),
-        abr,
-        output_path.as_str(),
-    )
-    .await;
-    delete_if_exists(&video_path);
-    delete_if_exists(&audio_path);
+        let res = FFMpeg::join_video_audio(&video.path, &audio.path, abr, &output.path).await;
 
         match res {
-        Ok(()) => Ok(output_path),
+            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_format(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/ffmpeg.rs

@@ -69,6 +69,9 @@ impl FFMpeg {
 
         Ok(())
     }
+
+    // TODO: implement function to generate thumbnails for videos
+    // ffmpeg -i test-video -vf "select=eq(n\,0)" -vf scale=320:-2 -q:v 3 test-video-thumbnail.jpg
 }
 
 #[cfg(test)]

src/dl/ffprobe.rs

@@ -0,0 +1,101 @@
+use super::spawn::{spawn, SpawnError};
+use serde::{de, de::Error, Deserialize, Deserializer};
+use std::fmt;
+
+fn duration_from_str<'de, D>(deserializer: D) -> Result<f64, D::Error>
+    where D: Deserializer<'de>
+{
+    let s = String::deserialize(deserializer)?;
+    
+    Ok(str::parse(&s).map_err(de::Error::custom)?)
+}
+
+#[derive(Deserialize, Debug)]
+pub struct FFProbeStream {
+    pub index: u32,
+    pub codec_name: String,
+    pub width: Option<u32>,
+    pub height: Option<u32>,
+    pub coded_width: Option<u32>,
+    pub coded_height: Option<u32>,
+    pub time_base: String,
+    pub duration_ts: u64,
+    
+    #[serde(deserialize_with = "duration_from_str")]
+    pub duration: f64
+}
+
+impl FFProbeStream {
+    pub fn is_video_stream(&self) -> bool {
+        self.width.is_some() && self.height.is_some()
+    }
+
+    pub fn get_video_resolution(&self) -> Option<(u32, u32)> {
+        if self.is_video_stream() {
+            return Some((self.width.unwrap(), self.height.unwrap()))
+        }
+
+        return None
+    }
+}
+
+#[derive(Deserialize, Debug)]
+pub struct FFProbeOutput {
+    pub streams: Vec<FFProbeStream>
+}
+
+impl FFProbeOutput {
+    pub fn parse(json: &[u8]) -> Result<FFProbeOutput, serde_json::Error> {
+        let output: FFProbeOutput = serde_json::from_slice(json)?;
+
+        Ok(output)
+    }
+
+    pub fn get_video_stream(&self) -> Option<&FFProbeStream> {
+        self.streams.iter().find(|&s| s.is_video_stream())
+    }
+}
+
+#[derive(Debug)]
+pub enum FFProbeError {
+    SpawnError(SpawnError),
+    JsonError
+}
+
+impl From<SpawnError> for FFProbeError {
+    fn from(value: SpawnError) -> Self {
+        Self::SpawnError(value)
+    }
+}
+
+impl From<serde_json::Error> for FFProbeError {
+    fn from(value: serde_json::Error) -> Self {
+        Self::JsonError
+    }
+}
+
+impl fmt::Display for FFProbeError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        use FFProbeError as FFPE;
+        match (self) {
+            FFPE::SpawnError(e) => write!(f, "{}", e),
+            FFPE::JsonError => write!(f, "ffprobe json error")
+        }
+    }
+}
+
+pub struct FFProbe {}
+impl FFProbe {
+    pub async fn probe(input_path: &str) -> Result<FFProbeOutput, FFProbeError> {
+        let output = spawn("ffprobe", &[
+            "-v", "quiet",
+            "-print_format", "json",
+            "-show_streams",
+            input_path
+        ]).await?;
+
+        let output = FFProbeOutput::parse(&output.stdout)?;
+
+        Ok(output)
+    }
+}

src/dl/spawn.rs

@@ -1,12 +1,18 @@
 use core::fmt;
 use std::process::Output;
 use std::str::Utf8Error;
+use std::{fs::OpenOptions, process::Stdio};
+use tokio::io::AsyncReadExt;
 use tokio::process::Command;
 use tracing::{event, Level};
 
+use super::tmpfile::TmpFile;
+
 #[derive(Debug)]
 pub enum SpawnError {
     CommandError(std::io::Error),
+    NoStdErr,
+    PipeError(std::io::Error),
     UtfError(Utf8Error),
     ErrorMessage(String),
 }
@@ -28,6 +34,8 @@ impl fmt::Display for SpawnError {
         use SpawnError as FE;
         match self {
             FE::CommandError(e) => write!(f, "Command::new - {}", e),
+            FE::NoStdErr => write!(f, "spawned process has closed stderr!"),
+            FE::PipeError(e) => write!(f, "pipe error - {}", e),
             FE::UtfError(_) => write!(f, "Error while decoding UTF8"),
             FE::ErrorMessage(msg) => write!(f, "ffmpeg error - {}", msg),
         }
@@ -42,7 +50,14 @@ pub async fn spawn(program: &str, args: &[&str]) -> Result<Output, SpawnError> {
         event!(Level::INFO, "{} {}", program, cmd_args);
     }
 
-    let output = Command::new(program).args(args).output().await?;
+    // TODO: python can't run without environment variables.
+    // TODO: I need to figure out which one are required for python to work
+    let output = Command::new(program)
+        .args(args)
+        .env_clear()
+        .env("PYTHONPATH", std::env::var("PYTHONPATH").unwrap())
+        .output()
+        .await?;
 
     if !output.status.success() {
         let message = std::str::from_utf8(&output.stderr)?;
@@ -51,3 +66,71 @@ pub async fn spawn(program: &str, args: &[&str]) -> Result<Output, SpawnError> {
 
     Ok(output)
 }
+
+pub async fn spawn_pipe(
+    program: &str,
+    args: &[&str],
+    output_file: &TmpFile,
+) -> Result<(), SpawnError> {
+    {
+        let cmd_args = args.join(" ");
+        event!(Level::INFO, "{} {}", program, cmd_args);
+    }
+
+    let output_file = OpenOptions::new()
+        .create(true)
+        .write(true)
+        .open(&output_file.path)
+        .map_err(|e| SpawnError::PipeError(e))?;
+
+    let mut process = Command::new(program)
+        .args(args)
+        .stdout(output_file)
+        .stderr(Stdio::piped())
+        .spawn()?;
+    let mut stderr = process.stderr.take().ok_or(SpawnError::NoStdErr)?;
+
+    let result = process.wait().await?;
+
+    if !result.success() {
+        let mut data: Vec<u8> = Vec::new();
+        stderr.read_to_end(&mut data).await?;
+
+        let message = std::str::from_utf8(&data)?;
+        return Err(SpawnError::ErrorMessage(message.to_string()));
+    }
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::dl::spawn::{spawn_pipe, SpawnError};
+    use crate::dl::tmpfile::TmpFile;
+    use std::fs;
+
+    #[tokio::test]
+    async fn test_spawn_pipe() {
+        let stdout_file = TmpFile::new("stdout.test").unwrap();
+        let result = spawn_pipe(
+            "python",
+            &[
+                "-c",
+                "import sys; print('stdout test', end=''); print('stderr test', file=sys.stderr, end=''); sys.exit(1)",
+            ],
+            &stdout_file,
+        )
+        .await;
+
+        let stdout = fs::read_to_string(&stdout_file.path).unwrap();
+        assert_eq!("stdout test", stdout);
+
+        assert_eq!(true, result.is_err());
+        if let Err(e) = result {
+            match e {
+                SpawnError::ErrorMessage(msg) => assert_eq!("stderr test", msg),
+                _ => panic!("SpawnError is not ErrorMessage!"),
+            }
+        }
+    }
+}

src/dl/tmpfile.rs

@@ -0,0 +1,45 @@
+use std::fs;
+use tracing::{event, Level};
+
+#[derive(Debug)]
+pub enum TmpFileError {
+    MakePathError,
+}
+
+pub struct TmpFile {
+    pub path: String,
+}
+
+impl TmpFile {
+    pub fn new(filename: &str) -> Result<Self, TmpFileError> {
+        let path = std::env::temp_dir()
+            .join(filename)
+            .into_os_string()
+            .into_string()
+            .map_err(|_| TmpFileError::MakePathError)?;
+
+        Ok(Self { path })
+    }
+
+    pub fn exists(&self) -> bool {
+        match fs::metadata(&self.path) {
+            Ok(_) => true,
+            Err(_) => false,
+        }
+    }
+
+    fn delete_if_exists(&self) {
+        if self.exists() {
+            if let Err(e) = fs::remove_file(&self.path) {
+                event!(Level::ERROR, "{}", e);
+            }
+        }
+        event!(Level::INFO, "deleted {}", self.path);
+    }
+}
+
+impl Drop for TmpFile {
+    fn drop(&mut self) {
+        self.delete_if_exists();
+    }
+}

src/dl/yt_dlp.rs

@@ -1,9 +1,9 @@
-use super::spawn::{spawn, SpawnError};
+use super::spawn::{spawn, spawn_pipe, SpawnError};
+use super::tmpfile::{TmpFile, TmpFileError};
 use core::fmt;
 use ordered_float::OrderedFloat;
 use serde::Deserialize;
 use serde_json;
-use std::fs;
 use tracing::{event, Level};
 
 #[derive(Deserialize, Debug)]
@@ -20,6 +20,7 @@ pub struct YtDlpFormat {
     pub abr: Option<f32>,
 }
 
+#[allow(unused)]
 struct VideoFormat<'a> {
     pub format: &'a YtDlpFormat,
     pub format_note: &'a String,
@@ -102,50 +103,6 @@ impl YtDlpInfo {
         Ok(info)
     }
 
-    pub fn default_format(&self) -> Option<&YtDlpFormat> {
-        match self
-            .formats
-            .iter()
-            .filter(|f| f.height.is_some_and(|h| h <= Self::H_LIMIT))
-            .last()
-        {
-            Some(format) => Some(format),
-            None => self.formats.last(),
-        }
-    }
-
-    #[deprecated(
-        since = "0.1.1",
-        note = "for YouTube download audio and video separately"
-    )]
-    pub fn best_av_format(&self) -> Option<&YtDlpFormat> {
-        let format = self
-            .formats
-            .iter()
-            .filter_map(|f| {
-                if f.vcodec.is_some() && f.acodec.is_some() {
-                    Some(VideoFormat {
-                        format: &f,
-                        format_note: f.format_note.as_ref()?,
-                        width: f.width?,
-                        height: f.height?,
-                        vbr: f.vbr?,
-                    })
-                } else {
-                    None
-                }
-            })
-            .max_by_key(|f| (f.width, f.height));
-
-        match format {
-            Some(vf) => Some(vf.format),
-            None => {
-                event!(Level::ERROR, "no av format for {}", self.id);
-                None
-            }
-        }
-    }
-
     pub fn best_audio_format(&self) -> Option<&YtDlpFormat> {
         let format = self
             .formats
@@ -199,6 +156,7 @@ pub enum YtDlpError {
     ErrorMessage(String), // keep it separate type if we ever plan to parse yt-dlp errors
     JsonError,
     NoFormats,
+    MakePathError,
     NoFilePresent,
 }
 // ^(?:ERROR: \[.*\] \S* )(.*$) - regex for matching yt-dlp's youtube errors
@@ -212,6 +170,12 @@ impl From<SpawnError> for YtDlpError {
     }
 }
 
+impl From<TmpFileError> for YtDlpError {
+    fn from(_value: TmpFileError) -> Self {
+        Self::MakePathError
+    }
+}
+
 impl From<serde_json::Error> for YtDlpError {
     fn from(_value: serde_json::Error) -> Self {
         Self::JsonError
@@ -226,17 +190,16 @@ impl fmt::Display for YtDlpError {
             YTE::ErrorMessage(msg) => write!(f, "yt-dlp error - {}", msg),
             YTE::JsonError => write!(f, "json parsing error"),
             YTE::NoFormats => write!(f, "no formats were parsed"),
+            YTE::MakePathError => write!(f, "make path error"),
             YTE::NoFilePresent => write!(f, "downloaded file doesn't exists"),
         }
     }
 }
 
 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?;
+        let output = spawn("python", &["-m", "yt_dlp", url, "-j", "--no-exec"]).await?;
 
         let info = YtDlpInfo::parse(&output.stdout)?;
         if info.formats.is_empty() {
@@ -246,25 +209,62 @@ impl YtDlp {
         Ok(info)
     }
 
-    pub async fn download(url: &str, format_id: &str, output_path: &str) -> Result<(), YtDlpError> {
-        spawn(
+    pub async fn download(url: &str, info: &YtDlpInfo) -> Result<TmpFile, YtDlpError> {
+        let file = TmpFile::new(&info.id)?;
+
+        // since yt-dlp tend to randomly choose filename we can't rely on it,
+        // and instead output to stdout and then pipe to our file
+        // that way we can avoid bugs related to filename confusion
+        let output = spawn_pipe(
+            "python",
+            &[
+                "-m",
+                "yt_dlp",
+                url,
+                "-o",
+                "-",
+                "--force-overwrites",
+                "--no-exec",
+            ],
+            &file,
+        )
+        .await?;
+        dbg!(output);
+
+        match file.exists() {
+            true => Ok(file),
+            false => Err(YtDlpError::NoFilePresent),
+        }
+    }
+
+    pub async fn download_format(
+        url: &str,
+        info: &YtDlpInfo,
+        format: &YtDlpFormat,
+    ) -> Result<TmpFile, YtDlpError> {
+        let file =
+            TmpFile::new(format!("{}_{}.{}", info.id, format.format_id, format.ext).as_str())?;
+
+        spawn_pipe(
             "python",
             &[
                 "-m",
                 "yt_dlp",
                 url,
                 "-f",
-                format_id,
+                &format.format_id,
                 "-o",
-                output_path,
+                "-",
                 "--force-overwrites",
+                "--no-exec",
             ],
+            &file,
         )
         .await?;
 
-        match fs::metadata(output_path) {
-            Ok(_) => Ok(()),
-            Err(_) => Err(YtDlpError::NoFilePresent),
+        match file.exists() {
+            true => Ok(file),
+            false => Err(YtDlpError::NoFilePresent),
         }
     }
 }
@@ -274,16 +274,6 @@ mod tests {
     use super::YtDlp;
     use std::env;
 
-    #[tokio::test]
-    async fn best_av_format() {
-        dotenv::from_filename(".env.test").unwrap();
-        let info = YtDlp::load_info(env::var("TEST_URL").unwrap().as_str())
-            .await
-            .unwrap();
-        let video = info.best_av_format().unwrap();
-        assert_eq!(video.format_id, "22");
-    }
-
     #[tokio::test]
     async fn best_audio_format() {
         dotenv::from_filename(".env.test").unwrap();

src/main.rs

@@ -13,6 +13,8 @@ use log::log_init;
 mod db;
 use db::db_init;
 
+mod security;
+
 rust_i18n::i18n!("locales");
 
 #[tokio::main]

src/security.rs

@@ -0,0 +1 @@
+pub mod sanitize;

src/security/sanitize.rs

@@ -3,7 +3,7 @@ use url::Url;
 
 // https://stackoverflow.com/questions/6038061/regular-expression-to-find-urls-within-a-string
 const RE_URL: &str =
-    r"(http|ftp|https):\/\/([\w_-]+(?:(?:\.[\w_-]+)+))([\w.,@?^=%&:\/~+#-]*[\w@?^=%&\/~+#-])";
+    r"(http|https):\/\/([\w_-]+(?:(?:\.[\w_-]+)+))([\w.,@?^=%&:\/~+#-]*[\w@?^=%&\/~+#-])";
 
 pub fn extract_url(text: &str) -> Option<&str> {
     let re = Regex::new(RE_URL).unwrap();
@@ -19,7 +19,7 @@ pub fn parse_url(url: &str) -> Option<Url> {
 
 #[cfg(test)]
 mod tests {
-    use crate::bot::sanitize::{extract_url, parse_url};
+    use crate::security::sanitize::{extract_url, parse_url};
 
     #[test]
     fn test_extract_url() {