Skip to content

Commit 46be54f

Browse files
committed
Add tries to scan_socket.
By default the number of try is 1, it can be set on the command line by using --tries option. This should close #38.
1 parent 5b35a09 commit 46be54f

File tree

4 files changed

+116
-41
lines changed

4 files changed

+116
-41
lines changed

README.md

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -242,6 +242,7 @@ OPTIONS:
242242
ascending order while the "random" option will scan ports randomly [default:
243243
serial] [possible values: Serial, Random]
244244
-t, --timeout <timeout> The timeout in milliseconds before a port is assumed to be closed [default: 1500]
245+
--tries <tries> The number of tries before a port is assumed to be closed [default: 1]
245246
-u, --ulimit <ulimit> Automatically ups the ULIMIT with the value you provided
246247

247248
ARGS:
@@ -260,7 +261,7 @@ Your operating system may not support this, but it is worth it to play around an
260261

261262
## Configuration file
262263

263-
This binary accepts a configuration file that is read from the home directory of the user. It follows the TOML format
264+
This binary accepts a configuration file, named `.rustscan.toml`, that is read from the home directory of the user. It follows the TOML format
264265
and accepts the following fields:
265266

266267
- `addresses`
@@ -272,6 +273,7 @@ and accepts the following fields:
272273
- `greppable`
273274
- `batch-size`
274275
- `timeout`
276+
- `tries`
275277
- `ulimit`
276278

277279
### Format example
@@ -286,6 +288,7 @@ accessible = true
286288
scan_order = "Serial"
287289
batch_size = 1000
288290
timeout = 1000
291+
tries = 3
289292
ulimit = 1000
290293
```
291294

src/input.rs

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,10 @@ pub struct Opts {
9595
#[structopt(short, long, default_value = "1500")]
9696
pub timeout: u32,
9797

98+
/// The number of tries before a port is assumed to be closed.
99+
#[structopt(long, default_value = "1")]
100+
pub tries: u8,
101+
98102
/// Automatically ups the ULIMIT with the value you provided.
99103
#[structopt(short, long)]
100104
pub ulimit: Option<rlimit::rlim>,
@@ -153,7 +157,9 @@ impl Opts {
153157
}
154158
}
155159

156-
merge_required!(addresses, greppable, accessible, batch_size, timeout, scan_order, command);
160+
merge_required!(
161+
addresses, greppable, accessible, batch_size, timeout, tries, scan_order, command
162+
);
157163
}
158164

159165
fn merge_optional(&mut self, config: &Config) {
@@ -195,6 +201,7 @@ pub struct Config {
195201
accessible: Option<bool>,
196202
batch_size: Option<u16>,
197203
timeout: Option<u32>,
204+
tries: Option<u8>,
198205
no_nmap: Option<bool>,
199206
ulimit: Option<rlimit::rlim>,
200207
scan_order: Option<ScanOrder>,
@@ -252,6 +259,7 @@ mod tests {
252259
greppable: Some(true),
253260
batch_size: Some(25_000),
254261
timeout: Some(1_000),
262+
tries: Some(1),
255263
ulimit: None,
256264
no_nmap: Some(false),
257265
command: Some(vec!["-A".to_owned()]),
@@ -270,6 +278,7 @@ mod tests {
270278
greppable: true,
271279
batch_size: 0,
272280
timeout: 0,
281+
tries: 0,
273282
ulimit: None,
274283
command: vec![],
275284
accessible: false,

src/main.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ fn main() {
7777
&ips,
7878
batch_size,
7979
Duration::from_millis(opts.timeout.into()),
80+
opts.tries,
8081
opts.greppable,
8182
PortStrategy::pick(opts.range, opts.ports, opts.scan_order),
8283
opts.accessible,

src/scanner/mod.rs

Lines changed: 101 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use super::PortStrategy;
1+
use super::{warning, PortStrategy};
22

33
mod socket_iterator;
44
use socket_iterator::SocketIterator;
@@ -11,6 +11,7 @@ use futures::stream::FuturesUnordered;
1111
use std::{
1212
io::ErrorKind,
1313
net::{IpAddr, Shutdown, SocketAddr},
14+
num::NonZeroU8,
1415
time::Duration,
1516
};
1617

@@ -26,6 +27,7 @@ pub struct Scanner {
2627
ips: Vec<IpAddr>,
2728
batch_size: u16,
2829
timeout: Duration,
30+
tries: NonZeroU8,
2931
greppable: bool,
3032
port_strategy: PortStrategy,
3133
accessible: bool,
@@ -36,13 +38,20 @@ impl Scanner {
3638
ips: &[IpAddr],
3739
batch_size: u16,
3840
timeout: Duration,
41+
tries: u8,
3942
greppable: bool,
4043
port_strategy: PortStrategy,
4144
accessible: bool,
4245
) -> Self {
4346
Self {
4447
batch_size,
4548
timeout,
49+
tries: if tries == 0 {
50+
warning!("tries is 0, corrected to 1!");
51+
NonZeroU8::new(1).unwrap()
52+
} else {
53+
NonZeroU8::new(tries).expect("tries cannot be 0!")
54+
},
4655
greppable,
4756
port_strategy,
4857
ips: ips.iter().map(|ip| ip.to_owned()).collect(),
@@ -87,54 +96,67 @@ impl Scanner {
8796
open_sockets
8897
}
8998

90-
/// Given a port, scan it.
99+
/// Given a socket, scan it self.tries times.
91100
/// Turns the address into a SocketAddr
92101
/// Deals with the <result> type
93102
/// If it experiences error ErrorKind::Other then too many files are open and it Panics!
94-
/// ese any other error, it returns the error in Result as a string
95-
/// If no errors occur, it returns the port number in Result to signify the port is open.
103+
/// Else any other error, it returns the error in Result as a string
104+
/// If no errors occur, it returns the port number in Result to signify the port is open.
96105
/// This function mainly deals with the logic of Results handling.
97106
/// # Example
98107
///
99-
/// self.scan_port(10:u16)
108+
/// self.scan_socket(socket)
100109
///
101110
/// Note: `self` must contain `self.ip`.
102111
async fn scan_socket(&self, socket: SocketAddr) -> io::Result<SocketAddr> {
103-
match self.connect(socket).await {
104-
Ok(x) => {
105-
debug!(
106-
"Connection was successful, shutting down stream {}",
107-
&socket
108-
);
109-
match x.shutdown(Shutdown::Both) {
110-
Err(e) => debug!("Shutdown stream error {}", &e),
111-
_ => {}
112-
}
113-
if !self.greppable {
114-
if self.accessible {
115-
println!("Open {}", socket.to_string());
116-
} else {
117-
println!("Open {}", socket.to_string().purple());
112+
let tries = self.tries.get();
113+
114+
debug!("self.tries: {}", tries);
115+
116+
for nr_try in 1..=tries {
117+
debug!("Try number: {}", nr_try);
118+
119+
match self.connect(socket).await {
120+
Ok(x) => {
121+
debug!(
122+
"Connection was successful, shutting down stream {}",
123+
&socket
124+
);
125+
if let Err(e) = x.shutdown(Shutdown::Both) {
126+
debug!("Shutdown stream error {}", &e);
127+
}
128+
if !self.greppable {
129+
if self.accessible {
130+
println!("Open {}", socket.to_string());
131+
} else {
132+
println!("Open {}", socket.to_string().purple());
133+
}
118134
}
135+
136+
debug!("Return Ok after {} tries", nr_try);
137+
return Ok(socket);
119138
}
139+
Err(e) => {
140+
let error_string = e.to_string();
120141

121-
Ok(socket)
122-
}
123-
Err(e) => match e.kind() {
124-
ErrorKind::Other => {
125-
if e.to_string().contains("No route to host")
126-
|| e.to_string().contains("Network is unreachable")
127-
{
128-
debug!("Socket connect error: {} {}", &e.to_string(), &socket);
129-
Err(io::Error::new(io::ErrorKind::Other, e.to_string()))
130-
} else {
131-
debug!("Socket connect error: {} {}", &e.to_string(), &socket);
132-
panic!("Too many open files. Please reduce batch size. The default is 5000. Try -b 2500.");
142+
if e.kind() == ErrorKind::Other {
143+
debug!("Socket connect error: {} {}", &error_string, &socket);
144+
145+
if !error_string.contains("No route to host")
146+
&& !error_string.contains("Network is unreachable")
147+
{
148+
panic!("Too many open files. Please reduce batch size. The default is 5000. Try -b 2500.");
149+
}
150+
}
151+
152+
if nr_try == tries {
153+
debug!("Return Err after {} tries", tries);
154+
return Err(io::Error::new(io::ErrorKind::Other, error_string));
133155
}
134156
}
135-
_ => Err(io::Error::new(io::ErrorKind::Other, e.to_string())),
136-
},
157+
};
137158
}
159+
unreachable!();
138160
}
139161

140162
/// Performs the connection to the socket with timeout
@@ -174,7 +196,15 @@ mod tests {
174196
end: 1_000,
175197
};
176198
let strategy = PortStrategy::pick(Some(range), None, ScanOrder::Random);
177-
let scanner = Scanner::new(&addrs, 10, Duration::from_millis(100), true, strategy, true);
199+
let scanner = Scanner::new(
200+
&addrs,
201+
10,
202+
Duration::from_millis(100),
203+
1,
204+
true,
205+
strategy,
206+
true,
207+
);
178208
block_on(scanner.run());
179209
// if the scan fails, it wouldn't be able to assert_eq! as it panicked!
180210
assert_eq!(1, 1);
@@ -188,7 +218,15 @@ mod tests {
188218
end: 1_000,
189219
};
190220
let strategy = PortStrategy::pick(Some(range), None, ScanOrder::Random);
191-
let scanner = Scanner::new(&addrs, 10, Duration::from_millis(100), true, strategy, true);
221+
let scanner = Scanner::new(
222+
&addrs,
223+
10,
224+
Duration::from_millis(100),
225+
1,
226+
true,
227+
strategy,
228+
true,
229+
);
192230
block_on(scanner.run());
193231
// if the scan fails, it wouldn't be able to assert_eq! as it panicked!
194232
assert_eq!(1, 1);
@@ -201,7 +239,15 @@ mod tests {
201239
end: 1_000,
202240
};
203241
let strategy = PortStrategy::pick(Some(range), None, ScanOrder::Random);
204-
let scanner = Scanner::new(&addrs, 10, Duration::from_millis(100), true, strategy, true);
242+
let scanner = Scanner::new(
243+
&addrs,
244+
10,
245+
Duration::from_millis(100),
246+
1,
247+
true,
248+
strategy,
249+
true,
250+
);
205251
block_on(scanner.run());
206252
assert_eq!(1, 1);
207253
}
@@ -213,7 +259,15 @@ mod tests {
213259
end: 445,
214260
};
215261
let strategy = PortStrategy::pick(Some(range), None, ScanOrder::Random);
216-
let scanner = Scanner::new(&addrs, 10, Duration::from_millis(100), true, strategy, true);
262+
let scanner = Scanner::new(
263+
&addrs,
264+
10,
265+
Duration::from_millis(100),
266+
1,
267+
true,
268+
strategy,
269+
true,
270+
);
217271
block_on(scanner.run());
218272
assert_eq!(1, 1);
219273
}
@@ -228,7 +282,15 @@ mod tests {
228282
end: 600,
229283
};
230284
let strategy = PortStrategy::pick(Some(range), None, ScanOrder::Random);
231-
let scanner = Scanner::new(&addrs, 10, Duration::from_millis(100), true, strategy, true);
285+
let scanner = Scanner::new(
286+
&addrs,
287+
10,
288+
Duration::from_millis(100),
289+
1,
290+
true,
291+
strategy,
292+
true,
293+
);
232294
block_on(scanner.run());
233295
assert_eq!(1, 1);
234296
}

0 commit comments

Comments
 (0)