From 031e9ed49154c59aa0f51a2025443a7758ad6fe1 Mon Sep 17 00:00:00 2001
From: "E. Almqvist" <elalmqvist@gmail.com>
Date: Thu, 11 Aug 2022 20:29:47 +0200
Subject: [PATCH] Refactor & stuff

---
 src/ipv4.rs    | 46 +++++++++++++++++-------
 src/scanner.rs | 98 ++++++++++++++++++++++++--------------------------
 2 files changed, 80 insertions(+), 64 deletions(-)

diff --git a/src/ipv4.rs b/src/ipv4.rs
index ea60063..be011aa 100644
--- a/src/ipv4.rs
+++ b/src/ipv4.rs
@@ -4,23 +4,11 @@ use convert_base::Convert;
 use std::convert::TryInto;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 
-/*
-Algorithm: O(n)
-
-let i = 0 .. u32:max_value()
-
-# Convert each i to base 256 and we get all the ipv4 addresses
-# This is waaaay better than a stupid loop
-*/
-
 #[derive(Debug, Copy, Clone)]
 pub struct IPv4 {
     pub id: u64,
     pub ip: [u8; 4],
     pub ignore: bool,
-//     pub enum __Generator {
-//         Start(&'static)
-//     }
 }
 
 impl IPv4 {
@@ -66,6 +54,40 @@ impl IPv4 {
     }
 }
 
+pub struct IPv4_Range {
+    pub id_start: u32,
+    pub id_end: u32,
+    pub id_ignore: Vec<u32>,
+}
+
+impl IPv4_Range {
+    pub fn new(from: u32, to: u32, id_ignore: Option<Vec<u32>>) -> Self {
+        Self {
+            id_start: from,
+            id_end: to,
+            id_ignore: id_ignore.unwrap_or(Vec::new()),
+        }
+    }
+}
+
+impl Iterator for IPv4_Range {
+    type Item = u32;
+
+    fn next(&mut self) -> Option<u32> {
+        if self.id_start == self.id_end {
+            None
+        } else {
+            let res = Some(self.id_start);
+            self.id_start += 1;
+
+            if self.id_ignore.contains(&res.unwrap()) {
+                return self.next();
+            }
+            res
+        }
+    }
+}
+
 pub fn get_all(ignorelist: Option<Vec<u64>>) -> Result<Vec<IPv4>> {
     // Ignore those that we know
     let ignorelist = ignorelist.unwrap_or(Vec::new());
diff --git a/src/scanner.rs b/src/scanner.rs
index ebf9c36..c499030 100644
--- a/src/scanner.rs
+++ b/src/scanner.rs
@@ -9,70 +9,64 @@ fn tcp_scan(mut target: ipv4::IPv4, target_port: u16) -> bool {
     let dest = target.to_socketaddr(target_port).unwrap();
 
     false
-// TODO: FIX
-//     if let Ok(res) = TcpStream::connect(dest) {
-//         info!("* Got TCP ack from: {:?} | {:?}", dest, res);
-//         return true;
-//     }
-//     false
+    // TODO: FIX
+    //     if let Ok(res) = TcpStream::connect(dest) {
+    //         info!("* Got TCP ack from: {:?} | {:?}", dest, res);
+    //         return true;
+    //     }
+    //     false
 }
 
-fn create_scan_thread(
-    ip_list: Vec<ipv4::IPv4>,
-    thread_id: u32,
-    ips_per_thread: u32,
-    target_port: u16,
-) -> JoinHandle<Vec<bool>> {
-    thread::spawn(move || {
-        info!("Starting thread worker #{}", thread_id);
-        let mut results: Vec<bool> = Vec::new();
-
-        // do the scan thing
-        for i in 0..ips_per_thread {
-            let id = (thread_id * ips_per_thread) + i;
-            let ref target = ip_list[id as usize];
-            let result = tcp_scan(*target, target_port);
-            results.push(result);
-        }
-
-        results
-    })
-}
+// fn create_scan_thread(
+//     ip_list: Vec<ipv4::IPv4>,
+//     thread_id: u32,
+//     ips_per_thread: u32,
+//     target_port: u16,
+// ) -> JoinHandle<Vec<bool>> {
+//     thread::spawn(move || {
+//         info!("Starting thread worker #{}", thread_id);
+//         let mut results: Vec<bool> = Vec::new();
+// 
+//         // do the scan thing
+//         for i in 0..ips_per_thread {
+//             let id = (thread_id * ips_per_thread) + i;
+//             let ref target = ip_list[id as usize];
+//             let result = tcp_scan(*target, target_port);
+//             results.push(result);
+//         }
+// 
+//         results
+//     })
+// }
 
-pub fn start_scan(target_port: u16, num_threads: u32, ignorelist: Option<Vec<u64>>) -> Result<()> {
+pub fn start_scan(
+    from: u32,
+    to: u32,
+    target_port: u16,
+    num_threads: u32,
+    ignorelist: Option<Vec<u32>>,
+) -> Result<()> {
     println!("Starting wwmap...");
-    let ip_list = ipv4::get_all(ignorelist)?;
-
-    println!("Post list gen");
+    //let ip_list = ipv4::get_all(ignorelist)?;
 
-    let ips_per_thread = ((ip_list.len() 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 = 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 result_list: Vec<bool> = Vec::new();
+    let ip_ranges: Vec<ipv4::IPv4_Range> = Vec::new();
 
-    // create all of our threads
     for thread_id in 0..num_threads {
-        threads.push(create_scan_thread(
-            ip_list.clone(),
-            thread_id,
-            ips_per_thread,
-            target_port,
-        ));
-    }
+        let (f, t) = ((thread_id * ips_per_thread), ((thread_id + 1) * ips_per_thread));
+        let range = ipv4::IPv4_Range::new(f, t, ignorelist); 
 
-    // create the last thread to do the job
-    if ips_left > 0 {
-        threads.push(create_scan_thread(
-            ip_list,
-            num_threads,
-            ips_left,
-            target_port,
-        ));
+        ip_ranges.push(range);
     }
 
-    // Append all of the results into the list
+    // container for all of our threads
+    let mut threads: Vec<JoinHandle<Vec<bool>>> = Vec::new();
+
+
+    // container for all the results
+    let mut result_list: Vec<bool> = Vec::new();
 
     Ok(())
 }