diff --git a/Cargo.toml b/Cargo.toml index 18a89bc..7886c00 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/database.rs b/src/database.rs index d85542a..e8f290d 100644 --- a/src/database.rs +++ b/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 values: Vec, // Array of string values + pub id: String, // Row identifier + pub ports: Vec, // 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(); - for (i, value) in self.values.iter().enumerate() { - if i > 0 { - str += ", "; - } - str += format!("{}: \"{}\"", i, value).as_str(); - } - str += "]"; + str += format!("{} - ports: [{}]", self.id, join_nums(&self.ports, ",")).as_str(); str } + pub fn ports_to_string(&self) -> String { + return join_nums(&self.ports, ","); + } +} + +pub fn join_nums(nums: &Vec, sep: &str) -> String { + // 1. Convert numbers to strings + let str_nums: Vec = 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 { + if str.is_empty() { + return vec![]; + }; + + return str + .split(sep) + .map(|n| { + if let Ok(num) = n.parse::() { + 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 - - // 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()); - } + let column_families = vec!["default".to_string(), "ports".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) -> Result<(), Box> { let db = Arc::new(DB::open_cf(&self.options, &self.path, &self.columns)?); - let cf_default = db.cf_handle("default").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 cf_default = db.cf_handle(&self.columns[0]).unwrap(); + let cf_ports = db.cf_handle(&self.columns[1]).unwrap(); 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 = 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) - for (col_idx, value) in row.values.iter().enumerate() { - if col_idx < cf_columns_ref.len() { - // 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(), - ); - } - } + let idx_key = + format!("{}:{}", fast_escape(row.ports_to_string().as_str()), row.id); + batch.put_cf(cf_ports, idx_key.as_bytes(), row.id.as_bytes()); } batch @@ -304,17 +303,19 @@ fn decode_row_binary(key: &str, data: &[u8]) -> Option { 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, 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); diff --git a/src/main.rs b/src/main.rs index 87870ef..a6d15e5 100644 --- a/src/main.rs +++ b/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); } } diff --git a/src/online_scan/online_scan.rs b/src/online_scan/online_scan.rs index fa9e04a..9e5adbb 100644 --- a/src/online_scan/online_scan.rs +++ b/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![ - 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() - }, - ], + ports: vec![], } } } diff --git a/src/port_scan/port_scan.rs b/src/port_scan/port_scan.rs index 4abdd4f..a6fd29d 100644 --- a/src/port_scan/port_scan.rs +++ b/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, sep: &str) -> String { - // 1. Convert numbers to strings - let str_nums: Vec = 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) -} diff --git a/src/port_scan/tcp_scan.rs b/src/port_scan/tcp_scan.rs index e8cbad0..90f49e8 100644 --- a/src/port_scan/tcp_scan.rs +++ b/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::transport::{self, TransportChannelType}; +use pnet::packet::{tcp, Packet}; +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, ports: Vec, timeout: Duration) -> Vec { // 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, ports: Vec, 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() { - Ok((packet, addr)) => { + match iter.next_with_timeout(timeout) { + 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, ports: Vec, 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, ports: Vec, 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, ports: Vec, 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, ports: Vec, timeout: Duration) -> Vec ); tcp_header.set_checksum(checksum); - match tx.send_to(tcp_header, *target) { - Ok(_) => {} - Err(e) => eprintln!("Failed to send packet: {}", e), - } + send_tcp_packet(&mut tx, tcp_header, target); pb.inc(1); thread::sleep(Duration::from_micros(100)); @@ -122,6 +147,7 @@ pub fn tcp_scan(targets: Vec, ports: Vec, 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, ports: Vec, 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, ports: Vec, 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); + } + } + } +}