Make service scanning better

This commit is contained in:
Michael Mikovsky
2025-04-20 10:19:33 -06:00
parent 39539478b7
commit 06d3a97840
7 changed files with 113 additions and 179 deletions
+1
View File
@@ -19,3 +19,4 @@ tokio = "1.44.2"
craftping = "0.7.0"
sha256 = "1.6.0"
rayon = "1.10.0"
futures = "0.3.31"
-143
View File
@@ -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,
+1 -1
View File
@@ -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;
-2
View File
@@ -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();
+83 -23
View File
@@ -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 {
loop {
let mut hosts = thread_hosts.lock().unwrap();
// 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) {
result.open_ports.push(*port);
result.services.insert(*port, (service_name, banner));
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
}
+12 -3
View File
@@ -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()?
.read_to_string(&mut result);
.send()?;
// println!("HTTP reading");
let _ = r.read_to_string(&mut result)?;
// println!("HTTP stop");
Ok(result)
}
+12 -3
View File
@@ -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()?
.read_to_string(&mut result);
.send()?;
// println!("https read");
let _ = response.read_to_string(&mut result);
// println!("https stop");
// println!("{}", result);