diff --git a/src/event_handlers/scans.rs b/src/event_handlers/scans.rs index 5037e13d..b3640ff0 100644 --- a/src/event_handlers/scans.rs +++ b/src/event_handlers/scans.rs @@ -120,7 +120,10 @@ impl ScanHandler { pub fn initialize(handles: Arc) -> (Joiner, ScanHandle) { log::trace!("enter: initialize"); - let data = Arc::new(FeroxScans::new(handles.config.output_level, handles.config.limit_bars)); + let data = Arc::new(FeroxScans::new( + handles.config.output_level, + handles.config.limit_bars, + )); let (tx, rx): FeroxChannel = mpsc::unbounded_channel(); let max_depth = handles.config.depth; @@ -322,7 +325,9 @@ impl ScanHandler { let scan = if let Some(ferox_scan) = self.data.get_scan_by_url(&target) { ferox_scan // scan already known } else { - self.data.add_directory_scan(&target, order).1 // add the new target; return FeroxScan + self.data + .add_directory_scan(&target, order, self.handles.clone()) + .1 // add the new target; return FeroxScan }; if should_test_deny diff --git a/src/extractor/container.rs b/src/extractor/container.rs index ad9f4f15..70871c8e 100644 --- a/src/extractor/container.rs +++ b/src/extractor/container.rs @@ -228,8 +228,11 @@ impl<'a> Extractor<'a> { if resp.is_file() || !resp.is_directory() { log::debug!("Extracted File: {}", resp); - c_scanned_urls - .add_file_scan(resp.url().as_str(), ScanOrder::Latest); + c_scanned_urls.add_file_scan( + resp.url().as_str(), + ScanOrder::Latest, + c_handles.clone(), + ); if c_handles.config.collect_extensions { // no real reason this should fail diff --git a/src/extractor/tests.rs b/src/extractor/tests.rs index 68221f59..ae541218 100644 --- a/src/extractor/tests.rs +++ b/src/extractor/tests.rs @@ -4,8 +4,8 @@ use super::*; use crate::config::{Configuration, OutputLevel}; use crate::scan_manager::ScanOrder; use crate::{ - event_handlers::Handles, scan_manager::FeroxScans, utils::make_request, Command, FeroxChannel, - DEFAULT_METHOD, + event_handlers::Handles, event_handlers::Handles, scan_manager::FeroxScans, + utils::make_request, Command, FeroxChannel, DEFAULT_METHOD, }; use anyhow::Result; use httpmock::{Method::GET, MockServer}; @@ -386,7 +386,11 @@ async fn request_link_bails_on_seen_url() -> Result<()> { }); let scans = Arc::new(FeroxScans::default()); - scans.add_file_scan(&served, ScanOrder::Latest); + scans.add_file_scan( + &served, + ScanOrder::Latest, + Arc::new(Handles::for_testing(None, None)), + ); let robots = setup_extractor(ExtractionTarget::RobotsTxt, scans.clone()); let body = setup_extractor(ExtractionTarget::ResponseBody, scans); diff --git a/src/scan_manager/scan.rs b/src/scan_manager/scan.rs index 398aea07..78644efd 100644 --- a/src/scan_manager/scan.rs +++ b/src/scan_manager/scan.rs @@ -1,8 +1,10 @@ use super::*; use crate::{ config::OutputLevel, + event_handlers::Handles, progress::update_style, progress::{add_bar, BarType}, + scan_manager::utils::determine_bar_type, scanner::PolicyTrigger, }; use anyhow::Result; @@ -69,7 +71,7 @@ pub struct FeroxScan { pub(super) task: sync::Mutex>>, /// The progress bar associated with this scan - pub(super) progress_bar: Mutex>, + pub progress_bar: Mutex>, /// whether or not the user passed --silent|--quiet on the command line pub(super) output_level: OutputLevel, @@ -89,8 +91,8 @@ pub struct FeroxScan { /// whether the progress bar is currently visible or hidden pub(super) visible: AtomicBool, - /// stored value for Configuration.limit_bars - pub(super) bar_limit: usize, + /// handles object pointer + pub(super) handles: Option>, } /// Default implementation for FeroxScan @@ -103,7 +105,7 @@ impl Default for FeroxScan { id: new_id, task: sync::Mutex::new(None), // tokio mutex status: Mutex::new(ScanStatus::default()), - bar_limit: 0, + handles: None, num_requests: 0, requests_made_so_far: 0, scan_order: ScanOrder::Latest, @@ -128,29 +130,37 @@ impl FeroxScan { self.visible.load(Ordering::Relaxed) } - /// return the visibility of the scan as a boolean - pub fn running(&self) -> bool { - self.visible.load(Ordering::Relaxed) - } - - pub fn swap_visibility(&self, bar_type: BarType) { + pub fn swap_visibility(&self) { // fetch_xor toggles the boolean to its opposite and returns the previous value let visible = self.visible.fetch_xor(true, Ordering::Relaxed); - if !visible { - // visibility was false before we xor'd the value + let Ok(bar) = self.progress_bar.lock() else { + log::warn!("couldn't unlock progress bar for {}", self.url); + return; + }; - let Ok(bar) = self.progress_bar.lock() else { - log::warn!("couldn't unlock progress bar for {}", self.url); - return; - }; + let Some(handles) = self.handles.as_ref() else { + log::warn!("couldn't access handles pointer for {}", self.url); + return; + }; - update_style(bar.as_ref().unwrap(), bar_type); - } + let bar_type = if !visible { + // visibility was false before we xor'd the value + match handles.config.output_level { + OutputLevel::Default => BarType::Default, + OutputLevel::Quiet => BarType::Quiet, + OutputLevel::Silent | OutputLevel::SilentJSON => BarType::Hidden, + } + } else { + // visibility was true before we xor'd the value + BarType::Hidden + }; + + update_style(bar.as_ref().unwrap(), bar_type); } /// Stop a currently running scan - pub async fn abort(&self) -> Result<()> { + pub async fn abort(&self, active_bars: usize) -> Result<()> { log::trace!("enter: abort"); match self.task.try_lock() { @@ -159,8 +169,7 @@ impl FeroxScan { log::trace!("aborting {:?}", self); task.abort(); self.set_status(ScanStatus::Cancelled)?; - // todo: is this right? check it - self.stop_progress_bar(0); + self.stop_progress_bar(active_bars); } } Err(e) => { @@ -202,7 +211,13 @@ impl FeroxScan { if guard.is_some() { let pb = (*guard).as_ref().unwrap(); - if self.bar_limit > 0 && self.bar_limit < active_bars + 1 { + let bar_limit = if let Some(handles) = self.handles.as_ref() { + handles.config.limit_bars + } else { + 0 + }; + + if bar_limit > 0 && bar_limit < active_bars { pb.finish_and_clear(); return; } @@ -223,12 +238,18 @@ impl FeroxScan { if guard.is_some() { (*guard).as_ref().unwrap().clone() } else { - let bar_type = match self.output_level { - OutputLevel::Default => BarType::Default, - OutputLevel::Quiet => BarType::Quiet, - OutputLevel::Silent | OutputLevel::SilentJSON => BarType::Hidden, + let (active_bars, bar_limit) = if let Some(handles) = self.handles.as_ref() { + if let Ok(scans) = handles.ferox_scans() { + (scans.number_of_bars(), handles.config.limit_bars) + } else { + (0, handles.config.limit_bars) + } + } else { + (0, 0) }; - // todo: add-bar check for limit/current + + let bar_type = determine_bar_type(bar_limit, active_bars, self.output_level); + let pb = add_bar(&self.url, self.num_requests, bar_type); pb.reset_elapsed(); @@ -242,12 +263,18 @@ impl FeroxScan { Err(_) => { log::warn!("Could not unlock progress bar on {:?}", self); - let bar_type = match self.output_level { - OutputLevel::Default => BarType::Default, - OutputLevel::Quiet => BarType::Quiet, - OutputLevel::Silent | OutputLevel::SilentJSON => BarType::Hidden, + let (active_bars, bar_limit) = if let Some(handles) = self.handles.as_ref() { + if let Ok(scans) = handles.ferox_scans() { + (scans.number_of_bars(), handles.config.limit_bars) + } else { + (0, handles.config.limit_bars) + } + } else { + (0, 0) }; - // todo: add-bar check for limit/current + + let bar_type = determine_bar_type(bar_limit, active_bars, self.output_level); + let pb = add_bar(&self.url, self.num_requests, bar_type); pb.reset_elapsed(); @@ -266,7 +293,7 @@ impl FeroxScan { output_level: OutputLevel, pb: Option, visibility: bool, - bar_limit: usize, + handles: Arc, ) -> Arc { Arc::new(Self { url: url.to_string(), @@ -275,17 +302,17 @@ impl FeroxScan { scan_order, num_requests, output_level, - bar_limit, progress_bar: Mutex::new(pb), visible: AtomicBool::new(visibility), + handles: Some(handles), ..Default::default() }) } /// Mark the scan as complete and stop the scan's progress bar - pub fn finish(&self, n: usize) -> Result<()> { + pub fn finish(&self, active_bars: usize) -> Result<()> { self.set_status(ScanStatus::Complete)?; - self.stop_progress_bar(n); + self.stop_progress_bar(active_bars); Ok(()) } @@ -604,7 +631,6 @@ mod tests { normalized_url: String::from("/"), scan_type: ScanType::Directory, scan_order: ScanOrder::Initial, - bar_limit: 0, num_requests: 0, requests_made_so_far: 0, visible: AtomicBool::new(true), @@ -630,4 +656,32 @@ mod tests { scan.finish(0).unwrap(); assert_eq!(scan.requests_per_second(), 0); } + + #[test] + fn test_swap_visibility() { + let scan = FeroxScan::new( + "http://localhost", + ScanType::Directory, + ScanOrder::Latest, + 1000, + OutputLevel::Default, + None, + true, + 0, + ); + + assert!(scan.visible()); + + scan.swap_visibility(); + assert!(scan.visible()); + + scan.swap_visibility(); + assert!(!scan.visible()); + + scan.swap_visibility(); + assert!(!scan.visible()); + + scan.swap_visibility(); + assert!(scan.visible()); + } } diff --git a/src/scan_manager/scan_container.rs b/src/scan_manager/scan_container.rs index e3e651ac..d2a078e8 100644 --- a/src/scan_manager/scan_container.rs +++ b/src/scan_manager/scan_container.rs @@ -12,6 +12,7 @@ use crate::{ config::OutputLevel, progress::PROGRESS_PRINTER, progress::{add_bar, BarType}, + scan_manager::utils::determine_bar_type, scan_manager::{MenuCmd, MenuCmdResult}, scanner::RESPONSES, traits::FeroxSerialize, @@ -20,7 +21,6 @@ use crate::{ use anyhow::Result; use console::style; use reqwest::StatusCode; -use scan::Visibility; use serde::{ser::SerializeSeq, Serialize, Serializer}; use std::{ collections::HashSet, @@ -393,8 +393,9 @@ impl FeroxScans { if input == 'y' || input == '\n' { self.menu.println(&format!("Stopping {}...", selected.url)); + let active_bars = self.number_of_bars(); selected - .abort() + .abort(active_bars) .await .unwrap_or_else(|e| log::warn!("Could not cancel task: {}", e)); @@ -526,18 +527,24 @@ impl FeroxScans { /// if a resumed scan is already complete, display a completed progress bar to the user pub fn print_completed_bars(&self, bar_length: usize) -> Result<()> { - let bar_type = match self.output_level { - OutputLevel::Default => BarType::Message, - OutputLevel::Quiet => BarType::Quiet, - OutputLevel::Silent | OutputLevel::SilentJSON => return Ok(()), // fast exit when --silent was used - }; + if self.output_level == OutputLevel::SilentJSON || self.output_level == OutputLevel::Silent + { + // fast exit when --silent was used + return Ok(()); + } + + let bar_type: BarType = + determine_bar_type(self.bar_limit, self.number_of_bars(), self.output_level); if let Ok(scans) = self.scans.read() { for scan in scans.iter() { + if matches!(bar_type, BarType::Hidden) { + // no need to show hidden bars + continue; + } + if scan.is_complete() { // these scans are complete, and just need to be shown to the user - // todo: this may change if i get the limited view working - // todo: add-bar check for limit/current let pb = add_bar( &scan.url, bar_length.try_into().unwrap_or_default(), @@ -602,31 +609,6 @@ impl FeroxScans { } } - /// determine the type of progress bar to display - /// takes both --limit-bars and output-level (--quiet|--silent|etc) - /// into account to arrive at a `BarType` - fn determine_bar_type(&self) -> BarType { - let visibility = if self.bar_limit == 0 { - // no limit from cli, just set the value to visible - // this protects us from a mutex unlock in number_of_bars - // in the normal case - Visibility::Visible - } else if self.bar_limit < self.number_of_bars() { - // active bars exceed limit; hidden - Visibility::Hidden - } else { - Visibility::Visible - }; - - match (self.output_level, visibility) { - (OutputLevel::Default, Visibility::Visible) => BarType::Default, - (OutputLevel::Quiet, Visibility::Visible) => BarType::Quiet, - (OutputLevel::Default, Visibility::Hidden) => BarType::Hidden, - (OutputLevel::Quiet, Visibility::Hidden) => BarType::Hidden, - (OutputLevel::Silent | OutputLevel::SilentJSON, _) => BarType::Hidden, - } - } - /// Given a url, create a new `FeroxScan` and add it to `FeroxScans` /// /// If `FeroxScans` did not already contain the scan, return true; otherwise return false @@ -637,6 +619,7 @@ impl FeroxScans { url: &str, scan_type: ScanType, scan_order: ScanOrder, + handles: Arc, ) -> (bool, Arc) { let bar_length = if let Ok(guard) = self.bar_length.lock() { *guard @@ -644,7 +627,8 @@ impl FeroxScans { 0 }; - let bar_type = self.determine_bar_type(); + let active_bars = self.number_of_bars(); + let bar_type = determine_bar_type(self.bar_limit, active_bars, self.output_level); let bar = match scan_type { ScanType::Directory => { @@ -667,7 +651,7 @@ impl FeroxScans { self.output_level, bar, is_visible, - self.bar_limit, + handles, ); // If the set did not contain the scan, true is returned. @@ -682,9 +666,14 @@ impl FeroxScans { /// If `FeroxScans` did not already contain the scan, return true; otherwise return false /// /// Also return a reference to the new `FeroxScan` - pub fn add_directory_scan(&self, url: &str, scan_order: ScanOrder) -> (bool, Arc) { + pub fn add_directory_scan( + &self, + url: &str, + scan_order: ScanOrder, + handles: Arc, + ) -> (bool, Arc) { let normalized = format!("{}/", url.trim_end_matches('/')); - self.add_scan(&normalized, ScanType::Directory, scan_order) + self.add_scan(&normalized, ScanType::Directory, scan_order, handles) } /// Given a url, create a new `FeroxScan` and add it to `FeroxScans` as a File Scan @@ -692,8 +681,13 @@ impl FeroxScans { /// If `FeroxScans` did not already contain the scan, return true; otherwise return false /// /// Also return a reference to the new `FeroxScan` - pub fn add_file_scan(&self, url: &str, scan_order: ScanOrder) -> (bool, Arc) { - self.add_scan(url, ScanType::File, scan_order) + pub fn add_file_scan( + &self, + url: &str, + scan_order: ScanOrder, + handles: Arc, + ) -> (bool, Arc) { + self.add_scan(url, ScanType::File, scan_order, handles) } /// returns the number of active AND visible scans; supports --limit-bars functionality @@ -702,7 +696,9 @@ impl FeroxScans { return 0; }; - let mut count = 0; + // starting at one ensures we don't have an extra bar + // due to counting up from 0 when there's actually 1 bar + let mut count = 1; for scan in &*scans { if scan.is_active() && scan.visible() { @@ -731,8 +727,7 @@ impl FeroxScans { } if scan.is_running() { - let bar_type = self.determine_bar_type(); - scan.swap_visibility(bar_type); + scan.swap_visibility(); return; } @@ -742,8 +737,7 @@ impl FeroxScans { } if let Some(scan) = queued { - let bar_type = self.determine_bar_type(); - scan.swap_visibility(bar_type); + scan.swap_visibility(); } } } diff --git a/src/scan_manager/tests.rs b/src/scan_manager/tests.rs index b11dbfb8..ece23934 100644 --- a/src/scan_manager/tests.rs +++ b/src/scan_manager/tests.rs @@ -655,7 +655,7 @@ async fn ferox_scan_abort() { errors: Default::default(), }; - scan.abort().await.unwrap(); + scan.abort(0).await.unwrap(); assert!(matches!( *scan.status.lock().unwrap(), diff --git a/src/scan_manager/utils.rs b/src/scan_manager/utils.rs index 18449576..1c500ac8 100644 --- a/src/scan_manager/utils.rs +++ b/src/scan_manager/utils.rs @@ -1,7 +1,12 @@ #[cfg(not(test))] use crate::event_handlers::TermInputHandler; use crate::{ - config::Configuration, event_handlers::Handles, parser::TIMESPEC_REGEX, scanner::RESPONSES, + config::{Configuration, OutputLevel}, + event_handlers::Handles, + parser::TIMESPEC_REGEX, + progress::BarType, + scan_manager::scan::Visibility, + scanner::RESPONSES, }; use std::{fs::File, io::BufReader, sync::Arc}; @@ -90,3 +95,32 @@ pub fn resume_scan(filename: &str) -> Configuration { log::trace!("exit: resume_scan -> {:?}", config); config } + +/// determine the type of progress bar to display +/// takes both --limit-bars and output-level (--quiet|--silent|etc) +/// into account to arrive at a `BarType` +pub fn determine_bar_type( + bar_limit: usize, + number_of_bars: usize, + output_level: OutputLevel, +) -> BarType { + let visibility = if bar_limit == 0 { + // no limit from cli, just set the value to visible + // this protects us from a mutex unlock in number_of_bars + // in the normal case + Visibility::Visible + } else if bar_limit < number_of_bars { + // active bars exceed limit; hidden + Visibility::Hidden + } else { + Visibility::Visible + }; + + match (output_level, visibility) { + (OutputLevel::Default, Visibility::Visible) => BarType::Default, + (OutputLevel::Quiet, Visibility::Visible) => BarType::Quiet, + (OutputLevel::Default, Visibility::Hidden) => BarType::Hidden, + (OutputLevel::Quiet, Visibility::Hidden) => BarType::Hidden, + (OutputLevel::Silent | OutputLevel::SilentJSON, _) => BarType::Hidden, + } +} diff --git a/src/scanner/ferox_scanner.rs b/src/scanner/ferox_scanner.rs index 980c86e5..7959f606 100644 --- a/src/scanner/ferox_scanner.rs +++ b/src/scanner/ferox_scanner.rs @@ -11,7 +11,7 @@ use tokio::sync::Semaphore; use crate::filters::{create_similarity_filter, EmptyFilter, SimilarityFilter}; use crate::heuristics::WildcardResult; - use crate::Command::AddFilter; +use crate::Command::AddFilter; use crate::{ event_handlers::{ Command::{AddError, AddToF64Field, AddToUsizeField, SubtractFromUsizeField}, @@ -289,8 +289,6 @@ impl FeroxScanner { " (add {} to scan)", style("--scan-dir-listings").bright().yellow() )?; - } else { - // todo: need to not skip them } if !self.handles.config.extract_links { diff --git a/src/scanner/requester.rs b/src/scanner/requester.rs index b8c39e02..732344ec 100644 --- a/src/scanner/requester.rs +++ b/src/scanner/requester.rs @@ -313,9 +313,12 @@ impl Requester { .set_status(ScanStatus::Cancelled) .unwrap_or_else(|e| log::warn!("Could not set scan status: {}", e)); + let scans = self.handles.ferox_scans()?; + let active_bars = scans.number_of_bars(); + // kill the scan self.ferox_scan - .abort() + .abort(active_bars) .await .unwrap_or_else(|e| log::warn!("Could not bail on scan: {}", e));