Merge pull request #105 from rhysd/refactor-main

Refactoring for main.rs to address several issues

src/main.rs

@@ -11,49 +11,60 @@ 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::fs::File;
-use std::io::{Error, Write};
+use std::io::{self, Error, Write};
+use std::process;
 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() {
     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 cache = HashMap::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()
         .timeout(Duration::from_secs(10))
         .danger_accept_invalid_certs(app_args.insecure)
@@ -63,18 +74,18 @@ fn main() {
 
     // Retrieve root document
     let (data, final_url) = retrieve_asset(
-            cache,
+        &mut cache,
         &client,
         app_args.url_target.as_str(),
         false,
         "",
         app_args.silent,
     )
-        .unwrap();
+    .expect("Could not retrieve assets in HTML");
     let dom = html_to_dom(&data);
 
     walk_and_embed_assets(
-            cache,
+        &mut cache,
         &client,
         &final_url,
         &dom.document,
@@ -94,10 +105,7 @@ fn main() {
         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();
-        }
-    }
 }