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
1 changed files with 81 additions and 73 deletions

View File

@ -11,93 +11,101 @@ 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 { impl Output {
Ok(file) => file, fn new(file_path: &str) -> Result<Output, Error> {
Err(error) => return Err(error), if file_path.is_empty() {
}; Ok(Output::Stdout(io::stdout()))
} else {
if content != str!() { Ok(Output::File(File::create(file_path)?))
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(()) 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
let mut header_map = HeaderMap::new();
match HeaderValue::from_str(&app_args.user_agent) {
Ok(header) => header_map.insert(USER_AGENT, header),
Err(err) => {
eprintln!("Invalid user agent! {}", err);
return;
}
};
let client = Client::builder()
.timeout(Duration::from_secs(10))
.danger_accept_invalid_certs(app_args.insecure)
.default_headers(header_map)
.build()
.expect("Failed to initialize HTTP client");
// Retrieve root document // Initialize client
let (data, final_url) = retrieve_asset( let mut cache = HashMap::new();
cache, let mut header_map = HeaderMap::new();
&client, header_map.insert(
app_args.url_target.as_str(), USER_AGENT,
false, HeaderValue::from_str(&app_args.user_agent).expect("Invalid User-Agent header specified"),
"", );
app_args.silent,
)
.unwrap();
let dom = html_to_dom(&data);
walk_and_embed_assets( let client = Client::builder()
cache, .timeout(Duration::from_secs(10))
&client, .danger_accept_invalid_certs(app_args.insecure)
&final_url, .default_headers(header_map)
&dom.document, .build()
app_args.no_css, .expect("Failed to initialize HTTP client");
app_args.no_js,
app_args.no_images,
app_args.silent,
app_args.no_frames,
);
let html: String = stringify_document( // Retrieve root document
&dom.document, let (data, final_url) = retrieve_asset(
app_args.no_css, &mut cache,
app_args.no_frames, &client,
app_args.no_js, app_args.url_target.as_str(),
app_args.no_images, false,
app_args.isolate, "",
); app_args.silent,
)
.expect("Could not retrieve assets in HTML");
let dom = html_to_dom(&data);
if app_args.output == str!() { walk_and_embed_assets(
println!("{}", html); &mut cache,
} else { &client,
create_file(&app_args.output, html).unwrap(); &final_url,
} &dom.document,
} app_args.no_css,
app_args.no_js,
app_args.no_images,
app_args.silent,
app_args.no_frames,
);
let html: String = stringify_document(
&dom.document,
app_args.no_css,
app_args.no_frames,
app_args.no_js,
app_args.no_images,
app_args.isolate,
);
output
.writeln_str(&html)
.expect("Could not write HTML output");
} }