implemented pause|resume functionality

Cargo.toml

@@ -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"

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() {
+            // 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,
             )
         })