155 lines
6.9 KiB
Rust
155 lines
6.9 KiB
Rust
|
use crate::utils;
|
||
|
|
||
|
use std::fs::rename;
|
||
|
use std::path::Path;
|
||
|
use std::process::{exit, Command};
|
||
|
use std::collections::HashMap;
|
||
|
use clap::ArgMatches;
|
||
|
use reqwest::Client;
|
||
|
use reqwest::redirect::Policy;
|
||
|
|
||
|
const MAX_DOWNLOAD_ATTEMPTS: i32 = 5;
|
||
|
|
||
|
pub async fn download(arg_m: &ArgMatches<'_>) {
|
||
|
let print_only = arg_m.is_present("print");
|
||
|
let resolution = arg_m.value_of("resolution");
|
||
|
let ids = arg_m.values_of("id").unwrap().collect::<Vec<_>>();
|
||
|
let policy = Policy::custom(|attempt| {
|
||
|
if attempt.previous().len() > 10 {
|
||
|
attempt.error("too many redirects")
|
||
|
} else if attempt.url().path() == "/404" {
|
||
|
attempt.stop()
|
||
|
} else {
|
||
|
attempt.follow()
|
||
|
}
|
||
|
});
|
||
|
let client = Client::builder().redirect(policy).build().unwrap();
|
||
|
let mut return_fail = false;
|
||
|
for id in ids {
|
||
|
let hentai_info = utils::get_hentai(client.clone(), id).await;
|
||
|
match hentai_info {
|
||
|
Ok(hentai_info) => {
|
||
|
match hentai_info {
|
||
|
Some(hentai_info) => {
|
||
|
let slug = hentai_info.state.data.video.hentai_video.slug;
|
||
|
let filename = format!("{}.mkv", &slug);
|
||
|
if !print_only && Path::new(&filename).exists() {
|
||
|
continue;
|
||
|
}
|
||
|
let mut download_url = None;
|
||
|
let mut perm_urls: HashMap<String, String> = HashMap::new();
|
||
|
let mut temp_urls: HashMap<String, String> = HashMap::new();
|
||
|
for server in hentai_info.state.data.video.videos_manifest.servers {
|
||
|
let mut to_hashmap = match server.is_permanent {
|
||
|
true => perm_urls.clone(),
|
||
|
false => temp_urls.clone()
|
||
|
};
|
||
|
for stream in server.streams {
|
||
|
if stream.url.is_empty() {
|
||
|
continue;
|
||
|
}
|
||
|
if server.is_permanent && Some(stream.height.as_str()) == resolution {
|
||
|
download_url = Some(stream.url);
|
||
|
break;
|
||
|
}
|
||
|
if !to_hashmap.contains_key(&stream.height) {
|
||
|
to_hashmap.insert(stream.height, stream.url);
|
||
|
};
|
||
|
}
|
||
|
if download_url.is_some() {
|
||
|
break;
|
||
|
}
|
||
|
match server.is_permanent {
|
||
|
true => perm_urls.extend(to_hashmap),
|
||
|
false => temp_urls.extend(to_hashmap)
|
||
|
};
|
||
|
}
|
||
|
if download_url.is_none() {
|
||
|
if resolution.is_some() {
|
||
|
download_url = temp_urls.get(resolution.unwrap()).map(|i| i.to_string());
|
||
|
}
|
||
|
if download_url.is_none() {
|
||
|
download_url = magic_thing(perm_urls).or_else(|| { magic_thing(temp_urls) });
|
||
|
if download_url.is_none() {
|
||
|
eprintln!("Failed to get {}: cannot get download url", id);
|
||
|
return_fail = true;
|
||
|
continue;
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
let download_url = download_url.unwrap();
|
||
|
if print_only {
|
||
|
println!("{}", download_url);
|
||
|
} else {
|
||
|
let mut fail_dl = true;
|
||
|
let tmp_filename = format!("{}.tmp", &slug);
|
||
|
for i in 0..MAX_DOWNLOAD_ATTEMPTS {
|
||
|
eprintln!("Downloading {} (attempt {})", &filename, i);
|
||
|
match Command::new("ffmpeg").args(&["-v", "warning", "-stats", "-nostdin", "-y", "-i", &download_url, "-c", "copy", "-f", "matroska", &tmp_filename]).spawn() {
|
||
|
Ok(mut child) => {
|
||
|
match child.wait() {
|
||
|
Ok(exit_status) => {
|
||
|
if exit_status.success() {
|
||
|
fail_dl = false;
|
||
|
match rename(&tmp_filename, &filename) {
|
||
|
Ok(_) => (),
|
||
|
Err(err) => eprintln!("Failed to rename {} to {} due to {}", &tmp_filename, &filename, err)
|
||
|
};
|
||
|
break;
|
||
|
}
|
||
|
eprintln!("ffmpeg exited with {:?}", exit_status.code());
|
||
|
},
|
||
|
Err(err) => eprintln!("Failed to wait on ffmpeg process due to {}", err)
|
||
|
};
|
||
|
},
|
||
|
Err(err) => eprintln!("Failed to spawn ffmpeg process due to {}", err)
|
||
|
};
|
||
|
}
|
||
|
if fail_dl {
|
||
|
eprintln!("Failed to download {}", &filename);
|
||
|
return_fail = true;
|
||
|
}
|
||
|
}
|
||
|
},
|
||
|
None => {
|
||
|
eprintln!("Failed to get {}: does not exist", id);
|
||
|
return_fail = true;
|
||
|
}
|
||
|
};
|
||
|
},
|
||
|
Err(err) => {
|
||
|
eprintln!("Failed to get {}: {}", id, err);
|
||
|
return_fail = true;
|
||
|
}
|
||
|
};
|
||
|
}
|
||
|
if return_fail {
|
||
|
exit(1);
|
||
|
}
|
||
|
}
|
||
|
|
||
|
fn magic_thing(map: HashMap<String, String>) -> Option<String> {
|
||
|
let mut keys = Vec::new();
|
||
|
for i in map.keys() {
|
||
|
match i.parse::<i32>() {
|
||
|
Ok(i) => keys.push(i),
|
||
|
Err(_) => {
|
||
|
keys.clear();
|
||
|
break;
|
||
|
}
|
||
|
};
|
||
|
}
|
||
|
keys.sort();
|
||
|
match keys.pop() {
|
||
|
Some(key) => Some(map.get(&key.to_string()).unwrap().to_string()),
|
||
|
None => {
|
||
|
let mut keys: Vec<_> = map.keys().collect();
|
||
|
keys.sort();
|
||
|
match keys.pop() {
|
||
|
Some(key) => Some(map.get(&key.to_string()).unwrap().to_string()),
|
||
|
None => None
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
}
|