Add scan after search functionality

src/database.rs

@@ -109,6 +109,10 @@ impl DatabaseResult {
         }
         None
     }
+
+    pub fn sameas(&self, other: &DatabaseResult) -> bool {
+        return self.port == other.port && self.ip == other.ip;
+    }
 }
 
 pub fn join_nums(nums: &Vec<i32>, sep: &str) -> String {
@@ -238,7 +242,7 @@ impl ResultDatabase {
 
     pub fn add_data_row(
         &self,
-        results: Vec<DatabaseResult>,
+        results: &Vec<DatabaseResult>,
     ) -> Result<(), Box<dyn std::error::Error>> {
         let mut string_rows = Vec::with_capacity(results.len()); // Pre-allocate capacity
 
@@ -251,7 +255,7 @@ impl ResultDatabase {
 
     pub fn save_rows(
         &self,
-        string_rows: Vec<DatabaseResult>,
+        string_rows: Vec<&DatabaseResult>,
     ) -> Result<(), Box<dyn std::error::Error>> {
         let db = Arc::new(DB::open_cf(&self.options, &self.path, &self.columns)?);
 
@@ -273,7 +277,7 @@ impl ResultDatabase {
         let length = string_rows.len();
 
         // Split the rows into chunks for parallel processing
-        let chunks: Vec<Vec<DatabaseResult>> = string_rows
+        let chunks: Vec<Vec<&DatabaseResult>> = string_rows
             .chunks(BATCH_SIZE)
             .map(|chunk| chunk.to_vec())
             .collect();
@@ -290,7 +294,7 @@ impl ResultDatabase {
 
                     for row in chunk {
                         let key = row.get_addr();
-                        println!("{}", key);
+                        // println!("{}", key);
                         let key = key.as_bytes();
 
                         batch.put_cf(cf_addr, key, key);

src/main.rs

@@ -16,7 +16,7 @@ use rand::{
     seq::{IteratorRandom, SliceRandom},
 };
 use untitled::{
-    database::ResultDatabase,
+    database::{DatabaseResult, ResultDatabase},
     online_scan,
     parse_ip_range::{self, extract_ipv4_from_file, generate_random_ipv4_addresses},
     port_scan::tcp_scan,
@@ -76,7 +76,26 @@ enum Commands {
         query: Vec<String>,
         /// Select N random results
         #[arg(short, long, default_value_t = 0)]
-        random: usize,
+        num: usize,
+
+        /// Scan results again before printing
+        #[arg(short, long, default_value_t = false)]
+        rescan: bool,
+        /// (For Rescan only) Size of block of IPs to scan
+        #[arg(short, long, default_value_t = 4096)]
+        batch_size: usize,
+        /// (For Rescan only) The top N most common ports to scan
+        #[arg(short, long, default_value_t = 150)]
+        n_ports: usize,
+        /// (For Rescan only) Timeout for requests
+        #[arg(short, long, default_value_t = 3000)]
+        timeout_ms: u64,
+        /// (For Rescan only) Delay between icmp echo requests
+        #[arg(short, long, default_value_t = 80)]
+        ping_delay_micros: u64,
+        /// (For Rescan only) Delay between tcp syn packets
+        #[arg(short, long, default_value_t = 100)]
+        syn_tcp_delay_micros: u64,
     },
 }
 
@@ -319,30 +338,74 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                 },
             }
         }
-        Commands::Search { query, random } => {
+        Commands::Search {
+            query,
+            num,
+            rescan,
+
+            batch_size,
+            n_ports,
+            timeout_ms,
+            ping_delay_micros,
+            syn_tcp_delay_micros,
+        } => {
             let start = Instant::now();
             if let Ok(query) = query::search(query) {
                 let results = database.search(query);
                 if let Ok(results) = results {
                     let total_len = results.len();
-                    if random != 0 {
-                        let local_len = min(random, total_len);
-                        let results = results.iter().choose_multiple(&mut rng(), local_len);
-                        for result in results {
-                            println!("{}", result.to_string());
-                        }
-                        println!(
-                            "{} results in {}ms, selected {}",
-                            total_len,
-                            start.elapsed().as_millis(),
-                            local_len
-                        );
+
+                    let results = if num != 0 {
+                        results
+                            .into_iter()
+                            .choose_multiple(&mut rng(), min(num, total_len))
                     } else {
-                        for result in results {
-                            println!("{}", result.to_string());
-                        }
-                        println!("{} results in {}ms", total_len, start.elapsed().as_millis());
+                        results
+                    };
+
+                    let results = if rescan {
+                        let scan_results = scan(
+                            batch_size,
+                            &database,
+                            results
+                                .iter()
+                                .map(|r| IpAddr::from_str(&r.ip).unwrap())
+                                .collect::<Vec<IpAddr>>(),
+                            ports::PORTS[0..n_ports].to_vec(),
+                            Duration::from_millis(timeout_ms),
+                            Duration::from_micros(ping_delay_micros),
+                            Duration::from_micros(syn_tcp_delay_micros),
+                        )?;
+
+                        // Filter out only ips and ports which are the same
+                        scan_results
+                            .into_iter()
+                            .filter_map(|curresult| {
+                                for prevresult in &results {
+                                    if prevresult.sameas(&curresult) {
+                                        return Some(curresult);
+                                    }
+                                }
+
+                                None
+                            })
+                            .collect::<Vec<DatabaseResult>>()
+                    } else {
+                        results
+                    };
+
+                    let len = results.len();
+
+                    for result in results {
+                        println!("{}", result.to_string());
                     }
+
+                    println!(
+                        "{} results in {}ms, selected {}",
+                        total_len,
+                        start.elapsed().as_millis(),
+                        len
+                    );
                 }
             }
         }
@@ -359,12 +422,13 @@ fn scan(
     timeout: Duration,
     ping_delay: Duration,
     tcp_delay: Duration,
-) -> Result<(), Box<dyn std::error::Error>> {
+) -> Result<Vec<DatabaseResult>, Box<dyn std::error::Error>> {
     let start_time = Instant::now();
     let mut server_count = 0;
 
     let chunks = hosts.chunks(batch_size);
     let num_chunks = chunks.len();
+    let mut results: Vec<DatabaseResult> = Vec::new();
     for (i, hosts) in chunks.enumerate() {
         let hosts = hosts.to_vec();
         let length = hosts.len();
@@ -383,10 +447,12 @@ fn scan(
         let tcp_results = tcp_scan::tcp_scan(up_hosts, &ports, timeout, tcp_delay);
         println!("Finished port scan");
 
-        let service_results = scan_services(tcp_results, min(50, up_len), timeout);
+        let mut service_results = scan_services(tcp_results, min(50, up_len), timeout);
         println!("Finished service scan");
         server_count += service_results.len();
-        let _ = database.add_data_row(service_results);
+        let _ = database.add_data_row(&service_results)?;
+
+        results.append(&mut service_results);
     }
 
     println!("Total Servers: {}", server_count);
@@ -394,5 +460,5 @@ fn scan(
     println!("Total Elapsed: {} min", elapsed);
     println!("Rate: {} servers/min", (server_count as f32 / elapsed));
 
-    Ok(())
+    Ok(results)
 }

src/port_scan/tcp_scan.rs

@@ -65,12 +65,22 @@ pub fn tcp_scan(
         }
     }
 
+    let pb = Arc::new(
+        ProgressBar::new((targets.len() * ports.len()) as u64).with_style(
+            ProgressStyle::with_template(
+                "[{msg}] {wide_bar:.cyan/blue} {pos}/{len} ({eta_precise})",
+            )
+            .unwrap(),
+        ),
+    );
+
     let finished_sending_time = Arc::new(AtomicBool::new(false));
     let port_count = Arc::new(AtomicU32::new(0));
 
     let receiver_results = Arc::clone(&results);
     let receiver_finished_sending_time = Arc::clone(&finished_sending_time);
     let receiver_port_count = Arc::clone(&port_count);
+    let reciever_pb = Arc::clone(&pb);
     let receiver_handle = thread::spawn(move || {
         let mut finish_sending_time: Option<Instant> = None;
 
@@ -93,8 +103,11 @@ pub fn tcp_scan(
                 && receiver_finished_sending_time.load(std::sync::atomic::Ordering::Relaxed)
             {
                 finish_sending_time = Some(Instant::now());
-                // pb = Some(ProgressBar::new(TIMEOUT.as_millis() as u64));
-                println!("Waiting {} seconds for timeout...", timeout.as_secs_f32())
+
+                reciever_pb.set_message(format!(
+                    "Waiting {} seconds for timeout...",
+                    timeout.as_secs_f32()
+                ));
             }
 
             // println!("loop");
@@ -127,11 +140,6 @@ pub fn tcp_scan(
         // for (packet, addr) in tmp_results {}
     });
 
-    let pb = ProgressBar::new((targets.len() * ports.len()) as u64).with_style(
-        ProgressStyle::with_template("[{msg}] {wide_bar:.cyan/blue} {pos}/{len} ({eta_precise})")
-            .unwrap(),
-    );
-
     // println!("{:?}", interface.ips);
 
     let source_ip = std_to_pnet_ipv4(
@@ -194,7 +202,6 @@ pub fn tcp_scan(
         }
     }
 
-    pb.finish_with_message("Finished!");
     sender_finished_sending_time.swap(true, std::sync::atomic::Ordering::Relaxed);
     // Wait for receiver to finish
     // thread::sleep(timeout);
@@ -202,18 +209,24 @@ pub fn tcp_scan(
 
     // Convert results to the return format
     let results_map = results.lock().unwrap();
-    targets
+    let mut total_ips = 0;
+    let result = targets
         .iter()
         .map(|ip| {
             let mut open_ports = results_map.get(ip).cloned().unwrap_or_default();
             open_ports.sort();
             open_ports.dedup();
+            total_ips += open_ports.len();
             PortScanResult {
                 ip: *ip,
                 open_ports,
             }
         })
-        .collect()
+        .collect();
+
+    pb.finish_with_message(format!("Finished! {:?} ports", total_ips));
+
+    result
 }
 
 fn send_tcp_packet(tx: &mut TransportSender, tcp_header: MutableTcpPacket<'_>, target: &IpAddr) {

src/service_scan/service_scan.rs

@@ -95,7 +95,7 @@ pub fn scan_services(
 
     // Create a thread for each chunk of IPs
     // let chunks = split_ips_into_chunks(port_scan_results, num_threads);
-    for i in 0..=min(num_threads, host_port_count) {
+    for _ in 0..=min(num_threads, host_port_count) {
         // println!("Thread {},{}", i, chunk.len());
         // let chunk_hosts = chunk.clone();
         let thread_hosts = Arc::clone(&host_port);
@@ -134,7 +134,8 @@ pub fn scan_services(
 
                 if let Some(result) = result {
                     let mut results_guard = thread_results.lock().unwrap();
-                    println!("{}, {}", i, result.to_string());
+                    thread_pb.set_message(format!("{} Found", results_guard.len()));
+                    // println!("{}, {}", i, result.to_string());
                     results_guard.push(result);
                     std::mem::drop(results_guard);
                 }
@@ -155,12 +156,15 @@ pub fn scan_services(
         handle.join().unwrap();
     }
 
-    pb.clone().finish_with_message("Finished!");
-
-    Arc::try_unwrap(results)
+    let results = Arc::try_unwrap(results)
         .expect("Arc still has multiple owners")
         .into_inner()
-        .expect("Mutex poisoned")
+        .expect("Mutex poisoned");
+
+    pb.clone()
+        .finish_with_message(format!("Finished! {} Found", results.len()));
+
+    results
     // .into_iter()
     // .map(|a| {
     //     println!("{:?}", a);