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
parent
c1dc798ded
commit
4e4ebe9c98
1 changed files with 81 additions and 73 deletions
154
src/main.rs
154
src/main.rs
|
|
@ -11,93 +11,101 @@ use monolith::utils::is_valid_url;
|
|||
use reqwest::blocking::Client;
|
||||
use reqwest::header::{HeaderMap, HeaderValue, USER_AGENT};
|
||||
use std::collections::HashMap;
|
||||
use std::fs::{remove_file, File};
|
||||
use std::io::{Error, Write};
|
||||
use std::fs::File;
|
||||
use std::io::{self, Error, Write};
|
||||
use std::process;
|
||||
use std::time::Duration;
|
||||
|
||||
fn create_file(file_path: &String, content: String) -> Result<(), Error> {
|
||||
let file = File::create(file_path.as_str());
|
||||
enum Output {
|
||||
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())?;
|
||||
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)?))
|
||||
}
|
||||
}
|
||||
|
||||
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() {
|
||||
let app_args = AppArgs::get();
|
||||
let cache = &mut HashMap::new();
|
||||
|
||||
// Attempt to create output file
|
||||
if app_args.output != str!() {
|
||||
create_file(&app_args.output, str!()).unwrap();
|
||||
if !is_valid_url(app_args.url_target.as_str()) {
|
||||
eprintln!(
|
||||
"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()) {
|
||||
// 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");
|
||||
let mut output = Output::new(&app_args.output).expect("Could not prepare output");
|
||||
|
||||
// Retrieve root document
|
||||
let (data, final_url) = retrieve_asset(
|
||||
cache,
|
||||
&client,
|
||||
app_args.url_target.as_str(),
|
||||
false,
|
||||
"",
|
||||
app_args.silent,
|
||||
)
|
||||
.unwrap();
|
||||
let dom = html_to_dom(&data);
|
||||
// Initialize client
|
||||
let mut cache = HashMap::new();
|
||||
let mut header_map = HeaderMap::new();
|
||||
header_map.insert(
|
||||
USER_AGENT,
|
||||
HeaderValue::from_str(&app_args.user_agent).expect("Invalid User-Agent header specified"),
|
||||
);
|
||||
|
||||
walk_and_embed_assets(
|
||||
cache,
|
||||
&client,
|
||||
&final_url,
|
||||
&dom.document,
|
||||
app_args.no_css,
|
||||
app_args.no_js,
|
||||
app_args.no_images,
|
||||
app_args.silent,
|
||||
app_args.no_frames,
|
||||
);
|
||||
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");
|
||||
|
||||
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,
|
||||
);
|
||||
// Retrieve root document
|
||||
let (data, final_url) = retrieve_asset(
|
||||
&mut cache,
|
||||
&client,
|
||||
app_args.url_target.as_str(),
|
||||
false,
|
||||
"",
|
||||
app_args.silent,
|
||||
)
|
||||
.expect("Could not retrieve assets in HTML");
|
||||
let dom = html_to_dom(&data);
|
||||
|
||||
if app_args.output == str!() {
|
||||
println!("{}", html);
|
||||
} else {
|
||||
create_file(&app_args.output, html).unwrap();
|
||||
}
|
||||
}
|
||||
walk_and_embed_assets(
|
||||
&mut cache,
|
||||
&client,
|
||||
&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");
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue