Make service scanning better

Cargo.toml

@@ -19,3 +19,4 @@ tokio = "1.44.2"
 craftping = "0.7.0"
 sha256 = "1.6.0"
 rayon = "1.10.0"
+futures = "0.3.31"

src/database.rs

@@ -501,149 +501,6 @@ fn collect_all_keys(db: &DB, cf: &ColumnFamily) -> Vec<Vec<u8>> {
     keys
 }
 
-/// Filter keys by port queries
-fn filter_by_port_queries(
-    db: &DB,
-    cf_ports: &ColumnFamily,
-    keys: Vec<Vec<u8>>,
-    port_queries: &[QueryDataType],
-) -> Vec<Vec<u8>> {
-    keys.into_iter()
-        .filter(|key| {
-            // Get the ports string for this key
-            if let Ok(Some(ports_value)) = db.get_cf(cf_ports, key) {
-                if let Ok(ports_str) = std::str::from_utf8(&ports_value) {
-                    let ports: Vec<i32> = ports_str
-                        .split(',')
-                        .filter_map(|p| p.trim().parse::<i32>().ok())
-                        .collect();
-
-                    // Check if all port queries are satisfied
-                    port_queries.iter().all(|query| {
-                        if let QueryDataType::Port(query_type, port_num) = query {
-                            match query_type {
-                                QueryType::Equals => ports.contains(port_num),
-                                QueryType::NotEquals => !ports.contains(port_num),
-                                QueryType::Includes => ports.contains(port_num),
-                                QueryType::NotIncludes => !ports.contains(port_num),
-                            }
-                        } else {
-                            false // Not a port query
-                        }
-                    })
-                } else {
-                    false
-                }
-            } else {
-                false
-            }
-        })
-        .collect()
-}
-
-/// Filter keys by service queries
-fn filter_by_service_queries(
-    db: &DB,
-    cf_services: &ColumnFamily,
-    keys: Vec<Vec<u8>>,
-    service_queries: &[QueryDataType],
-) -> Vec<Vec<u8>> {
-    keys.into_iter()
-        .filter(|key| {
-            // Get the services string for this key
-            if let Ok(Some(services_value)) = db.get_cf(cf_services, key) {
-                if let Ok(services_str) = std::str::from_utf8(&services_value) {
-                    let services: Vec<&str> = services_str.split(',').map(|s| s.trim()).collect();
-
-                    // Get the responses hashmap for this key
-                    if let Ok(Some(responses_value)) =
-                        db.get_cf(db.cf_handle("responses").unwrap(), key)
-                    {
-                        if let Ok(responses_str) = std::str::from_utf8(&responses_value) {
-                            if let Ok(responses_map) = serde_json::from_str::<
-                                HashMap<String, (String, String)>,
-                            >(responses_str)
-                            {
-                                // Check if all service queries are satisfied
-                                service_queries.iter().all(|query| {
-                                    if let QueryDataType::Service(
-                                        query_type,
-                                        service_name,
-                                        data_str,
-                                    ) = query
-                                    {
-                                        // Check across all responses in the hashmap
-                                        responses_map.values().any(|(service, data)| {
-                                            match query_type {
-                                                QueryType::Equals => {
-                                                    service == service_name && data == data_str
-                                                }
-                                                QueryType::NotEquals => {
-                                                    service != service_name || data != data_str
-                                                }
-                                                QueryType::Includes => {
-                                                    service.contains(service_name)
-                                                        && data.contains(data_str)
-                                                }
-                                                QueryType::NotIncludes => {
-                                                    !service.contains(service_name)
-                                                        || !data.contains(data_str)
-                                                }
-                                            }
-                                        })
-                                    } else {
-                                        false // Not a service query
-                                    }
-                                })
-                            } else {
-                                false
-                            }
-                        } else {
-                            false
-                        }
-                    } else {
-                        false
-                    }
-                } else {
-                    false
-                }
-            } else {
-                false
-            }
-        })
-        .collect()
-}
-
-/// Filter keys by fulltext queries (most expensive operation)
-fn filter_by_fulltext_queries(
-    db: &DB,
-    cf_responses: &ColumnFamily,
-    keys: Vec<Vec<u8>>,
-    fulltext_queries: &[QueryDataType],
-) -> Vec<Vec<u8>> {
-    keys.into_iter()
-        .filter(|key| {
-            // Get the raw responses string for this key
-            if let Ok(Some(responses_value)) = db.get_cf(cf_responses, key) {
-                if let Ok(responses_str) = std::str::from_utf8(&responses_value) {
-                    // Check if all fulltext queries are satisfied
-                    fulltext_queries.iter().all(|query| {
-                        if let QueryDataType::FullTextIncludes(search_str) = query {
-                            responses_str.contains(search_str)
-                        } else {
-                            false // Not a fulltext query
-                        }
-                    })
-                } else {
-                    false
-                }
-            } else {
-                false
-            }
-        })
-        .collect()
-}
-
 /// Optimized search implementation with parallelism for large datasets
 pub fn search_parallel(
     db: &DB,

src/online_scan/ping_scanner.rs

@@ -5,7 +5,7 @@ use pnet::packet::{
     icmp::{IcmpTypes, echo_request::MutableEchoRequestPacket},
 };
 use pnet::transport::{
-    TransportChannelType, TransportProtocol, icmp_packet_iter, tcp_packet_iter, transport_channel,
+    TransportChannelType, TransportProtocol, icmp_packet_iter, transport_channel,
 };
 use pnet::util::checksum;
 use std::collections::HashMap;

src/port_scan/tcp_scan.rs

@@ -1,4 +1,3 @@
-use std::cmp::Ordering;
 use std::collections::HashMap;
 use std::net::{IpAddr, Ipv4Addr};
 use std::str::FromStr;
@@ -68,7 +67,6 @@ pub fn tcp_scan(targets: Vec<IpAddr>, ports: Vec<i32>, timeout: Duration) -> Vec
     let receiver_finished_sending_time = Arc::clone(&finished_sending_time);
     let receiver_port_count = Arc::clone(&port_count);
     let receiver_handle = thread::spawn(move || {
-        let start_time = std::time::Instant::now();
         let mut finish_sending_time: Option<Instant> = None;
 
         // let mut tmp_results: Vec<(TcpPacket<'_>, IpAddr)> = Vec::new();

src/service_scan/service_scan.rs

@@ -2,12 +2,13 @@ use std::{
     collections::HashMap,
     io::{Read, Write},
     net::{IpAddr, SocketAddr, TcpStream},
-    sync::{Arc, Mutex},
+    sync::{Arc, Mutex, MutexGuard},
     thread,
     time::Duration,
 };
 
 use indicatif::{ProgressBar, ProgressStyle};
+use rand::seq::SliceRandom;
 
 use crate::{
     database::DatabaseResult, port_scan::port_scan::PortScanResult, service_scan::tcp_http,
@@ -54,28 +55,42 @@ impl ServiceScanResult {
 
 pub fn identify(ip: IpAddr, port: &i32, timeout: Duration) -> (String, String) {
     let e = || {
-        let (service, data) =
+        // // println!("secondary1");
-            basic_identify(ip, port, timeout).unwrap_or(("tcp".to_string(), "".to_string()));
+        // let (service, data) =
+        //     basic_identify(ip, port, timeout).unwrap_or(("tcp".to_string(), "".to_string()));
 
-        (match service.as_str() {
+        // // println!("secondary2");
-            "http" => tuple_or_none("http", tcp_http::scan(ip, port, timeout)),
+
-            "https" => tuple_or_none("https", tcp_https::scan(ip, port, timeout)),
+        // (match service.as_str() {
-            "minecraft" => tuple_or_none("minecraft", tcp_minecraft::scan(ip, port, timeout)),
+        //     "http" => tuple_or_none("http", tcp_http::scan(ip, port, timeout)),
-            _ => None,
+        //     "https" => tuple_or_none("https", tcp_https::scan(ip, port, timeout)),
-        })
+        //     "minecraft" => tuple_or_none("minecraft", tcp_minecraft::scan(ip, port, timeout)),
-        .unwrap_or((service, data))
+        //     _ => None,
+        // })
+        // .unwrap_or((service, data))
+        basic_identify(ip, port, timeout).unwrap_or(("tcp".to_string(), "".to_string()))
     };
 
+    // println!("primary");
+
     (match port {
         80 | 8080 | 8081 | 8082 | 8083 | 8084 | 8085 | 8086 | 8087 | 8088 | 8089 => {
+            // println!("http");
             tuple_or_none("http", tcp_http::scan(ip, port, timeout))
         }
-        443 | 8443 => tuple_or_none("https", tcp_https::scan(ip, port, timeout)),
+        443 | 8443 => {
-        25565 | 25575 => tuple_or_none("minecraft", tcp_minecraft::scan(ip, port, timeout)),
+            // println!("https");
+            tuple_or_none("https", tcp_https::scan(ip, port, timeout))
+        }
+        25565 | 25575 => {
+            // println!("minecraft");
+            tuple_or_none("minecraft", tcp_minecraft::scan(ip, port, timeout))
+        }
 
         _ => None,
     })
     .unwrap_or(e())
+    // basic_identify(ip, port, timeout).unwrap_or(("tcp".to_string(), "".to_string()))
 }
 
 fn tuple_or_none(
@@ -105,6 +120,17 @@ pub fn scan_services(
             .collect(),
     ));
 
+    let mut host_port: Vec<(IpAddr, i32)> = Vec::with_capacity(host_port_count as usize);
+    for host in &port_scan_results {
+        for port in &host.open_ports {
+            host_port.push((host.ip, port.clone()));
+        }
+    }
+
+    host_port.shuffle(&mut rand::rng());
+
+    let host_port = Arc::new(Mutex::new(host_port));
+
     let mut handles = Vec::new();
     let pb = Arc::new(
         ProgressBar::new(host_port_count).with_style(
@@ -116,28 +142,57 @@ pub fn scan_services(
     );
 
     // Create a thread for each chunk of IPs
-    let chunks = split_ips_into_chunks(port_scan_results, num_threads);
+    // let chunks = split_ips_into_chunks(port_scan_results, num_threads);
-    for (i, chunk) in chunks.iter().enumerate() {
+    for i in 0..=num_threads {
-        let chunk_hosts = chunk.clone();
+        // println!("Thread {},{}", i, chunk.len());
+        // let chunk_hosts = chunk.clone();
+        let thread_hosts = Arc::clone(&host_port);
         let thread_results = Arc::clone(&results);
         let thread_timeout = timeout;
         let thread_pb = Arc::clone(&pb);
         handles.push(thread::spawn(move || {
-            for host in chunk_hosts {
+            loop {
-                let ports = &host.open_ports;
+                let mut hosts = thread_hosts.lock().unwrap();
-                for port in ports {
+                // println!("{}, {}, {}", i, hosts.len(), total_count);
+
+                if hosts.len() == 0 {
+                    // println!("Break thread {} A", i);
+                    break;
+                }
+
+                let host = hosts.pop();
+
+                std::mem::drop(hosts);
+
+                if host.is_none() {
+                    // println!("Break thread {} B", i);
+                    break;
+                }
+                let host = host.unwrap();
+
+                let ip = host.0;
+                let port = host.1;
+
+                println!("{}, {}, {}", i, ip, port);
+
                 // Try to identify the service on the port
-                    let (service_name, banner) = identify(host.ip, port, thread_timeout);
+                // println!("Thread {} stall 2", i);
+                let (service_name, banner) = identify(ip, &port, thread_timeout);
+                // println!("Thread {} stall 3", i);
 
                 let mut results_guard = thread_results.lock().unwrap();
-                    if let Some(result) = results_guard.iter_mut().find(|r| r.ip == host.ip) {
+                if let Some(result) = results_guard.iter_mut().find(|r| r.ip == ip) {
-                        result.open_ports.push(*port);
+                    result.open_ports.push(port);
-                        result.services.insert(*port, (service_name, banner));
+                    result.services.insert(port, (service_name, banner));
                 }
+                // println!("Thread {} stall 4", i);
 
                 thread_pb.inc(1);
+                // println!("Thread {}", i);
+
+                // total_count += 1;
             }
-            }
+            // println!("Thread {}", i);
             // println!("Finished chunk {}", i)
         }));
     }
@@ -222,6 +277,7 @@ fn try_connect(ip: IpAddr, port: &i32, timeout: Duration, probe: &[u8]) -> Optio
 }
 
 fn basic_identify(ip: IpAddr, port: &i32, timeout: Duration) -> Option<(String, String)> {
+    // println!("Start try_connect");
     // Try a simple connection with no probe as last resort
     if let Some(response) = try_connect(ip, port, timeout, b"\x00\n") {
         if !response.is_empty() {
@@ -233,10 +289,14 @@ fn basic_identify(ip: IpAddr, port: &i32, timeout: Duration) -> Option<(String,
             }
         }
 
+        // println!("End try_connect1");
+
         // Port is open but service couldn't be identified
         return Some(("tcp".to_string(), "".to_string()));
     }
 
+    // println!("Start try_connect2");
+
     None
 }
 

src/service_scan/tcp_http.rs

@@ -8,14 +8,23 @@ pub fn scan(
     timeout: Duration,
 ) -> Result<String, Box<dyn std::error::Error>> {
     let mut result = String::new();
-    let _ = reqwest::blocking::Client::builder()
+
+    // println!("HTTP start");
+
+    let mut r = reqwest::blocking::Client::builder()
         .redirect(Policy::none())
         .timeout(timeout)
+        .connect_timeout(timeout)
         .build()
         .unwrap()
         .get(format!("http://{}:{}", ip.to_string(), port))
-        .send()?
+        .send()?;
-        .read_to_string(&mut result);
+
+    // println!("HTTP reading");
+
+    let _ = r.read_to_string(&mut result)?;
+
+    // println!("HTTP stop");
 
     Ok(result)
 }

src/service_scan/tcp_https.rs

@@ -8,15 +8,24 @@ pub fn scan(
     timeout: Duration,
 ) -> Result<String, Box<dyn std::error::Error>> {
     let mut result = String::new();
-    let _response = reqwest::blocking::Client::builder()
+
+    // println!("https start");
+
+    let mut response = reqwest::blocking::Client::builder()
         .danger_accept_invalid_certs(true)
         .redirect(Policy::none())
         .timeout(timeout)
+        .connect_timeout(timeout)
         .build()
         .unwrap()
         .get(format!("https://{}:{}", ip.to_string(), port))
-        .send()?
+        .send()?;
-        .read_to_string(&mut result);
+
+    // println!("https read");
+
+    let _ = response.read_to_string(&mut result);
+
+    // println!("https stop");
 
     // println!("{}", result);