From bfb97bd062680b5ada796659d36ec2dad1359d66 Mon Sep 17 00:00:00 2001
From: Sunshine <sunshine@uberspace.net>
Date: Thu, 26 Dec 2019 00:41:03 -0500
Subject: [PATCH] add option for saving output to file

---
 README.md   |  3 ++-
 src/args.rs |  3 +++
 src/main.rs | 36 +++++++++++++++++++++++++++++++++++-
 3 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 4474d26..d5782bb 100644
--- a/README.md
+++ b/README.md
@@ -29,7 +29,7 @@ If compared to saving websites with `wget -mpk`, this tool embeds all assets as
     $ brew install monolith
 
 ## Usage
-    $ monolith https://lyrics.github.io/db/p/portishead/dummy/roads/ > portishead-roads-lyrics.html
+    $ monolith https://lyrics.github.io/db/p/portishead/dummy/roads/ -o portishead-roads-lyrics.html
 
 ## Options
  - `-c`: Ignore styles
@@ -38,6 +38,7 @@ If compared to saving websites with `wget -mpk`, this tool embeds all assets as
  - `-I`: Isolate document
  - `-j`: Exclude JavaScript
  - `-k`: Accept invalid X.509 (TLS) certificates
+ - `-o`: Write output to file
  - `-s`: Silent mode
  - `-u`: Specify custom User-Agent
 
diff --git a/src/args.rs b/src/args.rs
index 34fe5d0..1a67bcd 100644
--- a/src/args.rs
+++ b/src/args.rs
@@ -9,6 +9,7 @@ pub struct AppArgs {
     pub no_js: bool,
     pub insecure: bool,
     pub isolate: bool,
+    pub output: String,
     pub silent: bool,
     pub user_agent: String,
 }
@@ -36,6 +37,7 @@ impl AppArgs {
             .args_from_usage("-I, --isolate 'Cut off from the Internet'")
             .args_from_usage("-j, --no-js 'Exclude JavaScript'")
             .args_from_usage("-k, --insecure 'Accept invalid X.509 (TLS) certificates'")
+            .args_from_usage("-o, --output=[document.html] 'Write output to <file>'")
             .args_from_usage("-s, --silent 'Suppress verbosity'")
             .args_from_usage("-u, --user-agent=[Iceweasel] 'Custom User-Agent string'")
             // .args_from_usage("-v, --include-video 'Embed video sources'")
@@ -53,6 +55,7 @@ impl AppArgs {
         app_args.insecure = app.is_present("insecure");
         app_args.isolate = app.is_present("isolate");
         app_args.silent = app.is_present("silent");
+        app_args.output = app.value_of("output").unwrap_or("").to_string();
         app_args.user_agent = app
             .value_of("user-agent")
             .unwrap_or_else(|| DEFAULT_USER_AGENT)
diff --git a/src/main.rs b/src/main.rs
index b7298e0..1e4fe2c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,6 +4,7 @@ extern crate monolith;
 extern crate reqwest;
 
 mod args;
+mod macros;
 
 use args::AppArgs;
 use monolith::html::{html_to_dom, stringify_document, walk_and_embed_assets};
@@ -11,11 +12,39 @@ use monolith::http::retrieve_asset;
 use monolith::utils::is_valid_url;
 use reqwest::header::{HeaderMap, HeaderValue, USER_AGENT};
 use std::collections::HashMap;
+use std::fs::{remove_file, File};
+use std::io::{Error, Write};
 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),
+    };
+
+    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(())
+}
+
 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()) {
         // Initialize client
         let mut header_map = HeaderMap::new();
@@ -33,6 +62,7 @@ fn main() {
             .build()
             .expect("Failed to initialize HTTP client");
 
+        // Retrieve root document
         let (data, final_url) = retrieve_asset(
             cache,
             &client,
@@ -65,6 +95,10 @@ fn main() {
             app_args.isolate,
         );
 
-        println!("{}", html);
+        if app_args.output == str!() {
+            println!("{}", html);
+        } else {
+            create_file(&app_args.output, html).unwrap();
+        }
     }
 }