Merge pull request #135 from RustScan/ba-add-scan-strategy

Add port scanning strategy

src/main.rs

@@ -5,6 +5,9 @@ mod tui;
 mod scanner;
 use scanner::Scanner;
 
+mod port_strategy;
+use port_strategy::PortStrategy;
+
 use colorful::Color;
 use colorful::Colorful;
 use futures::executor::block_on;
@@ -13,7 +16,7 @@ use rlimit::{getrlimit, setrlimit};
 use std::collections::HashMap;
 use std::process::Command;
 use std::{net::IpAddr, time::Duration};
-use structopt::StructOpt;
+use structopt::{clap::arg_enum, StructOpt};
 
 extern crate colorful;
 extern crate dirs;
@@ -28,6 +31,14 @@ const AVERAGE_BATCH_SIZE: rlimit::rlim = 3000;
 #[macro_use]
 extern crate log;
 
+arg_enum! {
+    #[derive(Debug, StructOpt)]
+    pub enum ScanOrder {
+        Serial,
+        Random,
+    }
+}
+
 #[derive(StructOpt, Debug)]
 #[structopt(name = "rustscan", setting = structopt::clap::AppSettings::TrailingVarArg)]
 /// Fast Port Scanner built in Rust.
@@ -59,6 +70,12 @@ struct Opts {
     #[structopt(short, long)]
     ulimit: Option<rlimit::rlim>,
 
+    /// The order of scanning to be performed. The "serial" option will
+    /// scan ports in ascending order while the "random" option will scan
+    /// ports randomly.
+    #[structopt(long, possible_values = &ScanOrder::variants(), case_insensitive = true, default_value = "serial")]
+    scan_order: ScanOrder,
+
     /// The Nmap arguments to run.
     /// To use the argument -A, end RustScan's args with '-- -A'.
     /// Example: 'rustscan -T 1500 127.0.0.1 -- -A -sC'.
@@ -87,11 +104,10 @@ fn main() {
 
     let scanner = Scanner::new(
         &opts.ips,
-        LOWEST_PORT_NUMBER,
-        TOP_PORT_NUMBER,
         batch_size,
         Duration::from_millis(opts.timeout.into()),
         opts.quiet,
+        PortStrategy::pick(LOWEST_PORT_NUMBER, TOP_PORT_NUMBER, opts.scan_order),
     );
 
     let scan_result = block_on(scanner.run());
@@ -272,7 +288,7 @@ fn infer_batch_size(opts: &Opts, ulimit: rlimit::rlim) -> u16 {
 
 #[cfg(test)]
 mod tests {
-    use crate::{adjust_ulimit_size, infer_batch_size, print_opening, Opts};
+    use crate::{adjust_ulimit_size, infer_batch_size, print_opening, Opts, ScanOrder};
     use std::{net::IpAddr, str::FromStr};
 
     #[test]
@@ -284,6 +300,7 @@ mod tests {
             timeout: 1_000,
             ulimit: Some(2_000),
             command: Vec::new(),
+            scan_order: ScanOrder::Serial,
         };
         let batch_size = infer_batch_size(&opts, 120);
 
@@ -299,6 +316,7 @@ mod tests {
             timeout: 1_000,
             ulimit: Some(2_000),
             command: Vec::new(),
+            scan_order: ScanOrder::Serial,
         };
         let batch_size = infer_batch_size(&opts, 9_000);
 
@@ -315,6 +333,7 @@ mod tests {
             timeout: 1_000,
             ulimit: Some(2_000),
             command: Vec::new(),
+            scan_order: ScanOrder::Serial,
         };
         let batch_size = infer_batch_size(&opts, 5_000);
 
@@ -330,6 +349,7 @@ mod tests {
             timeout: 1_000,
             ulimit: Some(2_000),
             command: Vec::new(),
+            scan_order: ScanOrder::Serial,
         };
         let batch_size = adjust_ulimit_size(&opts);
 
@@ -350,6 +370,7 @@ mod tests {
             timeout: 1_000,
             ulimit: None,
             command: Vec::new(),
+            scan_order: ScanOrder::Serial,
         };
 
         infer_batch_size(&opts, 1_000_000);

src/port_strategy/mod.rs

@@ -0,0 +1,74 @@
+mod range_iterator;
+use super::ScanOrder;
+use range_iterator::RangeIterator;
+
+pub enum PortStrategy {
+    Serial(SerialRange),
+    Random(RandomRange),
+}
+
+impl PortStrategy {
+    pub fn pick(start: u16, end: u16, scan_type: ScanOrder) -> Self {
+        match scan_type {
+            ScanOrder::Serial => PortStrategy::Serial(SerialRange { start, end }),
+            ScanOrder::Random => PortStrategy::Random(RandomRange { start, end }),
+        }
+    }
+
+    pub fn order(&self) -> Vec<u16> {
+        match self {
+            PortStrategy::Serial(range) => range.generate(),
+            PortStrategy::Random(range) => range.generate(),
+        }
+    }
+}
+
+trait RangeOrder {
+    fn generate(&self) -> Vec<u16>;
+}
+
+pub struct SerialRange {
+    start: u16,
+    end: u16,
+}
+
+impl RangeOrder for SerialRange {
+    fn generate(&self) -> Vec<u16> {
+        (self.start..self.end).collect()
+    }
+}
+
+pub struct RandomRange {
+    start: u16,
+    end: u16,
+}
+
+impl RangeOrder for RandomRange {
+    fn generate(&self) -> Vec<u16> {
+        RangeIterator::new(self.start.into(), self.end.into()).collect()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::PortStrategy;
+    use crate::ScanOrder;
+
+    #[test]
+    fn serial_strategy() {
+        let strategy = PortStrategy::pick(1, 100, ScanOrder::Serial);
+        let result = strategy.order();
+        let expected_range = (1..100).into_iter().collect::<Vec<u16>>();
+        assert_eq!(expected_range, result);
+    }
+    #[test]
+    fn random_strategy() {
+        let strategy = PortStrategy::pick(1, 100, ScanOrder::Random);
+        let mut result = strategy.order();
+        let expected_range = (1..100).into_iter().collect::<Vec<u16>>();
+        assert_ne!(expected_range, result);
+
+        result.sort();
+        assert_eq!(expected_range, result);
+    }
+}

src/scanner/mod.rs

@@ -1,5 +1,4 @@
-mod range_iterator;
-use range_iterator::RangeIterator;
+use super::PortStrategy;
 
 use async_std::io;
 use async_std::net::TcpStream;
@@ -21,28 +20,25 @@ use std::{
 #[cfg(not(tarpaulin_include))]
 pub struct Scanner {
     hosts: Vec<IpAddr>,
-    start: u16,
-    end: u16,
     batch_size: u16,
     timeout: Duration,
     quiet: bool,
+    port_strategy: PortStrategy,
 }
 
 impl Scanner {
     pub fn new(
         hosts: &[IpAddr],
-        start: u16,
-        end: u16,
         batch_size: u16,
         timeout: Duration,
         quiet: bool,
+        port_strategy: PortStrategy,
     ) -> Self {
         Self {
-            start,
-            end,
             batch_size,
             timeout,
             quiet,
+            port_strategy,
             hosts: hosts.iter().map(|host| host.to_owned()).collect(),
         }
     }
@@ -51,38 +47,26 @@ impl Scanner {
     /// If you want to run RustScan normally, this is the entry point used
     /// Returns all open ports as Vec<u16>
     pub async fn run(&self) -> Vec<SocketAddr> {
-        let range_iterator = RangeIterator::new(self.start.into(), self.end.into());
-        let mut sockets: Vec<SocketAddr> = Vec::new();
+        let ports: Vec<u16> = self.port_strategy.order();
+        let batch_per_host: usize = self.batch_size as usize / self.hosts.len();
         let mut open_sockets: Vec<SocketAddr> = Vec::new();
 
-        for port in range_iterator {
-            for host in &self.hosts {
-                sockets.push(SocketAddr::new(*host, port));
-            }
-
-            if sockets.len() >= self.batch_size.into() {
-                let mut results = self.scan_range(&sockets).await;
-                open_sockets.append(&mut results);
-                sockets = Vec::new();
-            }
-        }
-
-        // This will happen if we have a number of sockets remaining
-        // that is lower than the batch size.
-        if !sockets.is_empty() {
-            let mut results = self.scan_range(&sockets).await;
-            open_sockets.append(&mut results);
+        for batch in ports.chunks(batch_per_host) {
+            let mut sockets = self.scan_ports(batch).await;
+            open_sockets.append(&mut sockets);
         }
 
         open_sockets
     }
 
     /// Given a slice of sockets, scan them all.
     /// Returns a vector of open sockets.
-    async fn scan_range(&self, sockets: &[SocketAddr]) -> Vec<SocketAddr> {
+    async fn scan_ports(&self, ports: &[u16]) -> Vec<SocketAddr> {
         let mut ftrs = FuturesUnordered::new();
-        for socket in sockets {
-            ftrs.push(self.scan_socket(socket));
+        for port in ports {
+            for host in &self.hosts {
+                ftrs.push(self.scan_socket(SocketAddr::new(*host, *port)));
+            }
         }
 
         let mut open_sockets: Vec<SocketAddr> = Vec::new();
@@ -108,8 +92,8 @@ impl Scanner {
     ///     self.scan_port(10:u16)
     ///
     /// Note: `self` must contain `self.host`.
-    async fn scan_socket(&self, socket: &SocketAddr) -> io::Result<SocketAddr> {
-        match self.connect(*socket).await {
+    async fn scan_socket(&self, socket: SocketAddr) -> io::Result<SocketAddr> {
+        match self.connect(socket).await {
             Ok(x) => {
                 // match stream_result.shutdown(Shutdown::Both)
                 info!("Shutting down stream");
@@ -120,7 +104,7 @@ impl Scanner {
                     println!("Open {}", socket.to_string().purple());
                 }
 
-                Ok(*socket)
+                Ok(socket)
             }
             Err(e) => match e.kind() {
                 ErrorKind::Other => {
@@ -155,14 +139,16 @@ impl Scanner {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::ScanOrder;
     use async_std::task::block_on;
     use std::{net::IpAddr, time::Duration};
 
     #[test]
     fn scanner_runs() {
         // Makes sure the program still runs and doesn't panic
         let addrs = vec!["127.0.0.1".parse::<IpAddr>().unwrap()];
-        let scanner = Scanner::new(&addrs, 1, 1_000, 10, Duration::from_millis(100), true);
+        let strategy = PortStrategy::pick(1, 1_000, ScanOrder::Random);
+        let scanner = Scanner::new(&addrs, 10, Duration::from_millis(100), true, strategy);
         block_on(scanner.run());
         // if the scan fails, it wouldn't be able to assert_eq! as it panicked!
         assert_eq!(1, 1);
@@ -171,22 +157,25 @@ mod tests {
     fn ipv6_scanner_runs() {
         // Makes sure the program still runs and doesn't panic
         let addrs = vec!["::1".parse::<IpAddr>().unwrap()];
-        let scanner = Scanner::new(&addrs, 1, 1_000, 10, Duration::from_millis(100), false);
+        let strategy = PortStrategy::pick(1, 1_000, ScanOrder::Random);
+        let scanner = Scanner::new(&addrs, 10, Duration::from_millis(100), true, strategy);
         block_on(scanner.run());
         // if the scan fails, it wouldn't be able to assert_eq! as it panicked!
         assert_eq!(1, 1);
     }
     #[test]
     fn quad_zero_scanner_runs() {
         let addrs = vec!["0.0.0.0".parse::<IpAddr>().unwrap()];
-        let scanner = Scanner::new(&addrs, 1, 1_000, 10, Duration::from_millis(500), true);
+        let strategy = PortStrategy::pick(1, 1_000, ScanOrder::Random);
+        let scanner = Scanner::new(&addrs, 10, Duration::from_millis(100), true, strategy);
         block_on(scanner.run());
         assert_eq!(1, 1);
     }
     #[test]
     fn google_dns_runs() {
         let addrs = vec!["8.8.8.8".parse::<IpAddr>().unwrap()];
-        let scanner = Scanner::new(&addrs, 400, 445, 10, Duration::from_millis(1_500), true);
+        let strategy = PortStrategy::pick(400, 445, ScanOrder::Random);
+        let scanner = Scanner::new(&addrs, 10, Duration::from_millis(100), true, strategy);
         block_on(scanner.run());
         assert_eq!(1, 1);
     }
@@ -196,7 +185,8 @@ mod tests {
         let addrs = vec!["8.8.8.8".parse::<IpAddr>().unwrap()];
 
         // mac should have this automatically scaled down
-        let scanner = Scanner::new(&addrs, 400, 600, 10_000, Duration::from_millis(1500), true);
+        let strategy = PortStrategy::pick(400, 600, ScanOrder::Random);
+        let scanner = Scanner::new(&addrs, 10, Duration::from_millis(100), true, strategy);
         block_on(scanner.run());
         assert_eq!(1, 1);
     }