From 004a045da28f4071046d0304e5b72d8d70bb3897 Mon Sep 17 00:00:00 2001 From: epi Date: Sat, 30 Jan 2021 08:26:33 -0600 Subject: [PATCH] added terminal input event handlers --- src/event_handlers/inputs.rs | 147 +++++++++++++++++++++++++++++++++++ src/event_handlers/mod.rs | 2 + src/main.rs | 57 ++------------ src/scan_manager.rs | 104 +++++++------------------ tests/test_scanner.rs | 4 +- 5 files changed, 187 insertions(+), 127 deletions(-) create mode 100644 src/event_handlers/inputs.rs diff --git a/src/event_handlers/inputs.rs b/src/event_handlers/inputs.rs new file mode 100644 index 0000000..ab92900 --- /dev/null +++ b/src/event_handlers/inputs.rs @@ -0,0 +1,147 @@ +use super::*; +use crate::{ + config::PROGRESS_PRINTER, + scan_manager::FeroxState, + scan_manager::PAUSE_SCAN, + scanner::RESPONSES, + utils::{open_file, write_to}, + SLEEP_DURATION, +}; +use anyhow::Result; +use console::style; +use crossterm::event::{self, Event, KeyCode}; +use std::{ + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, + thread::sleep, + time::Duration, + time::{SystemTime, UNIX_EPOCH}, +}; + +/// Atomic boolean flag, used to determine whether or not the terminal input handler should exit +pub static SCAN_COMPLETE: AtomicBool = AtomicBool::new(false); + +#[derive(Debug)] +pub struct TerminalInputHandle {} + +#[derive(Debug)] +/// Container for filters transmitter and FeroxFilters object +pub struct TermInputHandler { + /// handles to other handlers + handles: Arc, +} + +/// implementation of event handler for terminal input (used for cancel scan menu) todo update for correctness /// Initialize the ctrl+c handler that saves scan state to disk +impl TermInputHandler { + /// Create new event handler + pub fn new(handles: Arc) -> Self { + Self { handles } + } + + /// Initialize the sc side of an mpsc channel that is responsible for updates todo update for correctness + pub fn initialize(handles: Arc) { + log::trace!("enter: initialize({:?})", handles); + + let handler = Self::new(handles); + handler.start(); + + log::trace!("exit: initialize"); + } + + fn start(&self) { + tokio::task::spawn_blocking(Self::start_enter_handler); + + let cloned = self.handles.clone(); + + if self.handles.config.save_state { + // start the ctrl+c handler + let result = ctrlc::set_handler(move || { + let _ = Self::sigint_handler(cloned.clone()); + }); + + if result.is_err() { + log::error!("Could not set Ctrl+c handler"); + std::process::exit(1); + } + } + } + + /// Writes the current state of the program to disk (if save_state is true) and then exits + pub fn sigint_handler(handles: Arc) -> Result<()> { + log::trace!("enter: sigint_handler({:?})", handles); + + let ts = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs(); + + let slug = if !handles.config.target_url.is_empty() { + // target url populated + handles + .config + .target_url + .replace("://", "_") + .replace("/", "_") + .replace(".", "_") + } else { + // stdin used + "stdin".to_string() + }; + + let filename = format!("ferox-{}-{}.state", slug, ts); + let warning = format!( + "🚨 Caught {} 🚨 saving scan state to {} ...", + style("ctrl+c").yellow(), + filename + ); + + PROGRESS_PRINTER.println(warning); + + let state = FeroxState::new( + handles.ferox_scans()?, + handles.config.clone(), + &RESPONSES, + handles.stats.data.clone(), + ); + + let state_file = open_file(&filename); + + let mut buffered_file = state_file?; + write_to(&state, &mut buffered_file, true)?; + + log::trace!("exit: sigint_handler (end of program)"); + std::process::exit(1); + } + + /// Handles specific key events triggered by the user over stdin + fn start_enter_handler() { + // todo eventually move away from atomics, the blocking recv is the problem + log::trace!("enter: start"); + + loop { + if PAUSE_SCAN.load(Ordering::Relaxed) { + // if the scan is already paused, we don't want this event poller fighting the user + // over stdin + sleep(Duration::from_millis(SLEEP_DURATION)); + } else 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() { + // ignore any other keys + if key_pressed == Event::Key(KeyCode::Enter.into()) { + // if the user presses Enter, set PAUSE_SCAN to true. The interactive menu + // will be triggered and will handle setting PAUSE_SCAN to false + PAUSE_SCAN.store(true, 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: start"); + } +} diff --git a/src/event_handlers/mod.rs b/src/event_handlers/mod.rs index f5a19a1..aa85b6d 100644 --- a/src/event_handlers/mod.rs +++ b/src/event_handlers/mod.rs @@ -5,10 +5,12 @@ mod container; mod command; mod outputs; mod scans; +mod inputs; pub use self::command::Command; pub use self::container::{Handles, Tasks}; pub use self::filters::{FiltersHandle, FiltersHandler}; +pub use self::inputs::{TermInputHandler, SCAN_COMPLETE}; pub use self::outputs::{TermOutHandle, TermOutHandler}; pub use self::scans::{ScanHandle, ScanHandler}; pub use self::statistics::{StatsHandle, StatsHandler}; diff --git a/src/main.rs b/src/main.rs index 443508f..4b62b9f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,16 +2,10 @@ use std::{ collections::HashSet, fs::File, io::{stderr, BufRead, BufReader}, - sync::{ - atomic::{AtomicBool, Ordering}, - Arc, - }, - thread::sleep, - time::Duration, + sync::{atomic::Ordering, Arc}, }; use anyhow::{bail, Context, Result}; -use crossterm::event::{self, Event, KeyCode}; use futures::StreamExt; use tokio::{io, sync::oneshot}; use tokio_util::codec::{FramedRead, LinesCodec}; @@ -21,52 +15,17 @@ use feroxbuster::{ config::{Configuration, PROGRESS_BAR, PROGRESS_PRINTER}, event_handlers::{ Command::{CreateBar, Exit, JoinTasks, LoadStats, ScanInitialUrls, UpdateWordlist}, - FiltersHandler, Handles, ScanHandler, StatsHandler, Tasks, TermOutHandler, + FiltersHandler, Handles, ScanHandler, StatsHandler, Tasks, TermInputHandler, + TermOutHandler, SCAN_COMPLETE, }, heuristics, logger, - scan_manager::{self, PAUSE_SCAN}, + scan_manager::{self}, scanner, utils::fmt_err, - SLEEP_DURATION, }; #[cfg(not(target_os = "windows"))] use feroxbuster::{utils::set_open_file_limit, DEFAULT_OPEN_FILE_LIMIT}; -/// Atomic boolean flag, used to determine whether or not the terminal input handler should exit -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 PAUSE_SCAN.load(Ordering::Relaxed) { - // if the scan is already paused, we don't want this event poller fighting the user - // over stdin - sleep(Duration::from_millis(SLEEP_DURATION)); - } else 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() { - // ignore any other keys - if key_pressed == Event::Key(KeyCode::Enter.into()) { - // if the user presses Enter, set PAUSE_SCAN to true. The interactive menu - // will be triggered and will handle setting PAUSE_SCAN to false - PAUSE_SCAN.store(true, 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) -> Result>> { log::trace!("enter: get_unique_words_from_wordlist({})", path); @@ -236,12 +195,8 @@ async fn wrapped_main(config: Arc) -> Result<()> { // 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); - - if config.save_state { - // start the ctrl+c handler - scan_manager::initialize(handles.clone()); - } + // also starts ctrl+c handler + TermInputHandler::initialize(handles.clone()); if config.resumed { let scanned_urls = handles.ferox_scans()?; diff --git a/src/scan_manager.rs b/src/scan_manager.rs index 65b0197..a43062f 100644 --- a/src/scan_manager.rs +++ b/src/scan_manager.rs @@ -1,6 +1,7 @@ use std::{ cmp::PartialEq, collections::HashMap, + convert::TryInto, fmt, fs::File, io::BufReader, @@ -8,7 +9,6 @@ use std::{ sync::atomic::{AtomicBool, AtomicUsize, Ordering}, sync::{Arc, Mutex, RwLock}, thread::sleep, - time::{SystemTime, UNIX_EPOCH}, }; use anyhow::{Context, Result}; @@ -26,19 +26,16 @@ use tokio::{ }; use uuid::Uuid; -use crate::event_handlers::Handles; -use crate::utils::fmt_err; -use crate::utils::write_to; use crate::{ - config::{Configuration, CONFIGURATION, PROGRESS_BAR, PROGRESS_PRINTER}, + config::{Configuration, PROGRESS_BAR, PROGRESS_PRINTER}, + event_handlers::Handles, parser::TIMESPEC_REGEX, progress::{add_bar, BarType}, scanner::RESPONSES, statistics::Stats, - utils::open_file, + utils::fmt_err, FeroxResponse, FeroxSerialize, SLEEP_DURATION, }; -use std::convert::TryInto; /// Single atomic number that gets incremented once, used to track first thread to interact with /// when pausing a scan @@ -939,7 +936,7 @@ pub struct FeroxState { scans: Arc, /// Current running config - config: &'static Configuration, + config: Arc, /// Known responses responses: &'static FeroxResponses, @@ -948,6 +945,24 @@ pub struct FeroxState { statistics: Arc, } +/// implementation of FeroxState +impl FeroxState { + /// create new FeroxState object + pub fn new( + scans: Arc, + config: Arc, + responses: &'static FeroxResponses, + statistics: Arc, + ) -> Self { + Self { + scans, + config, + responses, + statistics, + } + } +} + /// FeroxSerialize implementation for FeroxState impl FeroxSerialize for FeroxState { /// Simply return debug format of FeroxState to satisfy as_str @@ -998,7 +1013,7 @@ pub async fn start_max_time_thread(handles: Arc) { #[cfg(test)] panic!(handles); #[cfg(not(test))] - let _ = sigint_handler(handles.clone()); + let _ = crate::event_handlers::TermInputHandler::sigint_handler(handles.clone()); } log::error!( @@ -1007,65 +1022,6 @@ pub async fn start_max_time_thread(handles: Arc) { ); } -/// Writes the current state of the program to disk (if save_state is true) and then exits -fn sigint_handler(handles: Arc) -> Result<()> { - log::trace!("enter: sigint_handler({:?})", handles); - - let ts = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs(); - - let slug = if !CONFIGURATION.target_url.is_empty() { - // target url populated - CONFIGURATION - .target_url - .replace("://", "_") - .replace("/", "_") - .replace(".", "_") - } else { - // stdin used - "stdin".to_string() - }; - - let filename = format!("ferox-{}-{}.state", slug, ts); - let warning = format!( - "🚨 Caught {} 🚨 saving scan state to {} ...", - style("ctrl+c").yellow(), - filename - ); - - PROGRESS_PRINTER.println(warning); - - let state = FeroxState { - config: &CONFIGURATION, - scans: handles.ferox_scans()?, - responses: &RESPONSES, - statistics: handles.stats.data.clone(), - }; - - let state_file = open_file(&filename); - - let mut buffered_file = state_file?; - write_to(&state, &mut buffered_file, true)?; - - log::trace!("exit: sigint_handler (end of program)"); - std::process::exit(1); -} - -/// Initialize the ctrl+c handler that saves scan state to disk -pub fn initialize(handles: Arc) { - log::trace!("enter: initialize({:?})", handles); - - let result = ctrlc::set_handler(move || { - let _ = sigint_handler(handles.clone()); - }); - - if result.is_err() { - log::error!("Could not set Ctrl+c handler"); - std::process::exit(1); - } - - log::trace!("exit: initialize"); -} - /// Primary logic used to load a Configuration from disk and populate the appropriate data /// structures pub fn resume_scan(filename: &str) -> Configuration { @@ -1433,12 +1389,12 @@ mod tests { let response: FeroxResponse = serde_json::from_str(json_response).unwrap(); RESPONSES.insert(response); - let ferox_state = FeroxState { - scans: Arc::new(ferox_scans), - responses: &RESPONSES, - config: &CONFIGURATION, - statistics: stats, - }; + let ferox_state = FeroxState::new( + Arc::new(ferox_scans), + Arc::new(Configuration::new().unwrap()), + &RESPONSES, + stats, + ); let expected_strs = predicates::str::contains("scans: FeroxScans").and( predicate::str::contains("config: Configuration") diff --git a/tests/test_scanner.rs b/tests/test_scanner.rs index 173a5fe..e5207f4 100644 --- a/tests/test_scanner.rs +++ b/tests/test_scanner.rs @@ -458,7 +458,7 @@ fn scanner_single_request_scan_with_debug_logging() { assert!(contents.contains("DBG")); assert!(contents.contains("INF")); assert!(contents.contains("feroxbuster All scans complete!")); - assert!(contents.contains("feroxbuster exit: terminal_input_handler")); + assert!(contents.contains("feroxbuster::event_handlers::inputs exit: start")); assert_eq!(mock.hits(), 1); teardown_tmp_directory(tmp_dir); @@ -497,8 +497,8 @@ fn scanner_single_request_scan_with_debug_logging_as_json() { assert!(contents.contains("\"level\":\"INFO\"")); assert!(contents.contains("time_offset")); assert!(contents.contains("\"module\":\"feroxbuster::scanner\"")); + assert!(contents.contains("\"module\":\"feroxbuster::event_handlers::inputs\"")); assert!(contents.contains("All scans complete!")); - assert!(contents.contains("exit: terminal_input_handler")); assert_eq!(mock.hits(), 1); teardown_tmp_directory(tmp_dir);