Improve saving format

Cargo.toml

@@ -4,7 +4,6 @@ version = "0.1.0"
 edition = "2024"
 
 [dependencies]
-bincode = { version = "2.0.1", features = ["serde"] }
 byteorder = "1.5.0"
 indicatif = "0.17.11"
 memchr = "2.7.4"

src/database.rs

@@ -1,12 +1,10 @@
 use std::{net::IpAddr, sync::Arc, time::Instant};
 
-use rocksdb::{Cache, ColumnFamily, DB, IteratorMode, Options, WriteBatch};
+use rocksdb::{Cache, ColumnFamily, IteratorMode, Options, WriteBatch, DB};
 use serde::{Deserialize, Serialize};
 
 use crate::port_scan::port_scan::ScanResult;
 
-static COLUMN_COUNT: usize = 5;
-
 // Global settings for optimal performance
 const BLOCK_CACHE_SIZE_MB: usize = 512; // 512MB block cache
 const WRITE_BUFFER_SIZE_MB: usize = 64; // 64MB write buffer
@@ -21,25 +19,49 @@ pub struct ResultDatabase {
 
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct StringRow {
-    pub id: String,          // Row identifier
+    pub id: String,      // Row identifier
-    pub values: Vec<String>, // Array of string values
+    pub ports: Vec<i32>, // Array of string values
 }
 
 impl StringRow {
     pub fn to_string(&self) -> String {
         let mut str = "".to_string();
 
-        str += format!("Row ID: {}, Values: [", self.id).as_str();
+        str += format!("{} - ports: [{}]", self.id, join_nums(&self.ports, ",")).as_str();
-        for (i, value) in self.values.iter().enumerate() {
-            if i > 0 {
-                str += ", ";
-            }
-            str += format!("{}: \"{}\"", i, value).as_str();
-        }
-        str += "]";
 
         str
     }
+    pub fn ports_to_string(&self) -> String {
+        return join_nums(&self.ports, ",");
+    }
+}
+
+pub fn join_nums(nums: &Vec<i32>, sep: &str) -> String {
+    // 1. Convert numbers to strings
+    let str_nums: Vec<String> = nums
+        .iter()
+        .map(|n| n.to_string()) // map every integer to a string
+        .collect(); // collect the strings into the vector
+
+    // 2. Join the strings. There's already a function for this.
+    str_nums.join(sep)
+}
+
+pub fn split_nums(str: &str, sep: &str) -> Vec<i32> {
+    if str.is_empty() {
+        return vec![];
+    };
+
+    return str
+        .split(sep)
+        .map(|n| {
+            if let Ok(num) = n.parse::<i32>() {
+                return num;
+            } else {
+                return 0;
+            }
+        })
+        .collect();
 }
 
 impl ResultDatabase {
@@ -74,13 +96,7 @@ impl ResultDatabase {
 
         // Define column families for different indexes
 
-        let mut column_families = vec!["default".to_string()]; // Main data store
+        let column_families = vec!["default".to_string(), "ports".to_string()];
-
-        // Add column families for each column index we might want to search by
-        // (for demo, we'll create indexes for 5 potential columns)
-        for i in 0..COLUMN_COUNT {
-            column_families.push(format!("col{}_idx", i).to_string());
-        }
 
         Self {
             path: path.to_string(),
@@ -98,7 +114,7 @@ impl ResultDatabase {
         for result in results {
             string_rows.push(StringRow {
                 id: result.to_string(),
-                values: vec![],
+                ports: vec![],
             });
         }
 
@@ -120,14 +136,8 @@ impl ResultDatabase {
 
     pub fn save_rows(&self, string_rows: Vec<StringRow>) -> Result<(), Box<dyn std::error::Error>> {
         let db = Arc::new(DB::open_cf(&self.options, &self.path, &self.columns)?);
-        let cf_default = db.cf_handle("default").unwrap();
+        let cf_default = db.cf_handle(&self.columns[0]).unwrap();
-
+        let cf_ports = db.cf_handle(&self.columns[1]).unwrap();
-        // Get handles to column index families
-        let mut cf_columns = Vec::new();
-        for i in 0..3 {
-            let cf = db.cf_handle(&format!("col{}_idx", i)).unwrap();
-            cf_columns.push(cf);
-        }
 
         let start = Instant::now();
         let length = string_rows.len();
@@ -142,7 +152,6 @@ impl ResultDatabase {
         let elapsed = {
             let db_ref = Arc::clone(&db);
             let cf_default_ref = cf_default;
-            let cf_columns_ref = &cf_columns;
 
             // Create batches in parallel but write them sequentially
             let batches: Vec<WriteBatch> = chunks
@@ -161,19 +170,9 @@ impl ResultDatabase {
                         // Store in main column family
                         batch.put_cf(cf_default_ref, row.id.as_bytes(), &data);
 
-                        // Create indexes only for searchable columns (0-2)
+                        let idx_key =
-                        for (col_idx, value) in row.values.iter().enumerate() {
+                            format!("{}:{}", fast_escape(row.ports_to_string().as_str()), row.id);
-                            if col_idx < cf_columns_ref.len() {
+                        batch.put_cf(cf_ports, idx_key.as_bytes(), row.id.as_bytes());
-                                // Create search-friendly keys: value:rowid
-                                // Use minimal escaping for better performance
-                                let idx_key = format!("{}:{}", fast_escape(value), row.id);
-                                batch.put_cf(
-                                    cf_columns_ref[col_idx],
-                                    idx_key.as_bytes(),
-                                    row.id.as_bytes(),
-                                );
-                            }
-                        }
                     }
 
                     batch
@@ -304,17 +303,19 @@ fn decode_row_binary(key: &str, data: &[u8]) -> Option<StringRow> {
 
     Some(StringRow {
         id: key.to_string(),
-        values,
+        ports: split_nums(values[0].as_str(), ","),
     })
 }
 
 // Binary encoding of row data for maximum performance
 fn encode_row_binary(buf: &mut Vec<u8>, row: &StringRow) {
+    let values = vec![row.ports_to_string()];
+
     // Write number of values
-    buf.extend_from_slice(&(row.values.len() as u32).to_le_bytes());
+    buf.extend_from_slice(&(values.len() as u32).to_le_bytes());
 
     // Write each value
-    for value in &row.values {
+    for value in vec![row.ports_to_string()] {
         let value_bytes = value.as_bytes();
         buf.extend_from_slice(&(value_bytes.len() as u32).to_le_bytes());
         buf.extend_from_slice(value_bytes);

src/main.rs

@@ -164,6 +164,7 @@ fn scan(
                 ];
 
                 let tcp_results = tcp_scan::tcp_scan(up_hosts, tcp_ports, Duration::from_secs(3));
+                println!("Saving Data...");
                 let _ = database.add_tcp_results(&tcp_results);
             }
         }

src/online_scan/online_scan.rs

@@ -24,18 +24,7 @@ impl PingResult {
     pub fn to_string_row(&self) -> StringRow {
         StringRow {
             id: self.host.to_string(),
-            values: vec![
+            ports: vec![],
-                if self.is_up {
-                    "up".to_string()
-                } else {
-                    "down".to_string()
-                },
-                if self.response_time.is_some() {
-                    self.response_time.unwrap().as_millis().to_string()
-                } else {
-                    "None".to_string()
-                },
-            ],
         }
     }
 }

src/port_scan/port_scan.rs

@@ -20,18 +20,7 @@ impl ScanResult {
     pub fn to_string_row(&self) -> StringRow {
         StringRow {
             id: self.ip.to_string(),
-            values: vec![join_nums(&self.open_ports, ",")],
+            ports: (*self.open_ports).to_vec(),
         }
     }
 }
-
-fn join_nums(nums: &Vec<i32>, sep: &str) -> String {
-    // 1. Convert numbers to strings
-    let str_nums: Vec<String> = nums
-        .iter()
-        .map(|n| n.to_string()) // map every integer to a string
-        .collect(); // collect the strings into the vector
-
-    // 2. Join the strings. There's already a function for this.
-    str_nums.join(sep)
-}

src/port_scan/tcp_scan.rs

@@ -6,11 +6,12 @@ use std::thread;
 use std::time::Duration;
 
 use indicatif::ProgressBar;
+use pnet::datalink::linux::interfaces;
 use pnet::datalink::{self};
 use pnet::packet::ip::IpNextHeaderProtocols;
 use pnet::packet::tcp::{MutableTcpPacket, TcpFlags, TcpPacket};
-use pnet::packet::{Packet, tcp};
+use pnet::packet::{tcp, Packet};
-use pnet::transport::{self, TransportChannelType};
+use pnet::transport::{self, TransportChannelType, TransportSender};
 use rand::random_range;
 
 use super::port_scan::ScanResult;
@@ -22,9 +23,20 @@ fn std_to_pnet_ipv4(previous: &IpAddr) -> Ipv4Addr {
 // Main scanning function
 pub fn tcp_scan(targets: Vec<IpAddr>, ports: Vec<i32>, timeout: Duration) -> Vec<ScanResult> {
     // Find network interface
-    let interface = datalink::interfaces()
+    let interfaces = datalink::interfaces();
+
+    println!("{}", interfaces.len());
+
+    let interface = interfaces
         .into_iter()
-        .find(|iface| iface.is_up() && !iface.is_loopback() && !iface.ips.is_empty())
+        .find(|iface| {
+            iface.is_up()
+                && !iface.is_loopback()
+                && !iface.ips.is_empty()
+                && !iface.is_dormant()
+                && iface.is_running()
+                && iface.is_point_to_point()
+        })
         .expect("No valid network interface found");
 
     // Create transport channel for sending and receiving
@@ -52,13 +64,19 @@ pub fn tcp_scan(targets: Vec<IpAddr>, ports: Vec<i32>, timeout: Duration) -> Vec
         let start_time = std::time::Instant::now();
 
         while start_time.elapsed() < timeout {
+            // println!("loop");
             let mut iter = transport::tcp_packet_iter(&mut rx);
 
-            match iter.next() {
+            match iter.next_with_timeout(timeout) {
-                Ok((packet, addr)) => {
+                Ok(Some((packet, addr))) => {
                     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()
+                            );
                             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);
@@ -66,6 +84,9 @@ pub fn tcp_scan(targets: Vec<IpAddr>, ports: Vec<i32>, timeout: Duration) -> Vec
                         }
                     }
                 }
+                Ok(None) => {
+                    break;
+                }
                 Err(_) => {
                     // Just continue on errors
                     thread::sleep(Duration::from_millis(1));
@@ -76,6 +97,8 @@ pub fn tcp_scan(targets: Vec<IpAddr>, ports: Vec<i32>, timeout: Duration) -> Vec
 
     let pb = ProgressBar::new((targets.len() * ports.len()) as u64);
 
+    println!("{:?}", interface.ips);
+
     let source_ip = interface
         .ips
         .iter()
@@ -83,6 +106,11 @@ pub fn tcp_scan(targets: Vec<IpAddr>, ports: Vec<i32>, timeout: Duration) -> Vec
         .expect("No IPv4 address found")
         .ip();
 
+    // let source_ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
+    // let source_ip = IpAddr::V4(Ipv4Addr::new(10, 0, 70, 4));
+
+    println!("Using IP: {}", source_ip.to_string());
+
     for target in &targets {
         for port in &ports {
             // let source_ip = Ipv4Addr::from_bits(random_range(0..=(0xffffffff)));
@@ -111,10 +139,7 @@ pub fn tcp_scan(targets: Vec<IpAddr>, ports: Vec<i32>, timeout: Duration) -> Vec
             );
             tcp_header.set_checksum(checksum);
 
-            match tx.send_to(tcp_header, *target) {
+            send_tcp_packet(&mut tx, tcp_header, target);
-                Ok(_) => {}
-                Err(e) => eprintln!("Failed to send packet: {}", e),
-            }
 
             pb.inc(1);
             thread::sleep(Duration::from_micros(100));
@@ -122,6 +147,7 @@ pub fn tcp_scan(targets: Vec<IpAddr>, ports: Vec<i32>, timeout: Duration) -> Vec
     }
 
     // Wait for receiver to finish
+    // thread::sleep(timeout);
     receiver_handle.join().unwrap();
 
     // Convert results to the return format
@@ -131,6 +157,7 @@ pub fn tcp_scan(targets: Vec<IpAddr>, ports: Vec<i32>, timeout: Duration) -> Vec
         .map(|ip| {
             let mut open_ports = results_map.get(ip).cloned().unwrap_or_default();
             open_ports.sort();
+            open_ports.dedup();
             ScanResult {
                 ip: *ip,
                 open_ports,
@@ -138,3 +165,19 @@ pub fn tcp_scan(targets: Vec<IpAddr>, ports: Vec<i32>, timeout: Duration) -> Vec
         })
         .collect()
 }
+
+fn send_tcp_packet(tx: &mut TransportSender, tcp_header: MutableTcpPacket<'_>, target: &IpAddr) {
+    match tx.send_to(&tcp_header, *target) {
+        Ok(_) => {}
+        Err(e) => {
+            if let Some(code) = e.raw_os_error() {
+                if code == 105 {
+                    thread::sleep(Duration::from_millis(500));
+                    send_tcp_packet(tx, tcp_header, target);
+                }
+            } else {
+                eprintln!("Failed to send packet: {}", e);
+            }
+        }
+    }
+}