Worker refactor

main
E. Almqvist 2 years ago
parent d13364fc45
commit 0411248a93
  1. 41
      src/scanner.rs

@ -2,8 +2,8 @@ use crate::ipv4::{IPv4, IPv4Range};
use anyhow::Result; use anyhow::Result;
use log::info; use log::info;
use std::net::{SocketAddr, TcpStream}; use std::net::{SocketAddr, TcpStream};
use std::thread;
use std::thread::JoinHandle; use std::thread::JoinHandle;
use std::{panic, thread};
fn tcp_scan(mut target: IPv4, target_port: u16) -> bool { fn tcp_scan(mut target: IPv4, target_port: u16) -> bool {
let _dest = target let _dest = target
@ -23,16 +23,16 @@ fn create_scan_thread(
thread_id: u32, thread_id: u32,
ip_range: IPv4Range, ip_range: IPv4Range,
target_port: u16, target_port: u16,
) -> JoinHandle<Vec<bool>> { ) -> JoinHandle<Vec<(u32, bool)>> {
thread::spawn(move || { thread::spawn(move || {
info!("Starting thread worker #{}", thread_id); info!("Starting thread worker #{}", thread_id);
let mut results: Vec<bool> = Vec::new(); let mut results: Vec<(u32, bool)> = Vec::new();
// do the scan thing // do the scan thing
ip_range.into_iter().for_each(|id| { ip_range.into_iter().for_each(|id| {
let target = IPv4::new(id as u64); let target = IPv4::new(id as u64);
let result = tcp_scan(target, target_port); let result = tcp_scan(target, target_port);
results.push(result); results.push((id, result));
}); });
results results
@ -44,13 +44,19 @@ fn create_scan_worker(
ips_per_thread: u32, ips_per_thread: u32,
target_port: u16, target_port: u16,
ignorelist: Vec<u32>, ignorelist: Vec<u32>,
) -> JoinHandle<Vec<bool>> { ) -> thread::Result<Vec<(u32, bool)>> {
let (f, t) = ( let (f, t) = (
(thread_id * ips_per_thread), (thread_id * ips_per_thread),
((thread_id + 1) * ips_per_thread), ((thread_id + 1) * ips_per_thread),
); );
let range = IPv4Range::new(f, t, ignorelist); let range = IPv4Range::new(f, t, ignorelist);
create_scan_thread(thread_id, range, target_port) let thread = create_scan_thread(thread_id, range, target_port);
let result = thread.join();
match result {
Ok(_) => return result,
Err(e) => panic!("** Worker panic! {:?}", e),
}
} }
fn create_scan_workers( fn create_scan_workers(
@ -59,35 +65,42 @@ fn create_scan_workers(
target_port: u16, target_port: u16,
num_threads: u32, num_threads: u32,
ignorelist: Option<Vec<u32>>, ignorelist: Option<Vec<u32>>,
) -> Vec<JoinHandle<Vec<bool>>> { ) -> Vec<(u32, bool)> {
println!("Starting wwmap..."); println!("Starting wwmap...");
let ips_per_thread = (((to - from) as f32) / num_threads as f32) as u32; let ips_per_thread = (((to - from) as f32) / num_threads as f32) as u32;
let ips_left = (to - from) - (num_threads * ips_per_thread); // how many ips we have left after the first threads let ips_left = (to - from) - (num_threads * ips_per_thread); // how many ips we have left after the first threads
// container for all of our threads // container for all of our threads
let mut threads: Vec<JoinHandle<Vec<bool>>> = Vec::new(); // let mut threads: Vec<JoinHandle<Vec<bool>>> = Vec::new();
// container for all of our results
let mut results: Vec<(u32, bool)> = Vec::new();
for thread_id in 0..num_threads { for thread_id in 0..num_threads {
let id_ignorelist = ignorelist.clone().unwrap_or_default(); let id_ignorelist = ignorelist.clone().unwrap_or_default();
// Create a worker // Create a worker
let worker = create_scan_worker(thread_id, ips_per_thread, target_port, id_ignorelist); let mut worker_results =
threads.push(worker); create_scan_worker(thread_id, ips_per_thread, target_port, id_ignorelist).unwrap();
results.append(&mut worker_results);
} }
// Clean up the rest // Clean up the rest
if ips_left > 0 { if ips_left > 0 {
let id_ignorelist = ignorelist.clone().unwrap_or_default(); let id_ignorelist = ignorelist.clone().unwrap_or_default();
threads.push(create_scan_worker( let mut worker_results = create_scan_worker(
threads.len() as u32 + 1, results.len() as u32 + 1,
ips_per_thread, ips_per_thread,
target_port, target_port,
id_ignorelist, id_ignorelist,
)); )
.unwrap();
results.append(&mut worker_results);
} }
threads results
} }
pub fn start_scan( pub fn start_scan(

Loading…
Cancel
Save