diff --git a/src/main.rs b/src/main.rs index 0c66705..3e5201d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,16 +2,15 @@ mod json; mod args; mod strings; mod process; + use process as proc; -use std::hash::{Hash, Hasher}; -use std::io::Write; +use args as tcargs; +use args::Args; + use std::pin::Pin; use std::sync::Arc; use tokio::sync::RwLock; -use args as tcargs; -use args::Args; - const VERSION: &str = include_str!("version"); #[tokio::main] @@ -97,7 +96,7 @@ async fn main() -> std::process::ExitCode { let stats = match json::parse(&ver) { Ok(j) => j, Err(e) => { - eprintln!("{}", print_cobalt_error(e, ver)); + eprintln!("{}", proc::print_json_error(e, ver)); return std::process::ExitCode::FAILURE; } }; @@ -112,7 +111,7 @@ async fn main() -> std::process::ExitCode { } async fn execute_get_media(args: Args, bulk: u16, debug: bool) -> bool { - let json = cobalt_args(&args); + let json = proc::cobalt_args(&args); let download_url: &str = args.c_url.as_ref().unwrap(); let request = reqwest::Client::new().post("https://co.wuk.sh/api/json") @@ -127,7 +126,7 @@ async fn execute_get_media(args: Args, bulk: u16, debug: bool) -> bool { let body = res.text().await.unwrap(); if debug { eprintln!("[DEBUG {download_url}] Response received, parsing json ...") }; - let json = attempt!(json::parse(&body), print_cobalt_error("{}".into(), body)); + let json = attempt!(json::parse(&body), proc::print_json_error("{}".into(), body)); let status = json.get("status".into()).unwrap().get_str().unwrap(); match status.as_str() { @@ -139,7 +138,7 @@ async fn execute_get_media(args: Args, bulk: u16, debug: bool) -> bool { "stream" | "redirect" | "success" | "picker" => { if debug { eprintln!("[DEBUG {download_url}] Cobalt returned a response") }; - let url = get_url(&args, &status, &json); + let url = proc::get_url(&args, &status, &json); let media = if args.c_audio_only { "audio" @@ -153,7 +152,7 @@ async fn execute_get_media(args: Args, bulk: u16, debug: bool) -> bool { let res = attempt!(stream_request.send().await, "Live renderer did not respond:\n\"{}\"\n(when downloading from {download_url})"); if debug { eprintln!("[DEBUG {download_url}] Response received from stream") }; - let filename = extract_filename(&args, res.headers(), bulk, debug); + let filename = proc::extract_filename(&args, res.headers(), bulk, debug); println!( "Downloading {} from {} ...", media, @@ -183,167 +182,6 @@ async fn execute_get_media(args: Args, bulk: u16, debug: bool) -> bool { true } -fn print_cobalt_error(error: String, body: String) -> String { - let mut text = String::new(); - text.push_str("Cobalt server returned improper JSON\n"); - text.push_str(&format!("JSON parse error: {error}\n")); - if std::env::var("TCOBALT_DEBUG").is_ok_and(|v| v == 1.to_string()) == true { - text.push_str(&format!("\n[DEBUG] Cobalt returned response:\n{body}\n\n")); - text.push_str("[DEBUG] If this response isn't proper JSON, please contact wukko about this error.\n"); - text.push_str("[DEBUG] If this looks like proper json, contact khyernet/khyerdev about his json parser not functioning right."); - } else { - text.push_str("Contact wukko about this error. Run with TCOBALT_DEBUG=1 to see the incorrect response.") - } - text -} - -fn get_url(args: &Args, status: &str, json: &std::collections::HashMap) -> String { - let media = if args.c_audio_only { - "audio" - } else { - "video" - }; - - if status == "picker" { - let urls = { - let mut urls: Vec = Vec::new(); - let picker_array = json.get("picker").unwrap().get_array().unwrap(); - for picker in picker_array.iter() { - urls.push(picker.get_object().unwrap().get("url").unwrap().get_str().unwrap()); - } - urls - }; - - let choice = if args.picker_choice == 0 { - loop { - let mut buf = String::new(); - print!("Choose which {media} to download [1-{}] >> ", urls.len()); - std::io::stdout().flush().unwrap(); - std::io::stdin().read_line(&mut buf).unwrap(); - if let Ok(int) = buf.trim().parse::() { - if int as usize <= urls.len() { - break int; - } - } - println!("Input must be an integer between 1 and {}", urls.len()); - } - } else { - args.picker_choice - }; - - urls.get((choice - 1) as usize).unwrap_or(&urls[0]).clone() - } else { - json.get("url").unwrap().get_str().unwrap() - } -} - -fn extract_filename(args: &Args, headers: &reqwest::header::HeaderMap, bulk: u16, debug: bool) -> String { - match &args.out_filename { - Some(name) => { - if bulk > 0 { - format!("{bulk}-{name}") - } else { - name.to_string() - } - }, - None => { - let download_url = args.c_url.clone().unwrap(); - if debug { eprintln!("[DEBUG {download_url}] Obtaining filename from headers") }; - match headers.get("Content-Disposition") { - Some(disposition) => { - let disposition = disposition.to_str().unwrap(); - let mut pass: u8 = 0; - let mut filename = String::new(); - for c in disposition.chars() { - if c == ';' || c == '\"' { - pass += 1; - continue; - } - if pass == 2 { - filename.push(c); - } - if pass == 3 { - break; - } - } - filename - }, - None => { - if debug { eprintln!("[DEBUG {download_url}] No filename specified, generating random filename ...") }; - let mut hasher = std::collections::hash_map::DefaultHasher::new(); - download_url.hash(&mut hasher); - let mut hash = format!("{:x}", hasher.finish()); - if args.c_twitter_gif { - hash.push_str(".gif"); - } else { - match args.c_video_codec { - tcargs::types::VideoCodec::AV1 | tcargs::types::VideoCodec::H264 => { - hash.push_str(".mp4"); - } - tcargs::types::VideoCodec::VP9 => { - hash.push_str(".webm"); - } - } - } - hash - } - } - } - } -} - -#[macro_export] -macro_rules! attempt { - ($try: expr, $error_msg_format: literal $(,$($extra:expr),*)?) => {{ - let result = $try; - if result.is_err() { - let e = result.unwrap_err().to_string(); - eprintln!($error_msg_format, e $(,$($extra)*)?); - return false; - } - result.unwrap() - }}; - ($try: expr, $error_string_generator: expr) => {{ - let result = $try; - if result.is_err() { - let e = result.unwrap_err().to_string(); - let diag = $error_string_generator; - eprintln!("{}", diag.to_string().replace("{}", &e)); - return false; - } - result.unwrap() - }}; -} - -const POST_TEMPLATE: &str = "{ - \"url\": \"\", - \"vCodec\": \"\", - \"vQuality\": \"\", - \"aFormat\": \"\", - \"filenamePattern\": \"\", - \"isAudioOnly\": , - \"isTTFullAudio\": , - \"tiktokH265\": , - \"isAudioMuted\": , - \"dubLang\": , - \"disableMetadata\": , - \"twitterGif\": -}"; -fn cobalt_args(args_in: &Args) -> String { - POST_TEMPLATE.to_string() - .replace("", &args_in.c_url.clone().unwrap()) - .replace("", &args_in.c_video_codec.print()) - .replace("", &args_in.c_video_quality.to_string()) - .replace("", &args_in.c_audio_format.print()) - .replace("", &args_in.c_fname_style.print()) - .replace("", &args_in.c_tt_full_audio.to_string()) - .replace("", &args_in.c_tt_h265.to_string()) - .replace("", &args_in.c_audio_only.to_string()) - .replace("", &args_in.c_audio_muted.to_string()) - .replace("", &args_in.c_dublang.to_string()) - .replace("", &args_in.c_disable_metadata.to_string()) - .replace("", &args_in.c_twitter_gif.to_string()) -} #[cfg(test)] mod tests; diff --git a/src/process.rs b/src/process.rs new file mode 100644 index 0000000..fc625f4 --- /dev/null +++ b/src/process.rs @@ -0,0 +1,165 @@ +use crate::{tcargs, args::Args, json}; +use std::io::Write; +use std::hash::{Hash, Hasher}; + +pub fn print_json_error(error: String, body: String) -> String { + let mut text = String::new(); + text.push_str("Cobalt server returned improper JSON\n"); + text.push_str(&format!("JSON parse error: {error}\n")); + if std::env::var("TCOBALT_DEBUG").is_ok_and(|v| v == 1.to_string()) == true { + text.push_str(&format!("\n[DEBUG] Cobalt returned response:\n{body}\n\n")); + text.push_str("[DEBUG] If this response isn't proper JSON, please contact wukko about this error.\n"); + text.push_str("[DEBUG] If this looks like proper json, contact khyernet/khyerdev about his json parser not functioning right."); + } else { + text.push_str("Contact wukko about this error. Run with TCOBALT_DEBUG=1 to see the incorrect response.") + } + text +} + +pub fn get_url(args: &Args, status: &str, json: &std::collections::HashMap) -> String { + let media = if args.c_audio_only { + "audio" + } else { + "video" + }; + + if status == "picker" { + let urls = { + let mut urls: Vec = Vec::new(); + let picker_array = json.get("picker").unwrap().get_array().unwrap(); + for picker in picker_array.iter() { + urls.push(picker.get_object().unwrap().get("url").unwrap().get_str().unwrap()); + } + urls + }; + + let choice = if args.picker_choice == 0 { + loop { + let mut buf = String::new(); + print!("Choose which {media} to download [1-{}] >> ", urls.len()); + std::io::stdout().flush().unwrap(); + std::io::stdin().read_line(&mut buf).unwrap(); + if let Ok(int) = buf.trim().parse::() { + if int as usize <= urls.len() { + break int; + } + } + println!("Input must be an integer between 1 and {}", urls.len()); + } + } else { + args.picker_choice + }; + + urls.get((choice - 1) as usize).unwrap_or(&urls[0]).clone() + } else { + json.get("url").unwrap().get_str().unwrap() + } +} + +pub fn extract_filename(args: &Args, headers: &reqwest::header::HeaderMap, bulk: u16, debug: bool) -> String { + match &args.out_filename { + Some(name) => { + if bulk > 0 { + format!("{bulk}-{name}") + } else { + name.to_string() + } + }, + None => { + let download_url = args.c_url.clone().unwrap(); + if debug { eprintln!("[DEBUG {download_url}] Obtaining filename from headers") }; + match headers.get("Content-Disposition") { + Some(disposition) => { + let disposition = disposition.to_str().unwrap(); + let mut pass: u8 = 0; + let mut filename = String::new(); + for c in disposition.chars() { + if c == ';' || c == '\"' { + pass += 1; + continue; + } + if pass == 2 { + filename.push(c); + } + if pass == 3 { + break; + } + } + filename + }, + None => { + if debug { eprintln!("[DEBUG {download_url}] No filename specified, generating random filename ...") }; + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + download_url.hash(&mut hasher); + let mut hash = format!("{:x}", hasher.finish()); + if args.c_twitter_gif { + hash.push_str(".gif"); + } else { + match args.c_video_codec { + tcargs::types::VideoCodec::AV1 | tcargs::types::VideoCodec::H264 => { + hash.push_str(".mp4"); + } + tcargs::types::VideoCodec::VP9 => { + hash.push_str(".webm"); + } + } + } + hash + } + } + } + } +} + +const POST_TEMPLATE: &str = "{ + \"url\": \"\", + \"vCodec\": \"\", + \"vQuality\": \"\", + \"aFormat\": \"\", + \"filenamePattern\": \"\", + \"isAudioOnly\": , + \"isTTFullAudio\": , + \"tiktokH265\": , + \"isAudioMuted\": , + \"dubLang\": , + \"disableMetadata\": , + \"twitterGif\": + }"; +pub fn cobalt_args(args_in: &Args) -> String { + POST_TEMPLATE.to_string() + .replace("", &args_in.c_url.clone().unwrap()) + .replace("", &args_in.c_video_codec.print()) + .replace("", &args_in.c_video_quality.to_string()) + .replace("", &args_in.c_audio_format.print()) + .replace("", &args_in.c_fname_style.print()) + .replace("", &args_in.c_tt_full_audio.to_string()) + .replace("", &args_in.c_tt_h265.to_string()) + .replace("", &args_in.c_audio_only.to_string()) + .replace("", &args_in.c_audio_muted.to_string()) + .replace("", &args_in.c_dublang.to_string()) + .replace("", &args_in.c_disable_metadata.to_string()) + .replace("", &args_in.c_twitter_gif.to_string()) +} + +#[macro_export] +macro_rules! attempt { + ($try: expr, $error_msg_format: literal $(,$($extra:expr),*)?) => {{ + let result = $try; + if result.is_err() { + let e = result.unwrap_err().to_string(); + eprintln!($error_msg_format, e $(,$($extra)*)?); + return false; + } + result.unwrap() + }}; + ($try: expr, $error_string_generator: expr) => {{ + let result = $try; + if result.is_err() { + let e = result.unwrap_err().to_string(); + let diag = $error_string_generator; + eprintln!("{}", diag.to_string().replace("{}", &e)); + return false; + } + result.unwrap() + }}; +}