From d97d2714ce62e1d6e30a79ae21b525911efbb9f9 Mon Sep 17 00:00:00 2001
From: epi <epibar052@gmail.com>
Date: Tue, 3 Nov 2020 12:46:21 -0600
Subject: [PATCH] fixed comments from review

---
 src/lib.rs         | 19 +++++++++++++++++--
 src/main.rs        | 30 ++++++++++++++++++++++--------
 src/scanner.rs     |  2 +-
 tests/test_main.rs |  2 +-
 4 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 8a43797..ccad808 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -15,11 +15,26 @@ use reqwest::{
     header::HeaderMap,
     {Response, StatusCode, Url},
 };
+use std::{error, fmt};
 use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
 
 /// Generic Result type to ease error handling in async contexts
-pub type FeroxResult<T> =
-    std::result::Result<T, Box<dyn std::error::Error + Send + Sync + 'static>>;
+pub type FeroxResult<T> = std::result::Result<T, Box<dyn error::Error + Send + Sync + 'static>>;
+
+/// Simple Error implementation to allow for custom error returns
+#[derive(Debug, Default)]
+pub struct FeroxError {
+    /// fancy string that can be printed via Display
+    pub message: String,
+}
+
+impl error::Error for FeroxError {}
+
+impl fmt::Display for FeroxError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}", &self.message)
+    }
+}
 
 /// Generic mpsc::unbounded_channel type to tidy up some code
 pub type FeroxChannel<T> = (UnboundedSender<T>, UnboundedReceiver<T>);
diff --git a/src/main.rs b/src/main.rs
index 4ffe459..de77116 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,8 +4,8 @@ use feroxbuster::{
     config::{CONFIGURATION, PROGRESS_BAR, PROGRESS_PRINTER},
     heuristics, logger, reporter,
     scanner::{scan_url, PAUSE_SCAN},
-    utils::{ferox_print, get_current_depth, module_colorizer},
-    FeroxResponse, FeroxResult, SLEEP_DURATION, VERSION,
+    utils::{ferox_print, get_current_depth, module_colorizer, status_colorizer},
+    FeroxError, FeroxResponse, FeroxResult, SLEEP_DURATION, VERSION,
 };
 use futures::StreamExt;
 use std::{
@@ -105,8 +105,9 @@ async fn scan(
             .await??;
 
     if words.len() == 0 {
-        eprintln!("Did not find any words in {}", CONFIGURATION.wordlist);
-        process::exit(1); // todo cleanup?
+        let mut err = FeroxError::default();
+        err.message = format!("Did not find any words in {}", CONFIGURATION.wordlist);
+        return Err(Box::new(err));
     }
 
     let mut tasks = vec![];
@@ -131,6 +132,7 @@ async fn scan(
     Ok(())
 }
 
+/// Get targets from either commandline or stdin, pass them back to the caller as a Result<Vec>
 async fn get_targets() -> FeroxResult<Vec<String>> {
     log::trace!("enter: get_targets");
 
@@ -154,7 +156,8 @@ async fn get_targets() -> FeroxResult<Vec<String>> {
     Ok(targets)
 }
 
-/// todo doc
+/// async main called from real main, broken out in this way to allow for some synchronous code
+/// to be executed before bringing the tokio runtime online
 async fn wrapped_main() {
     // join can only be called once, otherwise it causes the thread to panic
     tokio::task::spawn_blocking(move || {
@@ -215,7 +218,7 @@ async fn wrapped_main() {
         }
         Err(e) => {
             ferox_print(
-                &format!("An error occurred while scanning: {}", e),
+                &format!("{} while scanning: {}", status_colorizer("Error"), e),
                 &PROGRESS_PRINTER,
             );
             clean_up(tx_term, term_handle, tx_file, file_handle, save_output).await;
@@ -228,7 +231,8 @@ async fn wrapped_main() {
     log::trace!("exit: main");
 }
 
-/// todo doc
+/// Single cleanup function that handles all the necessary drops/finishes etc required to gracefully
+/// shutdown the program
 async fn clean_up(
     tx_term: UnboundedSender<FeroxResponse>,
     term_handle: JoinHandle<()>,
@@ -236,7 +240,15 @@ async fn clean_up(
     file_handle: Option<JoinHandle<()>>,
     save_output: bool,
 ) {
-    // todo trace
+    log::trace!(
+        "enter: clean_up({:?}, {:?}, {:?}, {:?}, {}",
+        tx_term,
+        term_handle,
+        tx_file,
+        file_handle,
+        save_output
+    );
+
     drop(tx_term);
     log::trace!("dropped terminal output handler's transmitter");
 
@@ -274,6 +286,8 @@ async fn clean_up(
     // clean-up function for the MultiProgress bar; must be called last in order to still see
     // the final trace messages above
     PROGRESS_PRINTER.finish();
+
+    log::trace!("exit: clean_up");
 }
 
 fn main() {
diff --git a/src/scanner.rs b/src/scanner.rs
index 4054a41..f783251 100644
--- a/src/scanner.rs
+++ b/src/scanner.rs
@@ -621,7 +621,7 @@ pub async fn scan_url(
     if CALL_COUNT.load(Ordering::Relaxed) == 0 {
         CALL_COUNT.fetch_add(1, Ordering::Relaxed);
 
-        // this protection around join also allows us to add the first scanned url to SCANNED_URLS
+        // this protection allows us to add the first scanned url to SCANNED_URLS
         // from within the scan_url function instead of the recursion handler
         add_url_to_list_of_scanned_urls(&target_url, &SCANNED_URLS);
 
diff --git a/tests/test_main.rs b/tests/test_main.rs
index fe403c4..04b37a4 100644
--- a/tests/test_main.rs
+++ b/tests/test_main.rs
@@ -55,7 +55,7 @@ fn main_use_empty_wordlist() -> Result<(), Box<dyn std::error::Error>> {
         .arg("-vvvv")
         .assert()
         .failure()
-        .stderr(predicate::str::contains("Did not find any words in"));
+        .stdout(predicate::str::contains("Did not find any words in"));
 
     assert_eq!(mock.times_called(), 1);