Skip to content

Commit f74cb45

Browse files
dmitrisbee-san
andauthored
socket_iterator: use itertools.iproducts for ip x ports set (#231)
Replace the current manual iterator-like implementation with the call to the itertools.iproduct macro to generate the cartesian product of the IPs and ports. stdlib reference: https://docs.rs/itertools/0.9.0/itertools/macro.iproduct.html Co-authored-by: Bee <[email protected]>
1 parent 492d2fb commit f74cb45

File tree

3 files changed

+34
-25
lines changed

3 files changed

+34
-25
lines changed

Cargo.lock

+16
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Cargo.toml

+1
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ toml = "0.5.6"
3434
serde = "1.0.116"
3535
serde_derive = "1.0.116"
3636
cidr-utils = "0.5.0"
37+
itertools = "0.9.0"
3738

3839
[package.metadata.deb]
3940
depends = "$auto, nmap"

src/scanner/socket_iterator.rs

+17-25
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,17 @@
1+
use itertools::{iproduct, Product};
12
use std::net::{IpAddr, SocketAddr};
23

34
pub struct SocketIterator<'s> {
4-
ips: &'s [IpAddr],
5-
ports: &'s [u16],
6-
ip_idx: usize,
7-
ip_len: usize,
8-
port_idx: usize,
9-
port_len: usize,
5+
// product_it is a cartesian product iterator over
6+
// the slices of ports and IP addresses.
7+
//
8+
// The IP/port order is intentionally reversed here since we want
9+
// the itertools::iproduct! macro below to generate the pairs with
10+
// all the IPs for one port before moving on to the next one
11+
// ("hold the port, go through all the IPs, then advance the port...").
12+
// See also the comments in the iterator implementation for an example.
13+
product_it:
14+
Product<Box<std::slice::Iter<'s, u16>>, Box<std::slice::Iter<'s, std::net::IpAddr>>>,
1015
}
1116

1217
/// An iterator that receives a slice of IPs and ports and returns a Socket
@@ -16,13 +21,10 @@ pub struct SocketIterator<'s> {
1621
/// generating a vector containing all these combinations.
1722
impl<'s> SocketIterator<'s> {
1823
pub fn new(ips: &'s [IpAddr], ports: &'s [u16]) -> Self {
24+
let ports_it = Box::new(ports.into_iter());
25+
let ips_it = Box::new(ips.into_iter());
1926
Self {
20-
ip_idx: 0,
21-
ip_len: ips.len(),
22-
port_idx: 0,
23-
port_len: ports.len(),
24-
ips,
25-
ports,
27+
product_it: iproduct!(ports_it, ips_it),
2628
}
2729
}
2830
}
@@ -41,20 +43,10 @@ impl<'s> Iterator for SocketIterator<'s> {
4143
/// it.next(); // 192.168.0.1:443
4244
/// it.next(); // None
4345
fn next(&mut self) -> Option<Self::Item> {
44-
if self.port_idx == self.port_len {
45-
return None;
46+
match self.product_it.next() {
47+
None => None,
48+
Some((port, ip)) => Some(SocketAddr::new(*ip, *port)),
4649
}
47-
48-
self.ip_idx = self.ip_idx % self.ip_len;
49-
50-
let socket = SocketAddr::new(self.ips[self.ip_idx], self.ports[self.port_idx]);
51-
self.ip_idx += 1;
52-
53-
if self.ip_idx == self.ip_len {
54-
self.port_idx += 1;
55-
}
56-
57-
Some(socket)
5850
}
5951
}
6052

0 commit comments

Comments
 (0)