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 log::info;
use std::net::{SocketAddr, TcpStream};
use std::thread;
use std::thread::JoinHandle;
use std::{panic, thread};
fn tcp_scan(mut target: IPv4, target_port: u16) -> bool {
let _dest = target
@ -23,16 +23,16 @@ fn create_scan_thread(
thread_id: u32,
ip_range: IPv4Range,
target_port: u16,
) -> JoinHandle<Vec<bool>> {
) -> JoinHandle<Vec<(u32, bool)>> {
thread::spawn(move || {
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
ip_range.into_iter().for_each(|id| {
let target = IPv4::new(id as u64);
let result = tcp_scan(target, target_port);
results.push(result);
results.push((id, result));
});
results
@ -44,13 +44,19 @@ fn create_scan_worker(
ips_per_thread: u32,
target_port: u16,
ignorelist: Vec<u32>,
) -> JoinHandle<Vec<bool>> {
) -> thread::Result<Vec<(u32, bool)>> {
let (f, t) = (
(thread_id * ips_per_thread),
((thread_id + 1) * ips_per_thread),
);
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(
@ -59,35 +65,42 @@ fn create_scan_workers(
target_port: u16,
num_threads: u32,
ignorelist: Option<Vec<u32>>,
) -> Vec<JoinHandle<Vec<bool>>> {
) -> Vec<(u32, bool)> {
println!("Starting wwmap...");
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
// 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 {
let id_ignorelist = ignorelist.clone().unwrap_or_default();
// Create a worker
let worker = create_scan_worker(thread_id, ips_per_thread, target_port, id_ignorelist);
threads.push(worker);
let mut worker_results =
create_scan_worker(thread_id, ips_per_thread, target_port, id_ignorelist).unwrap();
results.append(&mut worker_results);
}
// Clean up the rest
if ips_left > 0 {
let id_ignorelist = ignorelist.clone().unwrap_or_default();
threads.push(create_scan_worker(
threads.len() as u32 + 1,
let mut worker_results = create_scan_worker(
results.len() as u32 + 1,
ips_per_thread,
target_port,
id_ignorelist,
));
)
.unwrap();
results.append(&mut worker_results);
}
threads
results
}
pub fn start_scan(

Loading…
Cancel
Save