Add better searching

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"

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 id: String,
-    pub ports: Vec<i32>,  // Array of string values
+    pub ports: Vec<i32>,
-    pub services: String, // json services
+    pub services: Vec<String>,
+    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::<HashMap<i32, (String, String)>>(self.responses.as_str())
+            {
+                format!("{:?}", data)
+            } else {
+                self.responses.clone()
+            }
         )
         .as_str();
 
         str
     }
-    pub fn encode(&self, buf: &mut Vec<u8>) {
-        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<Self> {
-        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();
+            string_rows.push(result.to_database());
-            print!("{}", e.services);
-            string_rows.push(e);
         }
 
         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<DatabaseResult> {
-        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<DatabaseResult> {
+    pub fn get_rows_by_service(&self, service: &str) -> Vec<DatabaseResult> {
-        if let Ok(result) = self.search_substring_in_column(self.columns[2].as_str(), port) {
+        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<Vec<DatabaseResult>, 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<DatabaseResult> = 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<Vec<DatabaseResult>, 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<DatabaseResult> = 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<QueryDataType>,
+    ) -> Result<Vec<DatabaseResult>, 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<DatabaseResult> {
         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) {
+        if let Ok(Some(data)) = db.get_cf(cf, row_id) {
-            String::from_utf8_lossy(data).to_string()
+            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<QueryDataType>) -> Vec<Vec<u8>> {
+//     // 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<Vec<u8>> {
+    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<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,
+    queries: Vec<QueryDataType>,
+    cfs: &Vec<&ColumnFamily>,
+) -> Vec<Vec<u8>> {
+    // 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<Vec<u8>> = 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<i32> = ports_str
+                            .split(',')
+                            .filter_map(|p| p.trim().parse::<i32>().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::<HashMap<String, (String, String)>>(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
+}

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;

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<dyn std::error::Error>> {
             }
             let _ = scan(database, args[2].clone(), args[3].clone());
         }
-        "search" => {
+        // "search" => {
-            if args.len() != 4 {
+        //     if args.len() != 4 {
-                println!("Invalid Usage!");
+        //         println!("Invalid Usage!");
-                print_help(Some(args[1].as_str()));
+        //         print_help(Some(args[1].as_str()));
-                return Ok(());
+        //         return Ok(());
-            }
+        //     }
-            search(database, args[2].to_string(), args[3].to_string());
+        //     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<dyn std::error::Error>> {
             }
             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);
             }

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(),
         }
     }
 }

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(),
         }
     }
 }

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<IpAddr>, ports: Vec<i32>, 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<IpAddr>, ports: Vec<i32>, timeout: Duration) -> Vec
     )
     .expect("Failed to create transport channel");
 
-    // Shared results
     let results = Arc::new(Mutex::new(HashMap::<IpAddr, Vec<i32>>::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<IpAddr>, ports: Vec<i32>, 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<Instant> = 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<IpAddr>, ports: Vec<i32>, 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!(
+                            // println!(
-                                "Discovered open port {} on {}",
+                            //     "Discovered open port {} on {}",
-                                tcp.get_source(),
+                            //     tcp.get_source(),
-                                addr.to_string()
+                            //     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<IpAddr>, ports: Vec<i32>, 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<IpAddr>, ports: Vec<i32>, 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<IpAddr>, ports: Vec<i32>, 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);

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<Vec<QueryDataType>, Box<dyn std::error::Error>> {
+    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)
+}

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;

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()),
+        443 | 8443 => tuple_or_none("https", tcp_https::scan(ip, port, timeout)),
-        _ => e(),
+        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()
+    // .into_iter()
-        .map(|a| {
+    // .map(|a| {
-            println!("{:?}", a);
+    //     println!("{:?}", a);
-            a
+    //     a
-        })
+    // })
-        .collect()
+    // .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
+    // // For binary responses, check for pattern matches
-    // Check for SSL/TLS
+    // // Check for SSL/TLS
-    if response.len() >= 3 && response[0] == 0x16 && (response[1] == 0x03 || response[1] == 0x02) {
+    // if response.len() >= 3 && response[0] == 0x16 && (response[1] == 0x03 || response[1] == 0x02) {
-        return Some("ssl/tls");
+    //     return Some("ssl/tls");
-    }
+    // }
 
-    // Check for MySQL protocol
+    // // Check for MySQL protocol
-    if response.len() >= 5 && response[0] == 0x4a && response[1] == 0x00 {
+    // if response.len() >= 5 && response[0] == 0x4a && response[1] == 0x00 {
-        return Some("mysql");
+    //     return Some("mysql");
-    }
+    // }
 
-    // Check for MongoDB wire protocol
+    // // Check for MongoDB wire protocol
-    if response.len() >= 4
+    // if response.len() >= 4
-        && response[0] == 0x02
+    //     && response[0] == 0x02
-        && response[1] == 0x00
+    //     && response[1] == 0x00
-        && response[2] == 0x00
+    //     && response[2] == 0x00
-        && response[3] == 0x00
+    //     && response[3] == 0x00
-    {
+    // {
-        return Some("mongodb");
+    //     return Some("mongodb");
-    }
+    // }
 
     None
 }

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"<html.*>").unwrap(), "http"),

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<String, Box<dyn std::error::Error>> {
+    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
+    }))?)
+}