From 4e4ebe9c9890d554330f9b9e6ec56432f8ed13d2 Mon Sep 17 00:00:00 2001
From: rhysd <lin90162@yahoo.co.jp>
Date: Wed, 15 Jan 2020 12:26:04 +0900
Subject: [PATCH] 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
---
 src/main.rs | 154 +++++++++++++++++++++++++++-------------------------
 1 file changed, 81 insertions(+), 73 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 54c267c..a805133 100644
--- a/src/main.rs
+++ b/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());
-
-    let mut file = match file {
-        Ok(file) => file,
-        Err(error) => return Err(error),
-    };
+enum Output {
+    Stdout(io::Stdout),
+    File(File),
+}
 
-    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");
 }