refactor main to address several issues

Addressed issues:

- when specified URL is invalid, it exited successfully with doing
  nothing. There was no way why it does not work for users
- it exited successfully even if invalid User-Agent value is specified
- it created file twice on `--output` option specified. It may cause an
  issue when some file watcher (e.g. FsEvents on macOS) is watching

Improvements:
- handle errors with `Result::expect` consistently it correctly exits
  with non-zero status on error
- define `Output` enum for handling both stdout and file outputs
This commit is contained in:
rhysd 2020-01-15 12:26:04 +09:00
parent c1dc798ded
commit 4e4ebe9c98

View File

@ -11,49 +11,60 @@ use monolith::utils::is_valid_url;
use reqwest::blocking::Client; use reqwest::blocking::Client;
use reqwest::header::{HeaderMap, HeaderValue, USER_AGENT}; use reqwest::header::{HeaderMap, HeaderValue, USER_AGENT};
use std::collections::HashMap; use std::collections::HashMap;
use std::fs::{remove_file, File}; use std::fs::File;
use std::io::{Error, Write}; use std::io::{self, Error, Write};
use std::process;
use std::time::Duration; use std::time::Duration;
fn create_file(file_path: &String, content: String) -> Result<(), Error> { enum Output {
let file = File::create(file_path.as_str()); Stdout(io::Stdout),
File(File),
let mut file = match file {
Ok(file) => file,
Err(error) => return Err(error),
};
if content != str!() {
file.write_all(content.as_bytes())?;
file.write_all("\n".as_bytes())?;
file.sync_all()?;
} else {
// Remove the file right away if it had no content
remove_file(file_path.as_str())?;
} }
Ok(()) impl Output {
fn new(file_path: &str) -> Result<Output, Error> {
if file_path.is_empty() {
Ok(Output::Stdout(io::stdout()))
} else {
Ok(Output::File(File::create(file_path)?))
}
}
fn writeln_str(&mut self, s: &str) -> Result<(), Error> {
match self {
Output::Stdout(stdout) => {
writeln!(stdout, "{}", s)?;
stdout.flush()
}
Output::File(f) => {
writeln!(f, "{}", s)?;
f.flush()
}
}
}
} }
fn main() { fn main() {
let app_args = AppArgs::get(); let app_args = AppArgs::get();
let cache = &mut HashMap::new();
// Attempt to create output file if !is_valid_url(app_args.url_target.as_str()) {
if app_args.output != str!() { eprintln!(
create_file(&app_args.output, str!()).unwrap(); "Only HTTP and HTTPS URLs are allowed but got: {}",
&app_args.url_target
);
process::exit(1);
} }
if is_valid_url(app_args.url_target.as_str()) { let mut output = Output::new(&app_args.output).expect("Could not prepare output");
// Initialize client // Initialize client
let mut cache = HashMap::new();
let mut header_map = HeaderMap::new(); let mut header_map = HeaderMap::new();
match HeaderValue::from_str(&app_args.user_agent) { header_map.insert(
Ok(header) => header_map.insert(USER_AGENT, header), USER_AGENT,
Err(err) => { HeaderValue::from_str(&app_args.user_agent).expect("Invalid User-Agent header specified"),
eprintln!("Invalid user agent! {}", err); );
return;
}
};
let client = Client::builder() let client = Client::builder()
.timeout(Duration::from_secs(10)) .timeout(Duration::from_secs(10))
.danger_accept_invalid_certs(app_args.insecure) .danger_accept_invalid_certs(app_args.insecure)
@ -63,18 +74,18 @@ fn main() {
// Retrieve root document // Retrieve root document
let (data, final_url) = retrieve_asset( let (data, final_url) = retrieve_asset(
cache, &mut cache,
&client, &client,
app_args.url_target.as_str(), app_args.url_target.as_str(),
false, false,
"", "",
app_args.silent, app_args.silent,
) )
.unwrap(); .expect("Could not retrieve assets in HTML");
let dom = html_to_dom(&data); let dom = html_to_dom(&data);
walk_and_embed_assets( walk_and_embed_assets(
cache, &mut cache,
&client, &client,
&final_url, &final_url,
&dom.document, &dom.document,
@ -94,10 +105,7 @@ fn main() {
app_args.isolate, app_args.isolate,
); );
if app_args.output == str!() { output
println!("{}", html); .writeln_str(&html)
} else { .expect("Could not write HTML output");
create_file(&app_args.output, html).unwrap();
}
}
} }