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

189 lines
8.4 KiB

use crate::utils;
use clap::ArgMatches;
use reqwest::redirect::Policy;
use reqwest::Client;
use std::collections::HashMap;
use std::fs::rename;
use std::path::Path;
use std::process::{exit, Command};
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, i32)> = HashMap::new();
let mut temp_urls: HashMap<String, (String, i32)> = 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, stream.filesize_mbs));
break;
}
if !to_hashmap.contains_key(&stream.height) {
to_hashmap
.insert(stream.height, (stream.url, stream.filesize_mbs));
};
}
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.0.to_string(), i.1));
}
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, filesize_mbs) = 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 {} ({}MB, attempt {})",
&filename, filesize_mbs, 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, i32)>) -> Option<(String, i32)> {
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) => map.get(&key.to_string()).map(|i| i.clone()),
None => {
let mut keys: Vec<_> = map.keys().collect();
keys.sort();
match keys.pop() {
Some(key) => map.get(key.as_str()).map(|i| i.clone()),
None => None,
}
}
}
}