From b82b2ace4c7bc9aea18b560d399318290e598a3b Mon Sep 17 00:00:00 2001
From: mykola2312 <nikolay_tihonov@i.ua>
Date: Sun, 31 Mar 2024 14:55:14 +0300
Subject: [PATCH] just let yt-dlp choose default format, however we need to
 determine file extension or telegram may reject file

---
 Cargo.lock       | 31 +++++++++++++++++++++++++++++++
 Cargo.toml       |  1 +
 src/dl.rs        | 30 ++++++------------------------
 src/dl/yt_dlp.rs | 24 +++++++++++++++++++++++-
 src/main.rs      |  2 ++
 5 files changed, 63 insertions(+), 25 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 8e153f3..3cf1f91 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -489,6 +489,12 @@ version = "1.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8fcfdc7a0362c9f4444381a9e697c79d435fe65b52a37466fc2c1184cee9edc6"
 
+[[package]]
+name = "fixedbitset"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80"
+
 [[package]]
 name = "flume"
 version = "0.11.0"
@@ -1095,6 +1101,7 @@ dependencies = [
  "tracing",
  "tracing-appender",
  "tracing-subscriber",
+ "tree_magic_mini",
  "url",
 ]
 
@@ -1337,6 +1344,16 @@ version = "2.3.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e"
 
+[[package]]
+name = "petgraph"
+version = "0.6.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e1d3afd2628e69da2be385eb6f2fd57c8ac7977ceeff6dc166ff1657b0e386a9"
+dependencies = [
+ "fixedbitset",
+ "indexmap 2.2.3",
+]
+
 [[package]]
 name = "pin-project"
 version = "1.1.3"
@@ -2543,6 +2560,20 @@ dependencies = [
  "tracing-log",
 ]
 
+[[package]]
+name = "tree_magic_mini"
+version = "3.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "77ee137597cdb361b55a4746983e4ac1b35ab6024396a419944ad473bb915265"
+dependencies = [
+ "fnv",
+ "home",
+ "memchr",
+ "nom",
+ "once_cell",
+ "petgraph",
+]
+
 [[package]]
 name = "triomphe"
 version = "0.1.11"
diff --git a/Cargo.toml b/Cargo.toml
index 0b611d9..8505a08 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -20,3 +20,4 @@ tracing = { version = "0.1.40", features = ["async-await"] }
 tracing-appender = "0.2.3"
 tracing-subscriber = "0.3.18"
 rust-i18n = "3.0.1"
+tree_magic_mini = "3.1.4"
diff --git a/src/dl.rs b/src/dl.rs
index 4a61c42..ede2a7c 100644
--- a/src/dl.rs
+++ b/src/dl.rs
@@ -1,7 +1,7 @@
 use std::fmt;
 use tracing::{event, Level};
 
-use crate::bot::sanitize::{extract_url, parse_url};
+use crate::security::sanitize::{extract_url, parse_url};
 use crate::dl::ffmpeg::FFMpeg;
 
 use self::spawn::SpawnError;
@@ -9,7 +9,7 @@ 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;
 
@@ -71,25 +71,7 @@ const DOWNLOADERS: [(&'static str, Downloader); 4] = [
 
 impl Downloader {
     async fn default_download(url: &str, info: &YtDlpInfo) -> Result<TmpFile, 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);
-                    }
-                }
-            }
-        };
-
-        Ok(YtDlp::download(url, &info, &av).await?)
+        Ok(YtDlp::download(url, &info).await?)
     }
 
     async fn youtube_download(url: &str, info: &YtDlpInfo) -> Result<TmpFile, DownloadError> {
@@ -102,8 +84,8 @@ impl Downloader {
             None => return Err(DownloadError::NoFormatFound),
         };
 
-        let video = YtDlp::download(url, &info, &vf).await?;
-        let audio = YtDlp::download(url, &info, &af).await?;
+        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)
@@ -141,7 +123,7 @@ impl Downloader {
             .find(|f| f.format_id == "0")
             .ok_or(DownloadError::NoFormatFound)?;
 
-        Ok(YtDlp::download(url, info, original).await?)
+        Ok(YtDlp::download_format(url, info, original).await?)
     }
 
     pub async fn download(&self, url: &str, info: &YtDlpInfo) -> Result<TmpFile, DownloadError> {
diff --git a/src/dl/yt_dlp.rs b/src/dl/yt_dlp.rs
index b14fb83..968e25f 100644
--- a/src/dl/yt_dlp.rs
+++ b/src/dl/yt_dlp.rs
@@ -253,7 +253,29 @@ impl YtDlp {
         Ok(info)
     }
 
-    pub async fn download(
+    pub async fn download(url: &str, info: &YtDlpInfo) -> Result<TmpFile, YtDlpError> {
+        let file = TmpFile::new(&info.id)?;
+
+        spawn(
+            "python",
+            &[
+                "-m",
+                "yt_dlp",
+                url,
+                "-o",
+                &file.path,
+                "--force-overwrites",
+            ],
+        )
+        .await?;
+
+        match file.exists() {
+            true => Ok(file),
+            false => Err(YtDlpError::NoFilePresent),
+        }
+    }
+
+    pub async fn download_format(
         url: &str,
         info: &YtDlpInfo,
         format: &YtDlpFormat,
diff --git a/src/main.rs b/src/main.rs
index 54fe696..fa231ae 100644
--- a/src/main.rs
+++ b/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]