diff --git a/src/main.rs b/src/main.rs index a34ce98..61692aa 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,5 @@ use std::{ collections::HashSet, - convert::TryInto, fs::File, io::{stderr, BufRead, BufReader}, sync::{ @@ -26,8 +25,7 @@ use feroxbuster::{ FiltersHandler, Handles, ScanHandler, StatsHandler, Tasks, TermOutHandler, }, heuristics, logger, - progress::{add_bar, BarType}, - scan_manager::{self, ScanStatus, PAUSE_SCAN}, + scan_manager::{self, PAUSE_SCAN}, scanner::{self, SCANNED_URLS}, send_command, utils::fmt_err, @@ -137,7 +135,7 @@ async fn scan(targets: Vec, handles: Arc) -> Result<()> { handles.stats.send(LoadStats(from_here))?; scanned_urls.print_known_responses(); - scanned_urls.print_completed_bars(); + scanned_urls.print_completed_bars(words.len())?; } handles.send_scan_command(ScanInitialUrls(targets))?; @@ -166,14 +164,15 @@ async fn get_targets() -> Result> { // resume-from can't be used with --url, and --stdin is marked false for every resumed // scan, making it mutually exclusive from either of the other two options if let Ok(scans) = SCANNED_URLS.scans.read() { + // todo this block shouldn't be SCANNED_URLS for scan in scans.iter() { // SCANNED_URLS gets deserialized scans added to it at program start if --resume-from // is used, so scans that aren't marked complete still need to be scanned - if matches!(*scan.status.lock().unwrap(), ScanStatus::Complete) { - // todo + if scan.is_complete() { // this one's already done, ignore it continue; } + targets.push(scan.url.to_owned()); } } diff --git a/src/scan_manager.rs b/src/scan_manager.rs index 1570235..bd2735c 100644 --- a/src/scan_manager.rs +++ b/src/scan_manager.rs @@ -37,6 +37,7 @@ use crate::{ utils::open_file, 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 @@ -120,18 +121,18 @@ impl FeroxScan { } } - /// todo + /// small wrapper to set the JoinHandle pub async fn set_task(&self, task: JoinHandle<()>) -> Result<()> { let mut guard = self.task.lock().await; let _ = std::mem::replace(&mut *guard, Some(task)); Ok(()) } - /// todo + /// small wrapper to set ScanStatus pub fn set_status(&self, status: ScanStatus) -> Result<()> { - // todo unwrap? the ? throws a cannot be sent between threads - let mut guard = self.status.lock().unwrap(); - let _ = std::mem::replace(&mut *guard, status); + if let Ok(mut guard) = self.status.lock() { + let _ = std::mem::replace(&mut *guard, status); + } Ok(()) } @@ -180,12 +181,14 @@ impl FeroxScan { } /// Mark the scan as complete and stop the scan's progress bar - pub fn finish(&self) { - self.set_status(ScanStatus::Complete).unwrap(); // todo + pub fn finish(&self) -> Result<()> { + self.set_status(ScanStatus::Complete)?; self.stop_progress_bar(); + Ok(()) } - /// todo + /// small wrapper to inspect ScanType and ScanStatus to see if a Directory scan is running or + /// in the queue to be run pub fn is_active(&self) -> bool { if let Ok(guard) = self.status.lock() { return matches!( @@ -197,6 +200,14 @@ impl FeroxScan { false } + /// small wrapper to inspect ScanStatus and see if it's Complete + pub fn is_complete(&self) -> bool { + if let Ok(guard) = self.status.lock() { + return matches!(*guard, ScanStatus::Complete); + } + false + } + /// todo doc pub async fn join(&self) { log::trace!("enter join({:?})", self); @@ -653,20 +664,21 @@ impl FeroxScans { } /// if a resumed scan is already complete, display a completed progress bar to the user - pub fn print_completed_bars(&self) { + pub fn print_completed_bars(&self, bar_length: usize) -> Result<()> { if let Ok(scans) = self.scans.read() { for scan in scans.iter() { - if matches!(*scan.status.lock()?, ScanStatus::Complete) { + if scan.is_complete() { // these scans are complete, and just need to be shown to the user let pb = add_bar( &scan.url, - words.len().try_into().unwrap_or_default(), + bar_length.try_into().unwrap_or_default(), BarType::Message, ); pb.finish(); } } } + Ok(()) } /// Forced the calling thread into a busy loop @@ -1177,7 +1189,7 @@ mod tests { let scan = FeroxScan::new(url, ScanType::Directory, pb.length(), Some(pb)); let scan_two = FeroxScan::new(url_two, ScanType::Directory, pb_two.length(), Some(pb_two)); - scan_two.finish(); // one complete, one incomplete + scan_two.finish().unwrap(); // one complete, one incomplete scan_two .set_task(tokio::spawn(async move { sleep(Duration::from_millis(SLEEP_DURATION)); diff --git a/src/scanner.rs b/src/scanner.rs index 821b771..2c52666 100644 --- a/src/scanner.rs +++ b/src/scanner.rs @@ -537,7 +537,7 @@ pub async fn scan_url( // drop the current permit so the semaphore will allow another scan to proceed drop(permit); - ferox_scan.finish(); + ferox_scan.finish()?; // todo remove // // manually drop tx in order for the rx task's while loops to eval to false