You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
hanimers/src/commands/download.rs

154 lines
6.9 KiB

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
}
}
}
}