Merge pull request #104 from epi052/FEATURE-add-pause-resume-functionality

add pause|resume feature

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "feroxbuster"
-version = "1.3.0"
+version = "1.4.0"
 authors = ["Ben 'epi' Risher <epibar052@gmail.com>"]
 license = "MIT"
 edition = "2018"
@@ -32,6 +32,7 @@ console = "0.12"
 openssl = { version = "0.10", features = ["vendored"] }
 dirs = "3.0"
 regex = "1"
+crossterm = "0.18"
 
 [dev-dependencies]
 tempfile = "3.1"

README.md

@@ -71,6 +71,7 @@ This attack is also known as Predictable Resource Location, File Enumeration, Di
     - [ferox-config.toml](#ferox-configtoml)
     - [Command Line Parsing](#command-line-parsing)
 - [Example Usage](#-example-usage)
+    - [Pause and Resume Scans (new in `v1.4.0`)](#pause-and-resume-scans-new-in-v140)
     - [Multiple Values](#multiple-values)
     - [Extract Links from Response Body (new in `v1.1.0`)](#extract-links-from-response-body-new-in-v110)
     - [Include Headers](#include-headers)
@@ -351,6 +352,12 @@ OPTIONS:
 
 ## 🧰 Example Usage
 
+### Pause and Resume Scans (new in `v1.4.0`)
+
+Scans can be paused and resumed by pressing the ENTER key (shown below)
+
+![pause-resume-demo](img/pause-resume-demo.gif)
+
 ### Multiple Values
 
 Options that take multiple values are very flexible.  Consider the following ways of specifying extensions:

src/banner.rs

@@ -1,5 +1,6 @@
 use crate::config::{Configuration, CONFIGURATION};
 use crate::utils::{make_request, status_colorizer};
+use console::style;
 use reqwest::{Client, Url};
 use serde_json::Value;
 use std::io::Write;
@@ -144,6 +145,7 @@ by Ben "epi" Risher {}                  ver: {}"#,
     let status = needs_update(&CONFIGURATION.client, UPDATE_URL, version).await;
 
     let top = "───────────────────────────┬──────────────────────";
+    let addl_section = "──────────────────────────────────────────────────";
     let bottom = "───────────────────────────┴──────────────────────";
 
     writeln!(&mut writer, "{}", artwork).unwrap_or_default();
@@ -433,6 +435,16 @@ by Ben "epi" Risher {}                  ver: {}"#,
     }
 
     writeln!(&mut writer, "{}", bottom).unwrap_or_default();
+    // ⏯
+    writeln!(
+        &mut writer,
+        " \u{23ef}  Press [{}] to {}|{} your scan",
+        style("ENTER").yellow(),
+        style("pause").red(),
+        style("resume").green()
+    )
+    .unwrap_or_default();
+    writeln!(&mut writer, "{}", addl_section).unwrap_or_default();
 }
 
 #[cfg(test)]

src/heuristics.rs

@@ -1,10 +1,13 @@
-use crate::config::{CONFIGURATION, PROGRESS_PRINTER};
-use crate::filters::WildcardFilter;
-use crate::scanner::should_filter_response;
-use crate::utils::{
-    ferox_print, format_url, get_url_path_length, make_request, module_colorizer, status_colorizer,
+use crate::{
+    config::{CONFIGURATION, PROGRESS_PRINTER},
+    filters::WildcardFilter,
+    scanner::should_filter_response,
+    utils::{
+        ferox_print, format_url, get_url_path_length, make_request, module_colorizer,
+        status_colorizer,
+    },
+    FeroxResponse,
 };
-use crate::FeroxResponse;
 use console::style;
 use indicatif::ProgressBar;
 use std::process;

src/lib.rs

@@ -11,8 +11,10 @@ pub mod reporter;
 pub mod scanner;
 pub mod utils;
 
-use reqwest::header::HeaderMap;
-use reqwest::{Response, StatusCode, Url};
+use reqwest::{
+    header::HeaderMap,
+    {Response, StatusCode, Url},
+};
 use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
 
 /// Generic Result type to ease error handling in async contexts
@@ -33,6 +35,9 @@ pub const VERSION: &str = env!("CARGO_PKG_VERSION");
 pub const DEFAULT_WORDLIST: &str =
     "/usr/share/seclists/Discovery/Web-Content/raft-medium-directories.txt";
 
+/// Number of milliseconds to wait between polls of `PAUSE_SCAN` when user pauses a scan
+pub static SLEEP_DURATION: u64 = 500;
+
 /// Default list of status codes to report
 ///
 /// * 200 Ok

src/main.rs

@@ -1,17 +1,59 @@
-use feroxbuster::config::{CONFIGURATION, PROGRESS_PRINTER};
-use feroxbuster::scanner::scan_url;
-use feroxbuster::utils::{ferox_print, get_current_depth, module_colorizer, status_colorizer};
-use feroxbuster::{banner, heuristics, logger, reporter, FeroxResponse, FeroxResult, VERSION};
+use crossterm::event::{self, Event, KeyCode};
+use feroxbuster::{
+    banner,
+    config::{CONFIGURATION, PROGRESS_PRINTER},
+    heuristics, logger, reporter,
+    scanner::{scan_url, PAUSE_SCAN},
+    utils::{ferox_print, get_current_depth, module_colorizer, status_colorizer},
+    FeroxResponse, FeroxResult, SLEEP_DURATION, VERSION,
+};
 use futures::StreamExt;
-use std::collections::HashSet;
-use std::fs::File;
-use std::io::{stderr, BufRead, BufReader};
-use std::process;
-use std::sync::Arc;
-use tokio::io;
-use tokio::sync::mpsc::UnboundedSender;
+use std::{
+    collections::HashSet,
+    fs::File,
+    io::{stderr, BufRead, BufReader},
+    process,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+    time::Duration,
+};
+use tokio::{io, sync::mpsc::UnboundedSender};
 use tokio_util::codec::{FramedRead, LinesCodec};
 
+/// Atomic boolean flag, used to determine whether or not the terminal input handler should exit
+pub static SCAN_COMPLETE: AtomicBool = AtomicBool::new(false);
+
+/// Handles specific key events triggered by the user over stdin
+fn terminal_input_handler() {
+    log::trace!("enter: terminal_input_handler");
+
+    loop {
+        if event::poll(Duration::from_millis(SLEEP_DURATION)).unwrap_or(false) {
+            // It's guaranteed that the `read()` won't block when the `poll()`
+            // function returns `true`
+
+            if let Ok(key_pressed) = event::read() {
+                if key_pressed == Event::Key(KeyCode::Enter.into()) {
+                    // if the user presses Enter, toggle the value stored in PAUSE_SCAN
+                    // ignore any other keys
+                    let current = PAUSE_SCAN.load(Ordering::Acquire);
+
+                    PAUSE_SCAN.store(!current, Ordering::Release);
+                }
+            }
+        } else {
+            // Timeout expired and no `Event` is available; use the timeout to check SCAN_COMPLETE
+            if SCAN_COMPLETE.load(Ordering::Relaxed) {
+                // scan has been marked complete by main, time to exit the loop
+                break;
+            }
+        }
+    }
+    log::trace!("exit: terminal_input_handler");
+}
+
 /// Create a HashSet of Strings from the given wordlist then stores it inside an Arc
 fn get_unique_words_from_wordlist(path: &str) -> FeroxResult<Arc<HashSet<String>>> {
     log::trace!("enter: get_unique_words_from_wordlist({})", path);
@@ -132,6 +174,11 @@ async fn main() {
     log::trace!("enter: main");
     log::debug!("{:#?}", *CONFIGURATION);
 
+    // spawn a thread that listens for keyboard input on stdin, when a user presses enter
+    // the input handler will toggle PAUSE_SCAN, which in turn is used to pause and resume
+    // scans that are already running
+    tokio::task::spawn_blocking(terminal_input_handler);
+
     let save_output = !CONFIGURATION.output.is_empty(); // was -o used?
 
     let (tx_term, tx_file, term_handle, file_handle) =
@@ -205,6 +252,9 @@ async fn main() {
         log::trace!("done awaiting file output handler's receiver");
     }
 
+    // mark all scans complete so the terminal input handler will exit cleanly
+    SCAN_COMPLETE.store(true, Ordering::Relaxed);
+
     log::trace!("exit: main");
 
     // clean-up function for the MultiProgress bar; must be called last in order to still see

src/scanner.rs

@@ -1,28 +1,49 @@
-use crate::config::{CONFIGURATION, PROGRESS_BAR};
-use crate::extractor::get_links;
-use crate::filters::{FeroxFilter, StatusCodeFilter, WildcardFilter};
-use crate::utils::{format_url, get_current_depth, make_request};
-use crate::{heuristics, progress, FeroxChannel, FeroxResponse};
-use futures::future::{BoxFuture, FutureExt};
-use futures::{stream, StreamExt};
+use crate::{
+    config::{CONFIGURATION, PROGRESS_BAR},
+    extractor::get_links,
+    filters::{FeroxFilter, StatusCodeFilter, WildcardFilter},
+    heuristics, progress,
+    utils::{format_url, get_current_depth, make_request},
+    FeroxChannel, FeroxResponse, SLEEP_DURATION,
+};
+use console::style;
+use futures::{
+    future::{BoxFuture, FutureExt},
+    stream, StreamExt,
+};
+use indicatif::{ProgressBar, ProgressStyle};
 use lazy_static::lazy_static;
 use reqwest::Url;
-use std::collections::HashSet;
-use std::convert::TryInto;
-use std::ops::Deref;
-use std::sync::atomic::{AtomicUsize, Ordering};
-use std::sync::{Arc, RwLock};
-use tokio::sync::mpsc::{self, UnboundedReceiver, UnboundedSender};
-use tokio::sync::Semaphore;
-use tokio::task::JoinHandle;
+use std::{
+    collections::HashSet,
+    convert::TryInto,
+    io::{stderr, Write},
+    ops::Deref,
+    sync::atomic::{AtomicBool, AtomicUsize, Ordering},
+    sync::{Arc, RwLock},
+};
+use tokio::{
+    sync::{
+        mpsc::{self, UnboundedReceiver, UnboundedSender},
+        Semaphore,
+    },
+    task::JoinHandle,
+    time,
+};
 
 /// Single atomic number that gets incremented once, used to track first scan vs. all others
 static CALL_COUNT: AtomicUsize = AtomicUsize::new(0);
 
+/// Atomic boolean flag, used to determine whether or not a scan should pause or resume
+pub static PAUSE_SCAN: AtomicBool = AtomicBool::new(false);
+
 lazy_static! {
     /// Set of urls that have been sent to [scan_url](fn.scan_url.html), used for deduplication
     static ref SCANNED_URLS: RwLock<HashSet<String>> = RwLock::new(HashSet::new());
 
+    /// A clock spinner protected with a RwLock to allow for a single thread to use at a time
+    static ref SINGLE_SPINNER: RwLock<ProgressBar> = RwLock::new(get_single_spinner());
+
     /// Vector of implementors of the FeroxFilter trait
     static ref FILTERS: Arc<RwLock<Vec<Box<dyn FeroxFilter>>>> = Arc::new(RwLock::new(Vec::<Box<dyn FeroxFilter>>::new()));
 
@@ -30,6 +51,72 @@ lazy_static! {
     static ref SCAN_LIMITER: Semaphore = Semaphore::new(CONFIGURATION.scan_limit);
 }
 
+/// Return a clock spinner, used when scans are paused
+fn get_single_spinner() -> ProgressBar {
+    log::trace!("enter: get_single_spinner");
+
+    let spinner = ProgressBar::new_spinner().with_style(
+        ProgressStyle::default_spinner()
+            .tick_strings(&[
+                "🕛", "🕐", "🕑", "🕒", "🕓", "🕔", "🕕", "🕖", "🕗", "🕘", "🕙", "🕚",
+            ])
+            .template(&format!(
+                "\t-= All Scans {{spinner}} {} =-",
+                style("Paused").red()
+            )),
+    );
+
+    log::trace!("exit: get_single_spinner -> {:?}", spinner);
+    spinner
+}
+
+/// Forced the calling thread into a busy loop
+///
+/// Every `SLEEP_DURATION` milliseconds, the function examines the result stored in `PAUSE_SCAN`
+///
+/// When the value stored in `PAUSE_SCAN` becomes `false`, the function returns, exiting the busy
+/// loop
+async fn pause_scan() {
+    log::trace!("enter: pause_scan");
+    // function uses tokio::time, not std
+
+    // local testing showed a pretty slow increase (less than linear) in CPU usage as # of
+    // concurrent scans rose when SLEEP_DURATION was set to 500, using that as the default for now
+    let mut interval = time::interval(time::Duration::from_millis(SLEEP_DURATION));
+
+    // ignore any error returned
+    let _ = stderr().flush();
+
+    if SINGLE_SPINNER.read().unwrap().is_finished() {
+        // in order to not leave draw artifacts laying around in the terminal, we call
+        // finish_and_clear on the progress bar when resuming scans. For this reason, we need to
+        // check if the spinner is finished, and repopulate the RwLock with a new spinner if
+        // necessary
+        if let Ok(mut guard) = SINGLE_SPINNER.write() {
+            *guard = get_single_spinner();
+        }
+    }
+
+    if let Ok(spinner) = SINGLE_SPINNER.write() {
+        spinner.enable_steady_tick(120);
+    }
+
+    loop {
+        // first tick happens immediately, all others wait the specified duration
+        interval.tick().await;
+
+        if !PAUSE_SCAN.load(Ordering::Acquire) {
+            // PAUSE_SCAN is false, so we can exit the busy loop
+            if let Ok(spinner) = SINGLE_SPINNER.write() {
+                spinner.finish_and_clear();
+            }
+            let _ = stderr().flush();
+            log::trace!("exit: pause_scan");
+            return;
+        }
+    }
+}
+
 /// Adds the given url to `SCANNED_URLS`
 ///
 /// If `SCANNED_URLS` did not already contain the url, return true; otherwise return false
@@ -599,7 +686,15 @@ pub async fn scan_url(
             let pb = progress_bar.clone(); // progress bar is an Arc around internal state
             let tgt = target_url.to_string(); // done to satisfy 'static lifetime below
             (
-                tokio::spawn(async move { make_requests(&tgt, &word, base_depth, txd, txr).await }),
+                tokio::spawn(async move {
+                    if PAUSE_SCAN.load(Ordering::Acquire) {
+                        // for every word in the wordlist, check to see if PAUSE_SCAN is set to true
+                        // when true; enter a busy loop that only exits by setting PAUSE_SCAN back
+                        // to false
+                        pause_scan().await;
+                    }
+                    make_requests(&tgt, &word, base_depth, txd, txr).await
+                }),
                 pb,
             )
         })
@@ -768,4 +863,33 @@ mod tests {
 
         assert_eq!(add_url_to_list_of_scanned_urls(url, &urls), false);
     }
+
+    #[test]
+    /// test that get_single_spinner returns the correct spinner
+    fn scanner_get_single_spinner_returns_spinner() {
+        let spinner = get_single_spinner();
+        assert!(!spinner.is_finished());
+    }
+
+    #[tokio::test(core_threads = 1)]
+    /// tests that pause_scan pauses execution and releases execution when PAUSE_SCAN is toggled
+    /// the spinner used during the test has had .finish_and_clear called on it, meaning that
+    /// a new one will be created, taking the if branch within the function
+    async fn scanner_pause_scan_with_finished_spinner() {
+        let now = time::Instant::now();
+
+        PAUSE_SCAN.store(true, Ordering::Relaxed);
+        SINGLE_SPINNER.write().unwrap().finish_and_clear();
+
+        let expected = time::Duration::from_secs(2);
+
+        tokio::spawn(async move {
+            time::delay_for(expected).await;
+            PAUSE_SCAN.store(false, Ordering::Relaxed);
+        });
+
+        pause_scan().await;
+
+        assert!(now.elapsed() > expected);
+    }
 }