shrupl/src/bin/shrupl.rs

142 lines
4.1 KiB
Rust

use std::{
process,
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
};
use clap::Parser;
use console::{StyledObject, style};
use log::{info, trace};
use shrupl::{
AppState, Cli,
output::{self, Log, SHRUPL},
};
fn main() {
let args = Cli::parse();
env_logger::Builder::new()
.filter_module("shrupl", args.get_level_filter())
.parse_default_env()
.init();
info!("args: {args:#?}");
println!("{} to {}!", style("Welcome").magenta().bold(), *SHRUPL);
let resumed = AppState::try_resume(&args);
let check_ctrlc = {
let stop = Arc::new(AtomicBool::new(false));
let stop_ctrlc = stop.clone();
ctrlc::set_handler(move || {
stop_ctrlc.store(true, Ordering::SeqCst);
eprintln!("{} stopping as soon as possible!", *SHRUPL);
})
.expect("Error setting Ctrl-C handler");
move || {
if stop.load(Ordering::SeqCst) {
process::exit(255);
}
}
};
let mut state = resumed
.inspect_err(|e| {
AppState::clear_any(&args);
Log::handle(e);
info!("could not resume from hash {:?}: {e}", args.get_hash());
})
.ok()
.and_then(|state| output::prompt_continue().then_some(state))
.unwrap_or_else(|| match AppState::from_args(&args) {
Ok(state) => {
state.save().unwrap_or_else(|e| {
Log::warning(format_args!("Failed to save state: {e}"));
});
state
}
Err(e) => {
Log::handle(&e);
Log::error(format_args!("Failed to create state: {e}"));
}
});
check_ctrlc();
info!("continuing with state: {state:#?}");
let fns_magenta = output::style_all(&args.file_names(), StyledObject::magenta).join(", ");
println!("{} is uploading: {fns_magenta}", *SHRUPL);
let mut buffer = vec![0; args.chunk_size * 1024 * 1024];
let mut tries = 0;
loop {
if !args.may_retry(tries) {
Log::error("Retry limit reached!");
}
match state.upload_chunk(&mut buffer) {
Err(e) => {
Log::handle(&e);
if let Some(p) = e.get_invalid_param() {
match p {
// Error 404 (File not found)
shrupl::Parameter::FileID(fid) => {
info!("retrying file {fid:?}");
state.abort_upload();
}
// Error 404 (Share not found)
shrupl::Parameter::ShareID(sid) => {
output::prompt_rebuild_share();
info!("rebuilding share {sid:?}");
// rebuild share
let Ok(s) = state.rebuild_share(&args) else {
Log::error("Failed to rebuild share!");
};
state = s;
}
p => Log::error(format_args!("Unexpected {p}!")),
}
} else {
// retry chunk
let Some(s) = state.rewind_chunk() else {
Log::error("Failed to retry chunk!");
};
tries += 1;
info!("State rewound, retrying last chunk (tries: {tries})");
state = s;
}
}
Ok(false) => {
trace!("chunk uploaded");
tries = 0;
}
Ok(true) => {
info!("all uploads done");
break;
}
}
state.save().unwrap_or_else(|e| {
Log::warning(format_args!("Failed to save state: {e}"));
});
check_ctrlc();
}
state.discard().unwrap_or_else(|e| {
Log::warning(format_args!("Failed to remove state: {e}"));
});
println!("{} finished {}", *SHRUPL, style("successfully!").green());
}