From 3cb4d9c504f51fa094c52478ab53f30fff6c026b Mon Sep 17 00:00:00 2001 From: Michael Mikovsky <77305074+Astatin3@users.noreply.github.com> Date: Sat, 19 Apr 2025 12:07:22 -0600 Subject: [PATCH] Add better searching --- Cargo.toml | 5 +- src/database.rs | 621 +++++++++++++++++++++++++----- src/lib.rs | 1 + src/main.rs | 41 +- src/online_scan/online_scan.rs | 3 +- src/port_scan/port_scan.rs | 3 +- src/port_scan/tcp_scan.rs | 35 +- src/query.rs | 67 ++++ src/service_scan/mod.rs | 1 + src/service_scan/service_scan.rs | 107 +++-- src/service_scan/services.rs | 4 + src/service_scan/tcp_minecraft.rs | 42 ++ 12 files changed, 776 insertions(+), 154 deletions(-) create mode 100644 src/query.rs create mode 100644 src/service_scan/tcp_minecraft.rs diff --git a/Cargo.toml b/Cargo.toml index 9230501..91eb85a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,5 +14,8 @@ rand = "0.9.0" regex = "1.11.1" rocksdb = "0.23.0" serde = { version = "1.0.219", features = ["derive"] } -serde_json = "1.0.140" +serde_json = { version = "1.0.140", features = ["std"] } tokio = "1.44.2" +craftping = "0.7.0" +sha256 = "1.6.0" +rayon = "1.10.0" diff --git a/src/database.rs b/src/database.rs index 8d388ca..30ff284 100644 --- a/src/database.rs +++ b/src/database.rs @@ -1,8 +1,11 @@ -use std::{net::IpAddr, sync::Arc, time::Instant}; +use std::{collections::HashMap, net::IpAddr, sync::Arc, time::Instant}; +use regex::Regex; use rocksdb::{Cache, ColumnFamily, DB, IteratorMode, Options, WriteBatch}; use serde::{Deserialize, Serialize}; +use rayon::prelude::*; + use crate::{port_scan::port_scan::PortScanResult, service_scan::service_scan::ServiceScanResult}; // Global settings for optimal performance @@ -19,9 +22,10 @@ pub struct ResultDatabase { #[derive(Debug, Serialize, Deserialize, Clone)] pub struct DatabaseResult { - pub id: String, // Row identifier - pub ports: Vec, // Array of string values - pub services: String, // json services + pub id: String, + pub ports: Vec, + pub services: Vec, + pub responses: String, } impl DatabaseResult { @@ -29,82 +33,22 @@ impl DatabaseResult { let mut str = "".to_string(); str += format!( - "{} - ports: [{}] services: [{}]", + "\n{}\n- ports: [{}]\n- services: [{}]\n- responses: [{}]", self.id, join_nums(&self.ports, ","), - &self.services + self.services.join(", "), + if let Ok(data) = + serde_json::from_str::>(self.responses.as_str()) + { + format!("{:?}", data) + } else { + self.responses.clone() + } ) .as_str(); str } - pub fn encode(&self, buf: &mut Vec) { - let values = vec![self.ports_to_string(), self.services.clone()]; - - // Write number of values - buf.extend_from_slice(&(values.len() as u32).to_le_bytes()); - - // Write each value - for value in values { - let value_bytes = value.as_bytes(); - buf.extend_from_slice(&(value_bytes.len() as u32).to_le_bytes()); - buf.extend_from_slice(value_bytes); - } - } - // Binary decoding of row data - pub fn decode(key: &str, data: &[u8]) -> Option { - println!("{}", data.len()); - if data.len() < 8 { - return None; - } - - let mut pos = 0; - - if pos + 4 > data.len() { - return None; - } - let mut values_count_bytes = [0u8; 4]; - values_count_bytes.copy_from_slice(&data[pos..pos + 4]); - let values_count = u32::from_le_bytes(values_count_bytes) as usize; - pos += 4; - - // Read values - let mut values = Vec::with_capacity(values_count); - for _ in 0..values_count { - if pos + 4 > data.len() { - println!("error1!"); - return None; - } - - let mut value_len_bytes = [0u8; 4]; - value_len_bytes.copy_from_slice(&data[pos..pos + 4]); - let value_len = u32::from_le_bytes(value_len_bytes) as usize; - pos += 4; - - if pos + value_len > data.len() { - println!("error!"); - return None; - } - - let value = String::from_utf8_lossy(&data[pos..pos + value_len]).to_string(); - values.push(value); - pos += value_len; - } - - Some(DatabaseResult { - id: key.to_string(), - ports: if 1 > 0 { - split_nums(values[0].as_str(), ",") - } else { - Vec::new() - }, - services: if 1 > 1 { - values[1].clone() - } else { - String::new() - }, - }) - } pub fn ports_to_string(&self) -> String { return join_nums(&self.ports, ","); } @@ -174,6 +118,7 @@ impl ResultDatabase { "default".to_string(), "ports".to_string(), "services".to_string(), + "responses".to_string(), ]; Self { @@ -193,7 +138,8 @@ impl ResultDatabase { string_rows.push(DatabaseResult { id: result.to_string(), ports: vec![], - services: String::new(), + services: Vec::new(), + responses: String::new(), }); } @@ -220,9 +166,7 @@ impl ResultDatabase { let mut string_rows = Vec::with_capacity(results.len()); // Pre-allocate capacity for result in results { - let e = result.to_database(); - print!("{}", e.services); - string_rows.push(e); + string_rows.push(result.to_database()); } return self.save_rows(string_rows); @@ -236,6 +180,7 @@ impl ResultDatabase { let cf_default = db.cf_handle(&self.columns[0]).unwrap(); let cf_ports = db.cf_handle(&self.columns[1]).unwrap(); let cf_services = db.cf_handle(&self.columns[2]).unwrap(); + let cf_responses = db.cf_handle(&self.columns[3]).unwrap(); let start = Instant::now(); let length = string_rows.len(); @@ -268,7 +213,14 @@ impl ResultDatabase { ); // Services - batch.put_cf(cf_services, row.id.as_bytes(), row.services.as_bytes()); + batch.put_cf( + cf_services, + row.id.as_bytes(), + row.services.join(",").into_bytes(), + ); + + // Responses + batch.put_cf(cf_responses, row.id.as_bytes(), row.responses.into_bytes()); } batch @@ -302,21 +254,25 @@ impl ResultDatabase { db.cf_handle(&self.columns[0]).unwrap(), db.cf_handle(&self.columns[1]).unwrap(), db.cf_handle(&self.columns[2]).unwrap(), + db.cf_handle(&self.columns[3]).unwrap(), ]; return self.fetch_row(&db, row, &cfs); } pub fn get_rows_by_port(&self, port: &str) -> Vec { - if let Ok(result) = self.search_substring_in_column(self.columns[1].as_str(), port) { + if let Ok(result) = self.search_substring_in_column_regex( + self.columns[1].as_str(), + Regex::new(&format!(r"\b{}\b", port)).unwrap(), + ) { return result; } else { return Vec::new(); } } - pub fn get_rows_by_service(&self, port: &str) -> Vec { - if let Ok(result) = self.search_substring_in_column(self.columns[2].as_str(), port) { + pub fn get_rows_by_service(&self, service: &str) -> Vec { + if let Ok(result) = self.search_substring_in_column(self.columns[2].as_str(), service) { return result; } else { return Vec::new(); @@ -326,7 +282,7 @@ impl ResultDatabase { pub fn search_substring_in_column( &self, column: &str, - substring: &str, + string: &str, ) -> Result, rocksdb::Error> { let db = Arc::new(DB::open_cf(&self.options, &self.path, &self.columns)?); @@ -335,21 +291,17 @@ impl ResultDatabase { db.cf_handle(&self.columns[0]).unwrap(), db.cf_handle(&self.columns[1]).unwrap(), db.cf_handle(&self.columns[2]).unwrap(), + db.cf_handle(&self.columns[3]).unwrap(), ]; let mut matching_keys: Vec = Vec::new(); - // Use RocksDB's iterator for efficient scanning let iter = db.iterator_cf(cf, IteratorMode::Start); - - // Iterate through all key-value pairs in the column family for item in iter { let (key_bytes, value_bytes) = item?; - - // Convert value to string (assumes UTF-8 encoding) if let Ok(value_str) = std::str::from_utf8(&value_bytes) { // Check if the value contains the substring - if value_str.contains(substring) { + if value_str.contains(string) { // Convert key to string and add to results if let Ok(key_str) = std::str::from_utf8(&key_bytes) { if let Some(row) = self.fetch_row(&db, key_str, &cfs) { @@ -363,22 +315,509 @@ impl ResultDatabase { Ok(matching_keys) } + pub fn search_substring_in_column_regex( + &self, + column: &str, + regex: Regex, + ) -> Result, rocksdb::Error> { + let db = Arc::new(DB::open_cf(&self.options, &self.path, &self.columns)?); + + let cf = db.cf_handle(column).unwrap(); + let cfs = vec![ + db.cf_handle(&self.columns[0]).unwrap(), + db.cf_handle(&self.columns[1]).unwrap(), + db.cf_handle(&self.columns[2]).unwrap(), + db.cf_handle(&self.columns[3]).unwrap(), + ]; + + let mut matching_keys: Vec = Vec::new(); + + let iter = db.iterator_cf(cf, IteratorMode::Start); + for item in iter { + let (key_bytes, value_bytes) = item?; + if let Ok(value_str) = std::str::from_utf8(&value_bytes) { + // Check if the value contains the substring + if regex.is_match(value_str) { + // Convert key to string and add to results + if let Ok(key_str) = std::str::from_utf8(&key_bytes) { + if let Some(row) = self.fetch_row(&db, key_str, &cfs) { + matching_keys.push(row); + } + } + } + } + } + + Ok(matching_keys) + } + + pub fn search( + &self, + queries: Vec, + ) -> Result, rocksdb::Error> { + if queries.len() == 0 { + return Ok(Vec::new()); + } + if queries.len() == 1 { + // Return host if results include host + match queries[0] { + QueryDataType::Host(row) => { + return Ok(vec![ + self.get_row_by_host(row.to_string().as_str()) + .expect("Host Not Found"), + ]); + } + _ => {} + } + } + + let db = Arc::new(DB::open_cf(&self.options, &self.path, &self.columns)?); + + let cfs = vec![ + db.cf_handle(&self.columns[0]).unwrap(), + db.cf_handle(&self.columns[1]).unwrap(), + db.cf_handle(&self.columns[2]).unwrap(), + db.cf_handle(&self.columns[3]).unwrap(), + ]; + + let matching_key_bytes = search_parallel(&db, queries, &cfs); + let mut matching_rows = Vec::new(); + + for key_bytes in matching_key_bytes { + if let Ok(key_str) = std::str::from_utf8(&key_bytes) { + if let Some(row) = self.fetch_row(&db, &key_str, &cfs) { + matching_rows.push(row); + } + } + } + + Ok(matching_rows) + } + fn fetch_row(&self, db: &DB, row_id: &str, cfs: &Vec<&ColumnFamily>) -> Option { match db.get_cf(&cfs[0], row_id.as_bytes()) { Ok(Some(_)) => Some(DatabaseResult { id: row_id.to_string(), ports: split_nums(&self.row_to_string(db, row_id, &cfs[1]), ","), - services: self.row_to_string(db, row_id, &cfs[2]), + services: self + .row_to_string(db, row_id, &cfs[2]) + .split(",") + .map(|a| a.to_string()) + .collect(), + responses: self.row_to_string(db, row_id, &cfs[3]), }), _ => None, } } fn row_to_string(&self, db: &DB, row_id: &str, cf: &ColumnFamily) -> String { - if let Ok(Some(data)) = &db.get_cf(cf, row_id) { - String::from_utf8_lossy(data).to_string() + if let Ok(Some(data)) = db.get_cf(cf, row_id) { + String::from_utf8_lossy(&*data).to_string() } else { String::new() } } } + +#[derive(Debug)] +pub enum QueryDataType { + Host(IpAddr), + Port(QueryType, i32), + Service(QueryType, String, String), + FullTextIncludes(String), +} + +#[derive(Debug)] +pub enum QueryType { + Equals, + NotEquals, + Includes, + NotIncludes, +} + +// /// Search function that takes query constraints and returns matching keys +// pub fn search(db: &DB, queries: Vec) -> Vec> { +// // Get column family handles +// let cf_ports = db +// .cf_handle("ports") +// .expect("Column family 'ports' not found"); +// let cf_services = db +// .cf_handle("services") +// .expect("Column family 'services' not found"); +// let cf_responses = db +// .cf_handle("responses") +// .expect("Column family 'responses' not found"); + +// // Prepare search results +// let mut matching_keys = Vec::new(); +// let mut potential_keys = collect_all_keys(db, cf_ports); + +// // Partition queries by type to optimize processing +// let mut port_queries = Vec::new(); +// let mut service_queries = Vec::new(); +// let mut fulltext_queries = Vec::new(); + +// for query in queries { +// match query { +// QueryDataType::Port(_, _) => port_queries.push(query), +// QueryDataType::Service(_, _, _) => service_queries.push(query), +// QueryDataType::FullTextIncludes(_) => fulltext_queries.push(query), +// _ => {} +// } +// } + +// // Apply port queries first (typically most restrictive) +// if !port_queries.is_empty() { +// potential_keys = filter_by_port_queries(db, cf_ports, potential_keys, &port_queries); +// } + +// // Apply service queries +// if !service_queries.is_empty() { +// potential_keys = +// filter_by_service_queries(db, cf_services, potential_keys, &service_queries); +// } + +// // Apply fulltext queries last (typically most expensive) +// if !fulltext_queries.is_empty() { +// potential_keys = +// filter_by_fulltext_queries(db, cf_responses, potential_keys, &fulltext_queries); +// } + +// matching_keys = potential_keys; +// matching_keys +// } + +/// Collect all keys from the ports column family as potential candidates +fn collect_all_keys(db: &DB, cf: &ColumnFamily) -> Vec> { + let mut keys = Vec::new(); + let iter = db.iterator_cf(cf, rocksdb::IteratorMode::Start); + + for result in iter { + if let Ok((key, _)) = result { + keys.push(key.to_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, + queries: Vec, + cfs: &Vec<&ColumnFamily>, +) -> Vec> { + // Get column family handles + let cf_ports = cfs[1]; + let cf_services = cfs[2]; + let cf_responses = cfs[3]; + + // Collect all keys as potential candidates + let potential_keys = collect_all_keys(db, cf_ports); + + // Partition queries by type + let port_queries: Vec<_> = queries + .iter() + .filter_map(|q| { + if let QueryDataType::Port(_, _) = q { + Some(q) + } else { + None + } + }) + .collect(); + + let service_queries: Vec<_> = queries + .iter() + .filter_map(|q| { + if let QueryDataType::Service(_, _, _) = q { + Some(q) + } else { + None + } + }) + .collect(); + + let fulltext_queries: Vec<_> = queries + .iter() + .filter_map(|q| { + if let QueryDataType::FullTextIncludes(_) = q { + Some(q) + } else { + None + } + }) + .collect(); + + // Load all data for batch processing to minimize DB reads + let mut ports_data = HashMap::new(); + let mut services_data = HashMap::new(); + let mut responses_data = HashMap::new(); + + for key in &potential_keys { + if let Ok(Some(value)) = db.get_cf(cf_ports, key) { + ports_data.insert(key.clone(), value); + } + + if let Ok(Some(value)) = db.get_cf(cf_services, key) { + services_data.insert(key.clone(), value); + } + + if let Ok(Some(value)) = db.get_cf(cf_responses, key) { + responses_data.insert(key.clone(), value); + } + } + + // Process in parallel using rayon + let matching_keys: Vec> = potential_keys + .into_par_iter() + .filter(|key| { + // Check port queries + let ports_match = port_queries.is_empty() + || if let Some(ports_value) = ports_data.get(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(); + + port_queries.iter().all(|query| { + if let QueryDataType::Port(query_type, port_num) = *query { + match query_type { + QueryType::Equals => ports_str == port_num.to_string(), + QueryType::NotEquals => ports_str != port_num.to_string(), + QueryType::Includes => ports.contains(port_num), + QueryType::NotIncludes => !ports.contains(port_num), + } + } else { + false + } + }) + } else { + false + } + } else { + false + }; + + if !ports_match { + return false; + } + + // Check service queries + let services_match = service_queries.is_empty() + || if let (Some(services_value), Some(responses_value)) = + (services_data.get(key), responses_data.get(key)) + { + if let (Ok(services_str), Ok(responses_str)) = ( + std::str::from_utf8(services_value), + std::str::from_utf8(responses_value), + ) { + if let Ok(responses_map) = + serde_json::from_str::>(responses_str) + { + service_queries.iter().all(|query| { + if let QueryDataType::Service(query_type, service_name, data_str) = + *query + { + 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 == service_name && data.contains(data_str) + } + QueryType::NotIncludes => { + service != service_name || !data.contains(data_str) + } + }) + } else { + false + } + }) + } else { + false + } + } else { + false + } + } else { + false + }; + + if !services_match { + return false; + } + + // Check fulltext queries + let fulltext_match = fulltext_queries.is_empty() + || if let Some(responses_value) = responses_data.get(key) { + if let Ok(responses_str) = std::str::from_utf8(responses_value) { + fulltext_queries.iter().all(|query| { + if let QueryDataType::FullTextIncludes(search_str) = *query { + responses_str.contains(search_str) + } else { + false + } + }) + } else { + false + } + } else { + false + }; + + fulltext_match + }) + .collect(); + + matching_keys +} diff --git a/src/lib.rs b/src/lib.rs index 4ac70c9..5ef58d3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,4 +2,5 @@ pub mod database; pub mod online_scan; pub mod parse_ip_range; pub mod port_scan; +pub mod query; pub mod service_scan; diff --git a/src/main.rs b/src/main.rs index a2eb05d..5338eb2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,8 +1,13 @@ -use std::{cmp::min, env, net::IpAddr, time::Duration}; +use std::{ + cmp::min, + env, + net::IpAddr, + time::{Duration, Instant}, +}; use parse_ip_range::parse_ip_targets; use untitled::{ - database::ResultDatabase, online_scan, parse_ip_range, port_scan::tcp_scan, + database::ResultDatabase, online_scan, parse_ip_range, port_scan::tcp_scan, query, service_scan::service_scan::scan_services, }; @@ -26,14 +31,14 @@ fn main() -> Result<(), Box> { } let _ = scan(database, args[2].clone(), args[3].clone()); } - "search" => { - if args.len() != 4 { - println!("Invalid Usage!"); - print_help(Some(args[1].as_str())); - return Ok(()); - } - search(database, args[2].to_string(), args[3].to_string()); - } + // "search" => { + // if args.len() != 4 { + // println!("Invalid Usage!"); + // print_help(Some(args[1].as_str())); + // return Ok(()); + // } + // search(database, args[2].to_string(), args[3].to_string()); + // } "help" => { if args.len() != 3 { print_help(None); @@ -41,10 +46,22 @@ fn main() -> Result<(), Box> { } print_help(Some(args[2].as_str())); } + "test" => { + let start = Instant::now(); + if let Ok(query) = query::search(args[2..].join(" ")) { + let results = database.search(query); + if let Ok(results) = results { + let len = results.len(); + for result in results { + println!("{}", result.to_string()); + } + println!("{} results in {}ms", len, start.elapsed().as_millis()); + } + } + } _ => { println!("Invalid command!"); print_help(None); - return Ok(()); } } @@ -184,7 +201,7 @@ fn scan( let _ = database.add_tcp_results(&tcp_results); let service_results = - scan_services(tcp_results, min(500, up_len), Duration::from_secs(3)); + scan_services(tcp_results, min(100, up_len), Duration::from_secs(1)); println!("Finished service scan"); let _ = database.add_service_results(&service_results); } diff --git a/src/online_scan/online_scan.rs b/src/online_scan/online_scan.rs index 2f1cb3c..e679a0d 100644 --- a/src/online_scan/online_scan.rs +++ b/src/online_scan/online_scan.rs @@ -25,7 +25,8 @@ impl PingResult { DatabaseResult { id: self.host.to_string(), ports: vec![], - services: String::new(), + services: Vec::new(), + responses: String::new(), } } } diff --git a/src/port_scan/port_scan.rs b/src/port_scan/port_scan.rs index ef7dacb..daf1959 100644 --- a/src/port_scan/port_scan.rs +++ b/src/port_scan/port_scan.rs @@ -20,7 +20,8 @@ impl PortScanResult { DatabaseResult { id: self.ip.to_string(), ports: (*self.open_ports).to_vec(), - services: String::new(), + services: Vec::new(), + responses: String::new(), } } } diff --git a/src/port_scan/tcp_scan.rs b/src/port_scan/tcp_scan.rs index 2682dd8..cc985f6 100644 --- a/src/port_scan/tcp_scan.rs +++ b/src/port_scan/tcp_scan.rs @@ -2,12 +2,12 @@ use std::cmp::Ordering; use std::collections::HashMap; use std::net::{IpAddr, Ipv4Addr}; use std::str::FromStr; -use std::sync::atomic::AtomicBool; +use std::sync::atomic::{AtomicBool, AtomicU32}; use std::sync::{Arc, Mutex}; use std::thread; use std::time::{Duration, Instant}; -use indicatif::ProgressBar; +use indicatif::{ProgressBar, ProgressStyle}; use pnet::datalink::{self}; use pnet::packet::ip::IpNextHeaderProtocols; use pnet::packet::tcp::{MutableTcpPacket, TcpFlags, TcpPacket}; @@ -44,7 +44,6 @@ pub fn tcp_scan(targets: Vec, ports: Vec, timeout: Duration) -> Vec })) .expect("No valid network interface found"); - // Create transport channel for sending and receiving let (mut tx, mut rx) = transport::transport_channel( 65535, TransportChannelType::Layer4(pnet::transport::TransportProtocol::Ipv4( @@ -53,10 +52,8 @@ pub fn tcp_scan(targets: Vec, ports: Vec, timeout: Duration) -> Vec ) .expect("Failed to create transport channel"); - // Shared results let results = Arc::new(Mutex::new(HashMap::>::new())); - // Initialize results map { let mut results_map = results.lock().unwrap(); for ip in &targets { @@ -65,14 +62,16 @@ pub fn tcp_scan(targets: Vec, ports: Vec, timeout: Duration) -> Vec } 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 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(); + // let mut tmp_results: Vec<(TcpPacket<'_>, IpAddr)> = Vec::new(); let mut iter = transport::tcp_packet_iter(&mut rx); loop { @@ -102,15 +101,16 @@ pub fn tcp_scan(targets: Vec, ports: Vec, timeout: Duration) -> Vec if let Some(tcp) = TcpPacket::new(packet.packet()) { // Check for SYN+ACK flags (indicating open port) if tcp.get_flags() == TcpFlags::SYN | TcpFlags::ACK { - println!( - "Discovered open port {} on {}", - tcp.get_source(), - addr.to_string() - ); + // println!( + // "Discovered open port {} on {}", + // tcp.get_source(), + // addr.to_string() + // ); let mut results_map = receiver_results.lock().unwrap(); if let Some(open_ports) = results_map.get_mut(&addr) { open_ports.push(tcp.get_source() as i32); } + receiver_port_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed); } } } @@ -124,7 +124,10 @@ pub fn tcp_scan(targets: Vec, ports: Vec, timeout: Duration) -> Vec // for (packet, addr) in tmp_results {} }); - let pb = ProgressBar::new((targets.len() * ports.len()) as u64); + 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); @@ -143,6 +146,7 @@ pub fn tcp_scan(targets: Vec, ports: Vec, timeout: Duration) -> Vec // println!("Using IP: {}", source_ip.to_string()); let sender_finished_sending_time = Arc::clone(&finished_sending_time); + let sender_port_count = Arc::clone(&port_count); for target in &targets { for port in &ports { // let source_ip = Ipv4Addr::from_bits(random_range(0..=(0xffffffff))); @@ -177,10 +181,17 @@ pub fn tcp_scan(targets: Vec, ports: Vec, timeout: Duration) -> Vec send_tcp_packet(&mut tx, tcp_header, target); + pb.set_message(format!( + "{} ports", + sender_port_count.load(std::sync::atomic::Ordering::Relaxed), + )); pb.inc(1); + thread::sleep(Duration::from_micros(100)); } } + + pb.finish_with_message("Finished!"); sender_finished_sending_time.swap(true, std::sync::atomic::Ordering::Relaxed); // Wait for receiver to finish // thread::sleep(timeout); diff --git a/src/query.rs b/src/query.rs new file mode 100644 index 0000000..04b1099 --- /dev/null +++ b/src/query.rs @@ -0,0 +1,67 @@ +use std::{net::IpAddr, str::FromStr}; + +use regex::Regex; + +use crate::database::{QueryDataType, QueryType, split_nums}; + +pub fn search(query: String) -> Result, Box> { + if let Ok(ip) = IpAddr::from_str(&query) { + return Ok(vec![QueryDataType::Host(ip)]); + } + + let split = query.split(" "); + + let delim = Regex::new("(?:!=|[=:;])")?; + + let mut results = Vec::new(); + + for query in split { + if let Ok(ip) = IpAddr::from_str(&query) { + return Ok(vec![QueryDataType::Host(ip)]); + } + + if let Some(m) = delim.find(query) { + let tag = query[0..m.start()].to_string(); + let delim = query[m.start()..m.end()].to_string(); + let data = query[m.end()..query.len()].to_string(); + + fn get_equals_type(delim: &str) -> QueryType { + match delim { + ":" => Some(QueryType::Includes), + ";" => Some(QueryType::NotIncludes), + "=" => Some(QueryType::Equals), + "!=" => Some(QueryType::NotEquals), + _ => None, + } + .expect("Error parsing query") + } + + match tag.as_str() { + "port" => { + let mut ports = split_nums(&data, ","); + + ports.sort(); + ports.dedup(); + + for port in ports { + if port == 0 { + continue; + } + results.push(QueryDataType::Port(get_equals_type(&delim), port)); + } + } + _ => results.push(QueryDataType::Service(get_equals_type(&delim), tag, data)), + }; + } else { + results.push(QueryDataType::FullTextIncludes(query.to_string())); + } + + // (host, data) = + } + + for result in &results { + println!("{:?}", result); + } + + Ok(results) +} diff --git a/src/service_scan/mod.rs b/src/service_scan/mod.rs index c1da55d..a8e7172 100644 --- a/src/service_scan/mod.rs +++ b/src/service_scan/mod.rs @@ -2,3 +2,4 @@ pub mod service_scan; pub mod services; pub mod tcp_http; pub mod tcp_https; +pub mod tcp_minecraft; diff --git a/src/service_scan/service_scan.rs b/src/service_scan/service_scan.rs index 42ddd86..36d3e54 100644 --- a/src/service_scan/service_scan.rs +++ b/src/service_scan/service_scan.rs @@ -7,13 +7,13 @@ use std::{ time::Duration, }; -use indicatif::ProgressBar; +use indicatif::{ProgressBar, ProgressStyle}; use crate::{ database::DatabaseResult, port_scan::port_scan::PortScanResult, service_scan::tcp_http, }; -use super::{services::SERVICE_PATTERNS, tcp_https}; +use super::{services::SERVICE_PATTERNS, tcp_https, tcp_minecraft}; #[derive(Debug, Clone)] pub struct ServiceScanResult { @@ -31,24 +31,51 @@ impl ServiceScanResult { } } pub fn to_database(&self) -> DatabaseResult { + let data = serde_json::to_string(&self.services).unwrap_or(String::new()); + + let mut services = Vec::new(); + + for key in self.services.keys() { + services.push(self.services.get(key).unwrap().0.clone()); + } + + services.sort(); + services.dedup(); + + // println!("{}", data); DatabaseResult { id: self.ip.to_string(), ports: self.open_ports.clone(), - services: serde_json::to_string(&self.services).unwrap_or(String::new()), + services, + responses: data, } } } pub fn identify(ip: IpAddr, port: &i32, timeout: Duration) -> (String, String) { - let e = || basic_identify(ip, port, timeout).unwrap_or(("tcp".to_string(), "".to_string())); + let e = || { + let (service, data) = + basic_identify(ip, port, timeout).unwrap_or(("tcp".to_string(), "".to_string())); - match port { + (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)) + }; + + (match port { 80 | 8080 | 8081 | 8082 | 8083 | 8084 | 8085 | 8086 | 8087 | 8088 | 8089 => { - tuple_or_none("http", tcp_http::scan(ip, port, timeout)).unwrap_or(e()) + tuple_or_none("http", tcp_http::scan(ip, port, timeout)) } - 443 | 8443 => tuple_or_none("https", tcp_https::scan(ip, port, timeout)).unwrap_or(e()), - _ => e(), - } + 443 | 8443 => tuple_or_none("https", tcp_https::scan(ip, port, timeout)), + 25565 | 25575 => tuple_or_none("minecraft", tcp_minecraft::scan(ip, port, timeout)), + + _ => None, + }) + .unwrap_or(e()) } fn tuple_or_none( @@ -79,11 +106,18 @@ pub fn scan_services( )); let mut handles = Vec::new(); - let pb = Arc::new(ProgressBar::new(host_port_count)); + let pb = Arc::new( + ProgressBar::new(host_port_count).with_style( + ProgressStyle::with_template( + "[{msg}] {wide_bar:.magenta/red} {pos}/{len} ({eta_precise})", + ) + .unwrap(), + ), + ); // Create a thread for each chunk of IPs let chunks = split_ips_into_chunks(port_scan_results, num_threads); - for chunk in chunks { + for (i, chunk) in chunks.iter().enumerate() { let chunk_hosts = chunk.clone(); let thread_results = Arc::clone(&results); let thread_timeout = timeout; @@ -104,6 +138,7 @@ pub fn scan_services( thread_pb.inc(1); } } + // println!("Finished chunk {}", i) })); } @@ -111,18 +146,18 @@ pub fn scan_services( handle.join().unwrap(); } - pb.clone().finish_and_clear(); + pb.clone().finish_with_message("Finished!"); Arc::try_unwrap(results) .expect("Arc still has multiple owners") .into_inner() .expect("Mutex poisoned") - .into_iter() - .map(|a| { - println!("{:?}", a); - a - }) - .collect() + // .into_iter() + // .map(|a| { + // println!("{:?}", a); + // a + // }) + // .collect() } // Helper function to split the IPs into roughly equal chunks for threading @@ -216,26 +251,26 @@ fn identify_service_from_response(response: &[u8]) -> Option<&str> { } } - // For binary responses, check for pattern matches - // Check for SSL/TLS - if response.len() >= 3 && response[0] == 0x16 && (response[1] == 0x03 || response[1] == 0x02) { - return Some("ssl/tls"); - } + // // For binary responses, check for pattern matches + // // Check for SSL/TLS + // if response.len() >= 3 && response[0] == 0x16 && (response[1] == 0x03 || response[1] == 0x02) { + // return Some("ssl/tls"); + // } - // Check for MySQL protocol - if response.len() >= 5 && response[0] == 0x4a && response[1] == 0x00 { - return Some("mysql"); - } + // // Check for MySQL protocol + // if response.len() >= 5 && response[0] == 0x4a && response[1] == 0x00 { + // return Some("mysql"); + // } - // Check for MongoDB wire protocol - if response.len() >= 4 - && response[0] == 0x02 - && response[1] == 0x00 - && response[2] == 0x00 - && response[3] == 0x00 - { - return Some("mongodb"); - } + // // Check for MongoDB wire protocol + // if response.len() >= 4 + // && response[0] == 0x02 + // && response[1] == 0x00 + // && response[2] == 0x00 + // && response[3] == 0x00 + // { + // return Some("mongodb"); + // } None } diff --git a/src/service_scan/services.rs b/src/service_scan/services.rs index 43ae247..49b4217 100644 --- a/src/service_scan/services.rs +++ b/src/service_scan/services.rs @@ -3,7 +3,11 @@ use regex::Regex; lazy_static! { pub static ref SERVICE_PATTERNS: Vec<(Regex, &'static str)> = vec![ + // Unknown + (Regex::new(r"^$").unwrap(), "tcp"), + // HTTP and Web Services + (Regex::new(r"HTTP.*HTTPS").unwrap(), "https"), (Regex::new(r"^HTTP/\d").unwrap(), "http"), (Regex::new(r"Server:").unwrap(), "http"), (Regex::new(r"").unwrap(), "http"), diff --git a/src/service_scan/tcp_minecraft.rs b/src/service_scan/tcp_minecraft.rs new file mode 100644 index 0000000..74cd2ea --- /dev/null +++ b/src/service_scan/tcp_minecraft.rs @@ -0,0 +1,42 @@ +use craftping::sync::ping; +use serde_json::json; +use sha256::digest; +use std::{ + net::{IpAddr, SocketAddr, TcpStream}, + time::Duration, +}; + +pub fn scan( + ip: IpAddr, + port: &i32, + timeout: Duration, +) -> Result> { + let port = *port as u16; + let socket = SocketAddr::new(ip, port); + let ip = ip.to_string(); + + let mut stream = TcpStream::connect_timeout(&socket, timeout)?; + let pong = ping(&mut stream, &ip, port)?; + + let icon_hash = match pong.favicon { + Some(icon) => digest(icon), + None => "null".to_string(), + }; + + Ok(serde_json::to_string(&json!({ + "version": pong.version, + "protocol": pong.protocol, + "max_players": pong.max_players, + "online_players": pong.online_players, + "players_list": pong.sample, + + "description": pong.description, + "icon": icon_hash, + + "mod_info": pong.mod_info, + "forge_data": pong.forge_data, + + "enforces_secure_chat": pong.enforces_secure_chat, + "previews_chat": pong.previews_chat + }))?) +}