Cargo.toml

@@ -1,22 +1,22 @@
 [package]
 name = "mk-dl-bot"
-version = "0.1.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,12 +13,12 @@ node {
 
     stage('Push') {
         docker.withRegistry('https://registry.hub.docker.com', 'a2aa5264-dce1-4054-8828-8db95e3c6c3c') {
-            app.push('v0.1.0')
+            app.push('v1.0.1')
         }
     }
 
     stage('Rollout') {
         sh('kubectl apply -f k8s/')
-        sh('kubectl rollout restart deployment bot')
+        sh('kubectl rollout restart deployment bot -n mk-dl-bot')
     }
 }

k8s/00_namespace.yml

@@ -0,0 +1,6 @@
+apiVersion: v1
+kind: Namespace
+metadata:
+  name: mk-dl-bot
+  labels:
+    name: mk-dl-bot

k8s/02_deployment.yml

@@ -5,6 +5,7 @@ metadata:
     app: mk-dl-bot
     service: bot
   name: bot
+  namespace: mk-dl-bot
 spec:
   replicas: 1
   selector:
@@ -18,8 +19,8 @@ spec:
         service: bot
     spec:
       containers:
-        - image: mykola2312/mk-dl-bot:v0.1.0
+        - image: mykola2312/mk-dl-bot:v1.0.1
           name: bot
           envFrom:
             - secretRef:
-                name: secret
+                name: secret

k8s/03_db-stateful-set.yml

@@ -2,6 +2,7 @@ apiVersion: apps/v1
 kind: StatefulSet
 metadata:
   name: db
+  namespace: mk-dl-bot
 spec:
   selector:
     matchLabels:
@@ -33,4 +34,4 @@ spec:
           - ReadWriteOnce
         resources:
           requests:
-            storage: 128Mi
+            storage: 128Mi

k8s/04_db-service.yml

@@ -2,6 +2,7 @@ apiVersion: v1
 kind: Service
 metadata:
   name: db-service
+  namespace: mk-dl-bot
 spec:
   selector:
     app: mk-dl-bot
@@ -10,4 +11,4 @@ spec:
   ports:
   - name: db
     protocol: TCP
-    port: 5432
+    port: 5432

k8s/05_tg-stateful-set.yml

@@ -0,0 +1,54 @@
+apiVersion: apps/v1
+kind: StatefulSet
+metadata:
+  name: tg
+  namespace: mk-dl-bot
+spec:
+  selector:
+    matchLabels:
+      app: mk-dl-bot
+      service: tg
+  serviceName: "tg"
+  replicas: 1
+  template:
+    metadata:
+      labels:
+        app: mk-dl-bot
+        service: tg
+    spec:
+      terminationGracePeriodSeconds: 10
+      containers:
+      - name: tg
+        image: mykola2312/telegram-bot-api:latest
+        ports:
+        - containerPort: 8081
+          name: tg
+        env:
+          - name: API_ID
+            valueFrom:
+              secretKeyRef:
+                name: secret
+                key: BOT_API_ID
+          - name: API_HASH
+            valueFrom:
+              secretKeyRef:
+                name: secret
+                key: BOT_API_HASH
+          - name: HTTP_PORT
+            value: "8081"
+        command: ["/app/telegram-bot-api"]
+        args: ["--api-id=$(API_ID)", "--api-hash=$(API_HASH)", "--local", "--http-port=$(HTTP_PORT)", "--dir=/var/lib/telegram-bot-api"]
+        volumeMounts:
+        - name: tg-data
+          mountPath: /var/lib/telegram-bot-api
+  volumeClaimTemplates:
+  - metadata:
+      name: tg-data
+    spec:
+      accessModes:
+      - ReadWriteOnce
+      resources:
+        requests:
+          storage: 5Gi
+---
+

k8s/06_tg-service.yml

@@ -0,0 +1,14 @@
+apiVersion: v1
+kind: Service
+metadata:
+  name: tg-service
+  namespace: mk-dl-bot
+spec:
+  selector:
+    app: mk-dl-bot
+    service: tg
+  type: ClusterIP
+  ports:
+  - name: tg
+    protocol: TCP
+    port: 8081

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,15 +1,18 @@
 use anyhow;
 use rust_i18n::t;
-use std::str;
+use std::str::{self, FromStr};
 use std::time::Duration;
 use teloxide::dispatching::{dialogue, dialogue::InMemStorage, UpdateHandler};
-use teloxide::types::{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::version::cmd_version;
 use super::types::*;
+use super::version::cmd_version;
 use crate::db::DbPool;
 use crate::util::{parse_env, unwrap_env};
 
@@ -24,7 +27,13 @@ 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"));
+    let bot = if cfg!(debug_assertions) {
+        Bot::new(unwrap_env("BOT_TOKEN"))
+    } else {
+        // we use telegram bot api server only in production
+        Bot::new(unwrap_env("BOT_TOKEN")).set_api_url(Url::from_str(&unwrap_env("BOT_API_URL"))?)
+    };
+
     let listener = Polling::builder(bot.clone())
         .timeout(Duration::from_secs(parse_env("POLLING_TIMEOUT")))
         .limit(parse_env("POLLING_LIMIT"))
@@ -109,7 +118,7 @@ enum Command {
 
     #[command(alias = "dl")]
     Download(String),
-    
+
     #[command(alias = "op")]
     OP,
 
@@ -128,8 +137,14 @@ enum Command {
     DeclineChat(String),
 }
 
+use crate::dl::spawn::spawn;
+
 async fn cmd_test(bot: Bot, msg: Message, _db: DbPool) -> HandlerResult {
-    bot.send_message(msg.chat.id, t!("test_response")).await?;
+    //bot.send_message(msg.chat.id, t!("test_response")).await?;
+    let output = spawn("python", &["-c", "import os; print(os.environ)"])
+        .await
+        .unwrap();
+    println!("{}", std::str::from_utf8(&output.stdout[0..4095]).unwrap());
 
     Ok(())
 }

src/bot/dl.rs

@@ -3,12 +3,11 @@ use teloxide::types::InputFile;
 use tracing::{event, Level};
 
 use super::types::HandlerResult;
-use crate::dl::delete_if_exists;
 use crate::dl::download;
 
 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 +15,8 @@ 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));
-    }
-
+    bot.send_video(msg.chat.id, InputFile::file(&output.path))
+        .await?;
     Ok(())
 }
 

src/bot/notify.rs

@@ -6,10 +6,11 @@ use crate::db::{DbPool, User};
 use super::types::HandlerResult;
 
 pub async fn notify_admins(bot: &Bot, db: &DbPool, message: String) -> HandlerResult {
-    let admins: Vec<User> =
-        sqlx::query_as(r#"SELECT * FROM "user" WHERE is_admin = true AND has_private_chat = true;"#)
-            .fetch_all(db)
-            .await?;
+    let admins: Vec<User> = sqlx::query_as(
+        r#"SELECT * FROM "user" WHERE is_admin = true AND has_private_chat = true;"#,
+    )
+    .fetch_all(db)
+    .await?;
 
     for admin in admins {
         let res = bot

src/bot/request.rs

@@ -23,22 +23,25 @@ pub async fn cmd_request(bot: Bot, msg: Message, text: String, db: DbPool) -> Ha
             reply_i18n_and_return!(bot, msg.chat.id, "already_can_download");
         }
 
-        let requests: i64 = sqlx::query(r#"SELECT COUNT(1) FROM "request" WHERE requested_by = $1;"#)
-            .bind(user.id)
-            .fetch_one(&db)
-            .await?
-            .get(0);
+        let requests: i64 =
+            sqlx::query(r#"SELECT COUNT(1) FROM "request" WHERE requested_by = $1;"#)
+                .bind(user.id)
+                .fetch_one(&db)
+                .await?
+                .get(0);
         if requests > 0 {
             reply_i18n_and_return!(bot, msg.chat.id, "already_has_requested");
         }
 
         // put the request
-        sqlx::query(r#"INSERT INTO "request" (requested_by,message,is_approved) VALUES ($1,$2,$3);"#)
-            .bind(user.id)
-            .bind(text)
-            .bind(false)
-            .execute(&db)
-            .await?;
+        sqlx::query(
+            r#"INSERT INTO "request" (requested_by,message,is_approved) VALUES ($1,$2,$3);"#,
+        )
+        .bind(user.id)
+        .bind(text)
+        .bind(false)
+        .execute(&db)
+        .await?;
         event!(Level::INFO, "added request for {}", user);
 
         // notify admins

src/bot/request_chat.rs

@@ -133,11 +133,13 @@ pub async fn cmd_approve_chat(bot: Bot, msg: Message, id: String, db: DbPool) ->
         };
 
         // approve request
-        sqlx::query(r#"UPDATE "request_chat" SET approved_by = $1, is_approved = true WHERE id = $2;"#)
-            .bind(user.id)
-            .bind(request.request_id)
-            .execute(&db)
-            .await?;
+        sqlx::query(
+            r#"UPDATE "request_chat" SET approved_by = $1, is_approved = true WHERE id = $2;"#,
+        )
+        .bind(user.id)
+        .bind(request.request_id)
+        .execute(&db)
+        .await?;
         event!(
             Level::INFO,
             "approved chat request {} by {} for {}",

src/bot/version.rs

@@ -8,4 +8,4 @@ pub async fn cmd_version(bot: Bot, msg: Message) -> HandlerResult {
     bot.send_message(msg.chat.id, VERSION).await?;
 
     Ok(())
-}
+}

src/dl.rs

@@ -1,16 +1,21 @@
 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 spawn;
+mod tmpfile;
 pub mod yt_dlp;
 
 pub enum DownloadError {
     Message(String),
+    NotAnURL,
     NoFormatFound,
     MakePathError,
 }
@@ -27,11 +32,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"
@@ -41,43 +55,123 @@ impl fmt::Display for DownloadError {
     }
 }
 
-fn make_download_path(info: &YtDlpInfo, format: &YtDlpFormat) -> Result<String, DownloadError> {
-    std::env::temp_dir()
-        .join(format!("{}.{}", info.id, 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 youtube_download(url: &str, info: &YtDlpInfo) -> Result<TmpFile, DownloadError> {
+        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 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)
+        } else {
+            event!(
+                Level::ERROR,
+                "somehow url {} audio format {} doesnt have abr",
+                url,
+                af.format_id
+            );
+
+            192
+        };
+
+        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
+        );
+
+        let res = FFMpeg::join_video_audio(&video.path, &audio.path, abr, &output.path).await;
+
+        match res {
+            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,
         }
     }
 }
 
-pub async fn download(url: &str) -> Result<String, DownloadError> {
-    event!(Level::INFO, "url {}", url);
+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"),
+        }
+    }
+}
 
-    let info = YtDlp::load_info(url).await?;
-    let av = match info.best_av_format() {
-        Some(av) => av,
-        None => return Err(DownloadError::NoFormatFound),
+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?
+        }
     };
 
-    let output_path = make_download_path(&info, &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)
+    Ok(output)
 }

src/dl/ffmpeg.rs

@@ -37,6 +37,38 @@ impl FFMpeg {
 
         Ok(())
     }
+
+    pub async fn join_video_audio(
+        video_path: &str,
+        audio_path: &str,
+        abr: u16,
+        output_path: &str,
+    ) -> Result<(), SpawnError> {
+        let abr = format!("{}k", abr);
+        let output = spawn(
+            "ffmpeg",
+            &[
+                "-i",
+                video_path,
+                "-i",
+                audio_path,
+                "-c",
+                "copy",
+                "-map",
+                "0:v:0",
+                "-map",
+                "1:a:0",
+                "-c:a",
+                "aac",
+                "-b:a",
+                &abr,
+                output_path,
+            ],
+        )
+        .await?;
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]

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,10 +20,23 @@ pub struct YtDlpFormat {
     pub abr: Option<f32>,
 }
 
+#[allow(unused)]
 struct VideoFormat<'a> {
     pub format: &'a YtDlpFormat,
+    pub format_note: &'a String,
     pub width: u16,
     pub height: u16,
+    pub vbr: f32,
+}
+
+impl<'a> VideoFormat<'a> {
+    pub fn is_mp4(&self) -> bool {
+        self.format.ext == "mp4"
+    }
+
+    pub fn is_premium(&self) -> bool {
+        self.format_note.contains("Premium")
+    }
 }
 
 struct AudioFormat<'a> {
@@ -79,6 +92,8 @@ pub struct YtDlpInfo {
 }
 
 impl YtDlpInfo {
+    const H_LIMIT: u16 = 1080;
+
     pub fn parse(json: &[u8]) -> Result<YtDlpInfo, serde_json::Error> {
         let mut info: YtDlpInfo = serde_json::from_slice(json)?;
         for format in &mut info.formats {
@@ -88,32 +103,6 @@ impl YtDlpInfo {
         Ok(info)
     }
 
-    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,
-                        width: f.width?,
-                        height: f.height?,
-                    })
-                } 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
@@ -134,6 +123,31 @@ impl YtDlpInfo {
             }
         }
     }
+
+    pub fn best_video_format(&self) -> Option<&YtDlpFormat> {
+        let format = self
+            .formats
+            .iter()
+            .filter_map(|f| {
+                Some(VideoFormat {
+                    format: f,
+                    format_note: f.format_note.as_ref()?,
+                    width: f.width?,
+                    height: f.height?,
+                    vbr: f.vbr?,
+                })
+            })
+            .filter(|f| f.height <= Self::H_LIMIT && f.is_mp4() && !f.is_premium())
+            .max_by_key(|f| OrderedFloat(f.vbr));
+
+        match format {
+            Some(vf) => Some(vf.format),
+            None => {
+                event!(Level::ERROR, "no video format for {}", self.id);
+                None
+            }
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -141,6 +155,8 @@ pub enum YtDlpError {
     SpawnError(SpawnError),
     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
@@ -154,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
@@ -167,40 +189,82 @@ impl fmt::Display for YtDlpError {
             YTE::SpawnError(e) => write!(f, "{}", e),
             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?;
 
-        Ok(YtDlpInfo::parse(&output.stdout)?)
+        let info = YtDlpInfo::parse(&output.stdout)?;
+        if info.formats.is_empty() {
+            return Err(YtDlpError::NoFormats);
+        }
+
+        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),
         }
     }
 }
@@ -210,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();
@@ -229,4 +283,14 @@ mod tests {
         let video = info.best_audio_format().unwrap();
         assert_eq!(video.format_id, "140");
     }
+
+    #[tokio::test]
+    async fn best_video_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_video_format().unwrap();
+        assert_eq!(video.format_id, "137");
+    }
 }

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() {