diff --git a/Cargo.lock b/Cargo.lock index 56e844a..595a16e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -254,6 +254,12 @@ dependencies = [ "regex", ] +[[package]] +name = "env_home" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f84e12ccf0a7ddc17a6c41c93326024c42920d7ee630d04950e6926645c0fe" + [[package]] name = "env_logger" version = "0.11.6" @@ -441,10 +447,12 @@ dependencies = [ "log", "m3u8-rs", "reqwest", + "shellwords", "tokio", "tower", "tower-http-client", "tower-reqwest", + "which", ] [[package]] @@ -773,6 +781,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + [[package]] name = "libc" version = "0.2.169" @@ -1265,6 +1279,16 @@ dependencies = [ "serde", ] +[[package]] +name = "shellwords" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89e515aa4699a88148ed5ef96413ceef0048ce95b43fbc955a33bde0a70fcae6" +dependencies = [ + "lazy_static", + "regex", +] + [[package]] name = "shlex" version = "1.3.0" @@ -1722,6 +1746,18 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "which" +version = "7.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb4a9e33648339dc1642b0e36e21b3385e6148e289226f657c809dee59df5028" +dependencies = [ + "either", + "env_home", + "rustix", + "winsafe", +] + [[package]] name = "windows-registry" version = "0.2.0" @@ -1834,6 +1870,12 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "winsafe" +version = "0.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d135d17ab770252ad95e9a872d365cf3090e3be864a34ab46f48555993efc904" + [[package]] name = "write16" version = "1.0.0" diff --git a/Cargo.toml b/Cargo.toml index 17397b1..97750aa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,7 +13,9 @@ http-body-util = "0.1" log = "0.4" m3u8-rs = "6.0" reqwest = "0.12" +shellwords = "1.1" tokio = { version = "1.42", features = [ "macros", "rt", "rt-multi-thread" ] } tower = { version = "0.5", features = [ "limit", "timeout" ] } tower-http-client = "0.4" tower-reqwest = "0.4" +which = "7.0" diff --git a/src/main.rs b/src/main.rs index d836aae..0f80121 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,13 +1,22 @@ mod download_error; mod client; -use std::time::Duration; +use std::{ + env::{current_dir, set_current_dir}, + ffi::OsStr, + io::ErrorKind, + path::{Path, PathBuf}, + process::Command, + time::Duration +}; use anyhow::anyhow; use clap::Parser; use env_logger::Env; use http::Uri; -use tokio::task::JoinSet; +use shellwords::escape; +use tokio::{fs::{create_dir, remove_dir_all}, task::JoinSet}; +use which::which; use log::{log, Level}; @@ -21,6 +30,8 @@ struct DownloadMessage { #[derive(Debug, Parser)] struct Args { + #[arg(short, long)] + name: PathBuf, #[arg(short, long)] url: String, #[arg(short, long)] @@ -44,7 +55,29 @@ async fn main() -> anyhow::Result<()> { let args = Args::parse(); log!(Level::Debug, "Arguments: {:?}", args); - log!(Level::Info, "Creating an HTTP client with Tower layers..."); + let name = args.name.as_path(); + + if name.file_name() != Some(OsStr::new(name)) { + return Err(anyhow!("Name must not contain a path")); + } + + if name.extension() != Some(OsStr::new("mp4")) { + return Err(anyhow!("Only filenames with .mp4 extension are allowed")); + } + + let basename = name.file_stem() + . ok_or(anyhow!("No valid filename given"))?; + + if let Err(error) = create_dir(basename).await { + if ErrorKind::AlreadyExists == error.kind() { + if ! Path::new(basename).is_dir() { + return Err(anyhow!("Unable to create directory at {:?}", basename)); + } + } else { + return Err(anyhow!("Unable to create directory at {:?}", basename)); + } + } + set_current_dir(basename)?; let concurrency = args.concurrency.unwrap_or(10); let timeout = args.timeout.unwrap_or(10); @@ -53,6 +86,9 @@ async fn main() -> anyhow::Result<()> { let m3u8_uri = Uri::try_from(&args.url)?; let m3u8_path_and_query = m3u8_uri.path_and_query() . ok_or(anyhow!("Problem path and query in m3u8 uri"))?; + + log!(Level::Info, "Creating an HTTP client with Tower layers..."); + let mut state = client::State::new(&m3u8_uri, concurrency, timeout)?; log!(Level::Info, "Get segments..."); @@ -91,5 +127,45 @@ async fn main() -> anyhow::Result<()> { } } + let outfile = current_dir()?; + let outfile = outfile.parent().ok_or(anyhow!("Can't get output dir"))?; + let mut outfile = outfile.canonicalize()?; + outfile.push(name); + let outfile = outfile.as_os_str(); + + let ffmpeg = which("ffmpeg")?; + let ffmpeg = ffmpeg.as_path().as_os_str(); + + let index = Path::new(m3u8_uri.path()).file_name() + . ok_or(anyhow!("unable to get index filename from url"))?; + let index = Path::new(index).canonicalize()?; + let index = index.as_os_str(); + + log!(Level::Debug, "ffmpeg: {:?}", ffmpeg); + log!(Level::Debug, "index: {:?}", index); + log!(Level::Debug, "outfile: {:?}", outfile); + + log!( Level::Info + , "execute: {} -allowed_extensions ALL -i {} -c copy {}" + , escape(ffmpeg.try_into()?) + , escape(index.try_into()?) + , escape(outfile.try_into()?) ); + let mut child = Command::new(ffmpeg) + . arg("-allowed_extensions") + . arg("ALL") + . arg("-i") + . arg(index) + . arg("-c") + . arg("copy") + . arg(outfile) + . spawn()?; + + let status = child.wait()?; + + log!(Level::Info, "ffmpeg status: {}", status); + + set_current_dir("..")?; + remove_dir_all(basename).await?; + Ok(()) }