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:
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::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");
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in a new issue