diff --git a/Cargo.toml b/Cargo.toml index 91eb85a..3e82b89 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/database.rs b/src/database.rs index 78719e2..7157caf 100644 --- a/src/database.rs +++ b/src/database.rs @@ -501,149 +501,6 @@ fn collect_all_keys(db: &DB, cf: &ColumnFamily) -> Vec> { keys } -/// Filter keys by port queries -fn filter_by_port_queries( - db: &DB, - cf_ports: &ColumnFamily, - keys: Vec>, - port_queries: &[QueryDataType], -) -> Vec> { - 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 = ports_str - .split(',') - .filter_map(|p| p.trim().parse::().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>, - service_queries: &[QueryDataType], -) -> Vec> { - 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, - >(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>, - fulltext_queries: &[QueryDataType], -) -> Vec> { - 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, diff --git a/src/online_scan/ping_scanner.rs b/src/online_scan/ping_scanner.rs index c2ca879..e920410 100644 --- a/src/online_scan/ping_scanner.rs +++ b/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; diff --git a/src/port_scan/tcp_scan.rs b/src/port_scan/tcp_scan.rs index cc985f6..7598cd3 100644 --- a/src/port_scan/tcp_scan.rs +++ b/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, ports: Vec, 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 = None; // let mut tmp_results: Vec<(TcpPacket<'_>, IpAddr)> = Vec::new(); diff --git a/src/service_scan/service_scan.rs b/src/service_scan/service_scan.rs index 36d3e54..13a661e 100644 --- a/src/service_scan/service_scan.rs +++ b/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) = - basic_identify(ip, port, timeout).unwrap_or(("tcp".to_string(), "".to_string())); + // // println!("secondary1"); + // let (service, data) = + // basic_identify(ip, port, timeout).unwrap_or(("tcp".to_string(), "".to_string())); - (match service.as_str() { - "http" => tuple_or_none("http", tcp_http::scan(ip, port, timeout)), - "https" => tuple_or_none("https", tcp_https::scan(ip, port, timeout)), - "minecraft" => tuple_or_none("minecraft", tcp_minecraft::scan(ip, port, timeout)), - _ => None, - }) - .unwrap_or((service, data)) + // // println!("secondary2"); + + // (match service.as_str() { + // "http" => tuple_or_none("http", tcp_http::scan(ip, port, timeout)), + // "https" => tuple_or_none("https", tcp_https::scan(ip, port, timeout)), + // "minecraft" => tuple_or_none("minecraft", tcp_minecraft::scan(ip, port, timeout)), + // _ => 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)), - 25565 | 25575 => tuple_or_none("minecraft", tcp_minecraft::scan(ip, port, timeout)), + 443 | 8443 => { + // 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); - for (i, chunk) in chunks.iter().enumerate() { - let chunk_hosts = chunk.clone(); + // let chunks = split_ips_into_chunks(port_scan_results, num_threads); + for i in 0..=num_threads { + // 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 { - let ports = &host.open_ports; - for port in ports { - // Try to identify the service on the port - let (service_name, banner) = identify(host.ip, port, thread_timeout); + loop { + let mut hosts = thread_hosts.lock().unwrap(); + // println!("{}, {}, {}", i, hosts.len(), total_count); - let mut results_guard = thread_results.lock().unwrap(); - if let Some(result) = results_guard.iter_mut().find(|r| r.ip == host.ip) { - result.open_ports.push(*port); - result.services.insert(*port, (service_name, banner)); - } - - thread_pb.inc(1); + 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 + // 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 == ip) { + result.open_ports.push(port); + 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 } diff --git a/src/service_scan/tcp_http.rs b/src/service_scan/tcp_http.rs index d552ca5..2b61056 100644 --- a/src/service_scan/tcp_http.rs +++ b/src/service_scan/tcp_http.rs @@ -8,14 +8,23 @@ pub fn scan( timeout: Duration, ) -> Result> { 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()? - .read_to_string(&mut result); + .send()?; + + // println!("HTTP reading"); + + let _ = r.read_to_string(&mut result)?; + + // println!("HTTP stop"); Ok(result) } diff --git a/src/service_scan/tcp_https.rs b/src/service_scan/tcp_https.rs index e7efae4..237aae6 100644 --- a/src/service_scan/tcp_https.rs +++ b/src/service_scan/tcp_https.rs @@ -8,15 +8,24 @@ pub fn scan( timeout: Duration, ) -> Result> { 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()? - .read_to_string(&mut result); + .send()?; + + // println!("https read"); + + let _ = response.read_to_string(&mut result); + + // println!("https stop"); // println!("{}", result);