From 397753e90ee64974fddf5c56a84a7bbe275a2e1e Mon Sep 17 00:00:00 2001 From: gongzhengyang Date: Fri, 8 Dec 2023 15:12:39 +0800 Subject: [PATCH 1/3] feat: add error handles --- rscanner/Cargo.toml | 4 +-- rscanner/src/consts.rs | 0 rscanner/src/err.rs | 40 ++++++++++++++++++++------ rscanner/src/execute/arp.rs | 24 ++++++++++------ rscanner/src/execute/icmp/interface.rs | 9 ++++-- rscanner/src/execute/udp.rs | 30 ++++++++++++------- rscanner/src/interfaces.rs | 10 +++++-- rscanner/src/lib.rs | 1 + rscanner/src/monitor.rs | 11 ++++--- rscanner/src/setting/command.rs | 34 +++++++++++----------- rscanner/src/setting/parse.rs | 20 ++++++------- src/main.rs | 12 ++++---- 12 files changed, 120 insertions(+), 75 deletions(-) create mode 100644 rscanner/src/consts.rs diff --git a/rscanner/Cargo.toml b/rscanner/Cargo.toml index 0b8748a..e4a5e6a 100644 --- a/rscanner/Cargo.toml +++ b/rscanner/Cargo.toml @@ -17,12 +17,12 @@ async-trait = "0.1" clap = { version = "4.2", features = ["derive"] } hashbrown = "0.14" ipnetwork = "0.20" -itertools = "0.11" +itertools = "0.12" pnet = "0.34" pnet_transport = "0.34" rand = "0.8" rlimit = "0.10" serde = { version = "1", features = ["derive"] } -thiserror = "1.0" +snafu = "0.7" tokio = { version = "1", features = ["full", "tracing"] } tracing = "0.1" diff --git a/rscanner/src/consts.rs b/rscanner/src/consts.rs new file mode 100644 index 0000000..e69de29 diff --git a/rscanner/src/err.rs b/rscanner/src/err.rs index 92fa10e..cd6b3d8 100644 --- a/rscanner/src/err.rs +++ b/rscanner/src/err.rs @@ -1,18 +1,42 @@ use serde::Serialize; -use thiserror::Error; +use snafu::Location; +use snafu::prelude::*; -#[derive(Error, Debug, Serialize, PartialEq)] +#[derive(Snafu, Debug, Serialize, PartialEq)] +#[snafu(visibility(pub))] pub enum APPError { - #[error("split failed because ip or port format is error")] - SplitIpPortFailed, - #[error("port args format is error")] - PortFormatError, - #[error("port couldn't be empty")] + #[snafu(display("split failed because ip or port format is error {}", value))] + SplitIpPortFailed { + location: Location, + value: String + }, + + #[snafu(display("port args format is error {}", value))] + PortFormat { + location: Location, + value: String, + }, + + #[snafu(display("port couldn't be empty"))] PortIsEmpty, - #[error("unix ulimits soft limit is bigger than hard limit")] + + #[snafu(display("unix ulimits soft limit is bigger than hard limit"))] ULimitSoftBiggerThanHard, + + #[snafu(display("Option value is None"))] + OptionEmpty { + location: Location + }, + + #[snafu(display("common io error {}", source))] + CommonIo { + location: Location, + source: std::io::Error, + }, } +pub type Result = std::result::Result; + #[cfg(test)] mod tests { use crate::err::APPError; diff --git a/rscanner/src/execute/arp.rs b/rscanner/src/execute/arp.rs index 583aa13..525ba4d 100644 --- a/rscanner/src/execute/arp.rs +++ b/rscanner/src/execute/arp.rs @@ -1,13 +1,17 @@ use std::net::{IpAddr, Ipv4Addr}; use std::time::Duration; -use pnet::datalink::Config; -use pnet::datalink::{Channel, DataLinkReceiver, DataLinkSender, MacAddr, NetworkInterface}; -use pnet::packet::arp::{ArpHardwareTypes, ArpOperations, ArpPacket, MutableArpPacket}; -use pnet::packet::ethernet::EtherTypes; -use pnet::packet::ethernet::MutableEthernetPacket; -use pnet::packet::{MutablePacket, Packet}; +use pnet::datalink::{ + Channel, Config, DataLinkReceiver, DataLinkSender, MacAddr, NetworkInterface +}; +use pnet::packet::{ + arp::{ArpHardwareTypes, ArpOperations, ArpPacket, MutableArpPacket}, + ethernet::{EtherTypes, MutableEthernetPacket}, + MutablePacket, Packet +}; +use snafu::OptionExt; +use crate::err::OptionEmptySnafu; use crate::interfaces::{get_interface_ipv4, interface_normal_running}; use crate::monitor; use crate::setting::command::ScanOpts; @@ -17,7 +21,7 @@ async fn send_arp_packets( source_ip: Ipv4Addr, target_ips: Vec, ) -> anyhow::Result<()> { - let mac = interface.mac.unwrap(); + let mac = interface.mac.context(OptionEmptySnafu)?; let (mut sender, _) = get_sender_receiver(&interface); tracing::info!("Sent ARP request with interface: {interface:?}"); for target_ip in target_ips { @@ -26,14 +30,16 @@ async fn send_arp_packets( continue; } let mut ethernet_buffer = [0u8; 42]; - let mut ethernet_packet = MutableEthernetPacket::new(&mut ethernet_buffer).unwrap(); + let mut ethernet_packet = MutableEthernetPacket::new( + &mut ethernet_buffer).context(OptionEmptySnafu)?; ethernet_packet.set_destination(MacAddr::broadcast()); ethernet_packet.set_source(mac); ethernet_packet.set_ethertype(EtherTypes::Arp); let mut arp_buffer = [0u8; 28]; - let mut arp_packet = MutableArpPacket::new(&mut arp_buffer).unwrap(); + let mut arp_packet = MutableArpPacket::new( + &mut arp_buffer).context(OptionEmptySnafu)?; arp_packet.set_hardware_type(ArpHardwareTypes::Ethernet); arp_packet.set_protocol_type(EtherTypes::Ipv4); diff --git a/rscanner/src/execute/icmp/interface.rs b/rscanner/src/execute/icmp/interface.rs index 23125cb..85d5378 100644 --- a/rscanner/src/execute/icmp/interface.rs +++ b/rscanner/src/execute/icmp/interface.rs @@ -8,13 +8,15 @@ use pnet::packet::{ ipv4::MutableIpv4Packet, }; use pnet::packet::{MutablePacket, Packet}; +use snafu::OptionExt; +use crate::err::OptionEmptySnafu; use crate::interfaces; use crate::interfaces::interface_normal_running; - +use crate::err::Result; use super::common; -pub fn send_with_interface(target_ip: Ipv4Addr) { +pub fn send_with_interface(target_ip: Ipv4Addr) -> Result<()> { tracing::debug!("{target_ip} send by specific interface"); for interface in pnet::datalink::interfaces() { if !interface_normal_running(&interface) { @@ -22,7 +24,7 @@ pub fn send_with_interface(target_ip: Ipv4Addr) { } if let Some(source_ip) = interfaces::get_interface_ipv4(&interface) { let mut header = [0u8; common::ICMP_LEN]; - let mut icmp_packet = MutableEchoRequestPacket::new(&mut header).unwrap(); + let mut icmp_packet = MutableEchoRequestPacket::new(&mut header).context(OptionEmptySnafu)?; common::modify_icmp_packet(&mut icmp_packet); let mut ipv4_header = [0u8; common::IPV4_LEN]; @@ -58,4 +60,5 @@ pub fn send_with_interface(target_ip: Ipv4Addr) { .unwrap(); } } + Ok(()) } diff --git a/rscanner/src/execute/udp.rs b/rscanner/src/execute/udp.rs index 3f65537..8d445e6 100644 --- a/rscanner/src/execute/udp.rs +++ b/rscanner/src/execute/udp.rs @@ -1,17 +1,25 @@ use std::net::{IpAddr, SocketAddr}; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; use std::time::Duration; use async_trait::async_trait; use hashbrown::{HashMap, HashSet}; -use pnet::packet::icmp::destination_unreachable::IcmpCodes; -use pnet::packet::icmp::echo_reply::EchoReplyPacket; -use pnet::packet::ipv4::Ipv4Packet; -use pnet::packet::udp::UdpPacket; -use pnet::packet::Packet; -use tokio::net::UdpSocket; -use tokio::sync::OnceCell; -use tokio::time::Instant; +use pnet::packet::{ + icmp::{ + destination_unreachable::IcmpCodes, + echo_reply::EchoReplyPacket + }, + ipv4::Ipv4Packet, + Packet, + udp::UdpPacket +}; +use tokio::{ + net::UdpSocket, + sync::{ + Mutex, OnceCell + }, + time::Instant +}; use crate::setting::command::ScanOpts; @@ -60,12 +68,12 @@ async fn get_socket_manager() -> &'static Arc>> { } async fn add_socket_to_manager(socket: SocketAddr) { - let mut socket_manager = get_socket_manager().await.lock().unwrap(); + let mut socket_manager = get_socket_manager().await.lock().await; socket_manager.insert(socket); } async fn remove_socket_from_manager(socket: &SocketAddr) { - let mut socket_manager = get_socket_manager().await.lock().unwrap(); + let mut socket_manager = get_socket_manager().await.lock().await; let remove = socket_manager.remove(socket); tracing::debug!("remove {socket} {remove}"); } diff --git a/rscanner/src/interfaces.rs b/rscanner/src/interfaces.rs index 7eb551e..58b26bf 100644 --- a/rscanner/src/interfaces.rs +++ b/rscanner/src/interfaces.rs @@ -1,6 +1,9 @@ use std::net::{IpAddr, Ipv4Addr}; use pnet::datalink::NetworkInterface; +use snafu::OptionExt; + +use crate::err::{OptionEmptySnafu, Result}; pub fn get_interface_ipv4(interface: &NetworkInterface) -> Option { interface @@ -13,11 +16,12 @@ pub fn get_interface_ipv4(interface: &NetworkInterface) -> Option { }) } -pub fn get_interface_by_name(interface_name: &str) -> NetworkInterface { - pnet::datalink::interfaces() +pub fn get_interface_by_name(interface_name: &str) -> Result { + Ok(pnet::datalink::interfaces() .into_iter() .find(|interface| interface.name == interface_name) - .unwrap() + .context(OptionEmptySnafu)? + } } pub fn interface_normal_running(interface: &NetworkInterface) -> bool { diff --git a/rscanner/src/lib.rs b/rscanner/src/lib.rs index 88ae9e3..e35b892 100644 --- a/rscanner/src/lib.rs +++ b/rscanner/src/lib.rs @@ -4,3 +4,4 @@ pub mod interfaces; pub mod monitor; pub mod performance; pub mod setting; +pub mod consts; diff --git a/rscanner/src/monitor.rs b/rscanner/src/monitor.rs index 2d5a499..e8bc2df 100644 --- a/rscanner/src/monitor.rs +++ b/rscanner/src/monitor.rs @@ -1,8 +1,11 @@ use std::net::IpAddr; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; use hashbrown::HashSet; -use tokio::sync::OnceCell; +use tokio::sync::{ + Mutex, + OnceCell +}; static RECEIVE_PACKETS: OnceCell>>> = OnceCell::const_new(); @@ -16,11 +19,11 @@ pub async fn is_addr_received(addr: &IpAddr) -> bool { receive_packets_handle() .await .lock() - .unwrap() + .await .contains(addr) } pub async fn add_receive_ipaddr(addr: IpAddr) { - let mut receive_handle = receive_packets_handle().await.lock().unwrap(); + let mut receive_handle = receive_packets_handle().await.lock().await; receive_handle.insert(addr); } diff --git a/rscanner/src/setting/command.rs b/rscanner/src/setting/command.rs index 7cd4686..a7a8906 100644 --- a/rscanner/src/setting/command.rs +++ b/rscanner/src/setting/command.rs @@ -8,7 +8,7 @@ use crate::setting::sockets_iter::SocketIterator; use super::parse::{parse_hosts, parse_ports}; #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum, Debug, Default)] -pub enum Executes { +pub enum ScanType { #[default] /// send icmp to discover hosts reply icmp packets,all packets build by rust purely /// use can be `sudo rscan icmp 1.1.1.1/24` @@ -27,40 +27,38 @@ pub enum Executes { #[derive(Parser, Debug, Default, Clone)] #[command(author, version, about, long_about = None)] pub struct ScanOpts { - /// choose a executor + /// choose scan type #[arg(value_enum)] - pub execute: Executes, + pub execute: ScanType, - /// A list of comma separated CIDRs, IPs, or hosts to be scanned. Example: 1.1.1.1,2.2.2.2/24 + /// A list of comma separated CIDRs, IPs, or hosts to be scanned. + /// Eg: 1.1.1.1,2.2.2.2/24 #[arg(value_parser = parse_hosts)] pub hosts: Arc>, - /// A list of comma separed ports to be scanned. Example: 80,443,1-100,1-7. + /// A list of comma separated ports to be scanned. Example: 80,443,1-100,1-7. #[arg(short, long, value_parser = parse_ports, default_value = "")] pub ports: Arc>, - /// The batch size for port scanning, it increases or slows the speed of - /// scanning. Depends on the open file limit of your OS. If you do 65535 - /// it will do every port at the same time. Although, your OS may not - /// support this. - #[arg(long, default_value_t = 64)] + /// The batch size for port scanning. + /// It increases or slows the speed of scanning. + /// Depends on the open file limit of your OS. + /// If you do 65535, it will do every port at the same time. + /// Although, your OS may not support this. + #[arg(long, default_value_t = 1024)] pub batch_size: usize, /// The global timeout in seconds before a port is assumed to be closed. - #[arg(long, default_value_t = 30)] + #[arg(long, default_value_t = 10)] pub timeout: u64, /// The number of retries for sending icmp,tcp,udp packets to remote host - #[arg(long, default_value_t = 1)] + #[arg(long, default_value_t = 3)] pub retry: u64, /// The seconds retry interval when retry is set bigger than 0 - #[arg(long, default_value_t = 1)] - pub retry_interval: u64, - - /// every single operation timeout, tcp connect timeout ro udp timeout #[arg(long, default_value_t = 3)] - pub per_timeout: u64, + pub retry_interval: u64, } impl ScanOpts { @@ -84,7 +82,7 @@ mod tests { .map(|x| Ipv4Addr::from_str(x).unwrap()) .collect::>(); let scan_opts = ScanOpts { - execute: Executes::Tcp, + execute: ScanType::Tcp, hosts: Arc::new(hosts.clone()), ports: Arc::new(ports.clone()), ..Default::default() diff --git a/rscanner/src/setting/parse.rs b/rscanner/src/setting/parse.rs index 15bda61..1d6a52d 100644 --- a/rscanner/src/setting/parse.rs +++ b/rscanner/src/setting/parse.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use ipnetwork::Ipv4Network; -use crate::err::APPError; +use crate::err::{APPError, PortFormatSnafu}; /// parse input like /// single ip `192.168.1.1` @@ -14,18 +14,14 @@ pub fn parse_hosts(value: &str) -> anyhow::Result>> { let splits = value.split(','); let mut full_values = vec![]; for i in splits { - let results = parse_ipv4_cidr(i)?; + let results = Ipv4Network::from_str(i)? + .into_iter() + .collect::>(); full_values.extend(results); } Ok(Arc::new(full_values)) } -pub fn parse_ipv4_cidr(value: &str) -> anyhow::Result> { - Ok(Ipv4Network::from_str(value)? - .into_iter() - .collect::>()) -} - pub fn parse_ports(input: &str) -> anyhow::Result>> { if input.eq("") { return Ok(Arc::new(vec![])); @@ -50,7 +46,9 @@ pub fn parse_ports_range(input: &str) -> anyhow::Result> { let result = (*start..*end).collect::>(); return Ok(result); } - Err(APPError::PortFormatError.into()) + PortFormatSnafu { + value: input + }.into() } #[cfg(test)] @@ -75,7 +73,7 @@ mod tests { #[test] fn ip_cidr() { - let result = parse_ipv4_cidr("1.1.1.1/30").unwrap(); + let result = parse_hosts("1.1.1.1/30").unwrap(); assert_eq!(result.len(), 4); let min = u32::from(Ipv4Addr::from_str("1.1.1.0").unwrap()); let mut values = vec![]; @@ -88,7 +86,7 @@ mod tests { #[test] fn ip_single() { assert_eq!( - parse_ipv4_cidr("192.168.1.1").unwrap(), + parse_hosts("192.168.1.1").unwrap(), vec![Ipv4Addr::from_str("192.168.1.1").unwrap()] ); } diff --git a/src/main.rs b/src/main.rs index 32395c7..878dedf 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,7 +4,7 @@ use clap::Parser; use rscanner::execute; use rscanner::execute::common::SocketScanner; -use rscanner::setting::command::{Executes, ScanOpts}; +use rscanner::setting::command::{ScanType, ScanOpts}; #[tokio::main] async fn main() { @@ -17,13 +17,13 @@ async fn main() { tracing::info!("waiting for {} seconds", timeout); match scan_opts.execute { - Executes::Icmp => { + ScanType::Icmp => { tracing::info!("execute icmp scan"); tokio::spawn(async move { execute::icmp::scan(scan_opts.clone()).await.unwrap(); }); } - Executes::Tcp => { + ScanType::Tcp => { tracing::info!("execute tcp scan"); tokio::spawn(async move { execute::tcp::TcpSocketScanner::scan(scan_opts) @@ -31,7 +31,7 @@ async fn main() { .unwrap(); }); } - Executes::Udp => { + ScanType::Udp => { tracing::info!("execute udp scan"); tracing::warn!( "udp scan based on icmp reply with Port Unreachable with udp packets,\ @@ -43,13 +43,13 @@ async fn main() { .unwrap(); }); } - Executes::Arp => { + ScanType::Arp => { tracing::info!("execute arp scan"); tokio::spawn(async move { execute::arp::scan(scan_opts).await.unwrap(); }); } - Executes::Show => { + ScanType::Show => { println!("hosts len {}", scan_opts.hosts.len()); let socket_iter = scan_opts.iter_sockets().unwrap(); for socket in socket_iter { From e718808e8130f3f5d47991a3125053369378e2a1 Mon Sep 17 00:00:00 2001 From: gongzhengyang Date: Tue, 12 Dec 2023 15:07:18 +0800 Subject: [PATCH 2/3] fix: fix compile erro --- rscanner/Makefile | 2 +- rscanner/src/consts.rs | 1 + rscanner/src/err.rs | 40 +++++++------------------- rscanner/src/execute/arp.rs | 11 ++++--- rscanner/src/execute/common.rs | 4 +-- rscanner/src/execute/icmp/interface.rs | 9 +++--- rscanner/src/execute/icmp/send.rs | 13 ++++----- rscanner/src/execute/udp.rs | 23 ++++++--------- rscanner/src/interfaces.rs | 5 ++-- rscanner/src/lib.rs | 2 +- rscanner/src/monitor.rs | 11 ++----- rscanner/src/setting/parse.rs | 29 ++++++++++--------- 12 files changed, 59 insertions(+), 91 deletions(-) diff --git a/rscanner/Makefile b/rscanner/Makefile index 05e1f46..5a357c3 100644 --- a/rscanner/Makefile +++ b/rscanner/Makefile @@ -2,7 +2,7 @@ check: cargo fmt cargo tomlfmt - cargo clippy --all-targets + cargo clippy cargo install --locked cargo-outdated cargo outdated -R cargo install cargo-udeps --locked diff --git a/rscanner/src/consts.rs b/rscanner/src/consts.rs index e69de29..8b13789 100644 --- a/rscanner/src/consts.rs +++ b/rscanner/src/consts.rs @@ -0,0 +1 @@ + diff --git a/rscanner/src/err.rs b/rscanner/src/err.rs index cd6b3d8..070f355 100644 --- a/rscanner/src/err.rs +++ b/rscanner/src/err.rs @@ -1,19 +1,20 @@ -use serde::Serialize; -use snafu::Location; use snafu::prelude::*; +use snafu::Location; +use std::num::ParseIntError; -#[derive(Snafu, Debug, Serialize, PartialEq)] +#[derive(Snafu, Debug)] #[snafu(visibility(pub))] pub enum APPError { #[snafu(display("split failed because ip or port format is error {}", value))] - SplitIpPortFailed { - location: Location, - value: String - }, + SplitIpPortFailed { location: Location, value: String }, #[snafu(display("port args format is error {}", value))] - PortFormat { + PortFormat { location: Location, value: String }, + + #[snafu(display("port value parse failed {}", source))] + PortParse { location: Location, + source: ParseIntError, value: String, }, @@ -24,9 +25,7 @@ pub enum APPError { ULimitSoftBiggerThanHard, #[snafu(display("Option value is None"))] - OptionEmpty { - location: Location - }, + OptionEmpty { location: Location }, #[snafu(display("common io error {}", source))] CommonIo { @@ -36,22 +35,3 @@ pub enum APPError { } pub type Result = std::result::Result; - -#[cfg(test)] -mod tests { - use crate::err::APPError; - - #[test] - fn error_response() { - for i in 0..3 { - println!("{:?}", result_with_value(i)); - } - } - - fn result_with_value(value: u8) -> anyhow::Result<()> { - if value.eq(&1) { - return Err(APPError::SplitIpPortFailed.into()); - } - Ok(()) - } -} diff --git a/rscanner/src/execute/arp.rs b/rscanner/src/execute/arp.rs index 525ba4d..dc71fd4 100644 --- a/rscanner/src/execute/arp.rs +++ b/rscanner/src/execute/arp.rs @@ -2,12 +2,12 @@ use std::net::{IpAddr, Ipv4Addr}; use std::time::Duration; use pnet::datalink::{ - Channel, Config, DataLinkReceiver, DataLinkSender, MacAddr, NetworkInterface + Channel, Config, DataLinkReceiver, DataLinkSender, MacAddr, NetworkInterface, }; use pnet::packet::{ arp::{ArpHardwareTypes, ArpOperations, ArpPacket, MutableArpPacket}, ethernet::{EtherTypes, MutableEthernetPacket}, - MutablePacket, Packet + MutablePacket, Packet, }; use snafu::OptionExt; @@ -30,16 +30,15 @@ async fn send_arp_packets( continue; } let mut ethernet_buffer = [0u8; 42]; - let mut ethernet_packet = MutableEthernetPacket::new( - &mut ethernet_buffer).context(OptionEmptySnafu)?; + let mut ethernet_packet = + MutableEthernetPacket::new(&mut ethernet_buffer).context(OptionEmptySnafu)?; ethernet_packet.set_destination(MacAddr::broadcast()); ethernet_packet.set_source(mac); ethernet_packet.set_ethertype(EtherTypes::Arp); let mut arp_buffer = [0u8; 28]; - let mut arp_packet = MutableArpPacket::new( - &mut arp_buffer).context(OptionEmptySnafu)?; + let mut arp_packet = MutableArpPacket::new(&mut arp_buffer).context(OptionEmptySnafu)?; arp_packet.set_hardware_type(ArpHardwareTypes::Ethernet); arp_packet.set_protocol_type(EtherTypes::Ipv4); diff --git a/rscanner/src/execute/common.rs b/rscanner/src/execute/common.rs index b518fd2..0ff690c 100644 --- a/rscanner/src/execute/common.rs +++ b/rscanner/src/execute/common.rs @@ -22,11 +22,11 @@ pub trait SocketScanner { Self::pre_scan(&scan_opts).await?; for socket_addr in scan_opts.iter_sockets()? { - let per_timeout = scan_opts.per_timeout; + let retry_timeout = scan_opts.retry_interval; Self::pre_send_socket(&socket_addr) .await .unwrap_or_else(|e| tracing::error!("pre send socket error with {e:?}")); - tokio::spawn(async move { Self::socket_success(socket_addr, per_timeout).await }); + tokio::spawn(async move { Self::socket_success(socket_addr, retry_timeout).await }); } Self::after_scan().await?; diff --git a/rscanner/src/execute/icmp/interface.rs b/rscanner/src/execute/icmp/interface.rs index 85d5378..2cf1ff1 100644 --- a/rscanner/src/execute/icmp/interface.rs +++ b/rscanner/src/execute/icmp/interface.rs @@ -1,5 +1,6 @@ use std::net::Ipv4Addr; +use crate::err::OptionEmptySnafu; use pnet::datalink::{Channel, MacAddr}; use pnet::packet::ethernet::EtherTypes; use pnet::packet::ethernet::MutableEthernetPacket; @@ -9,12 +10,11 @@ use pnet::packet::{ }; use pnet::packet::{MutablePacket, Packet}; use snafu::OptionExt; -use crate::err::OptionEmptySnafu; +use super::common; +use crate::err::Result; use crate::interfaces; use crate::interfaces::interface_normal_running; -use crate::err::Result; -use super::common; pub fn send_with_interface(target_ip: Ipv4Addr) -> Result<()> { tracing::debug!("{target_ip} send by specific interface"); @@ -24,7 +24,8 @@ pub fn send_with_interface(target_ip: Ipv4Addr) -> Result<()> { } if let Some(source_ip) = interfaces::get_interface_ipv4(&interface) { let mut header = [0u8; common::ICMP_LEN]; - let mut icmp_packet = MutableEchoRequestPacket::new(&mut header).context(OptionEmptySnafu)?; + let mut icmp_packet = + MutableEchoRequestPacket::new(&mut header).context(OptionEmptySnafu)?; common::modify_icmp_packet(&mut icmp_packet); let mut ipv4_header = [0u8; common::IPV4_LEN]; diff --git a/rscanner/src/execute/icmp/send.rs b/rscanner/src/execute/icmp/send.rs index 6f84124..1c3f30e 100644 --- a/rscanner/src/execute/icmp/send.rs +++ b/rscanner/src/execute/icmp/send.rs @@ -5,11 +5,12 @@ use std::time::Duration; use pnet::packet::icmp::echo_request::MutableEchoRequestPacket; use tokio::time::MissedTickBehavior; +use crate::monitor; +use crate::setting::command::ScanOpts; + use super::common; use super::interface; use super::receive; -use crate::monitor; -use crate::setting::command::ScanOpts; static R: AtomicU64 = AtomicU64::new(0); @@ -36,11 +37,7 @@ pub async fn scan(scan_opts: ScanOpts) -> anyhow::Result<()> { Err(_) => { tracing::info!("receive packets thread over because timeout"); let send_count = R.load(Ordering::Relaxed); - let total_received = monitor::receive_packets_handle() - .await - .lock() - .unwrap() - .len(); + let total_received = monitor::receive_packets_handle().await.lock().await.len(); println!("send {send_count} ips, receive packets from {total_received} ips"); } Ok(e) => { @@ -62,7 +59,7 @@ pub async fn icmp_ips_chunks(hosts: Vec) -> anyhow::Result<()> { continue; } tx.send_to(icmp_packet, target).unwrap_or_else(|_| { - interface::send_with_interface(host); + interface::send_with_interface(host).unwrap(); 0 }); // tx.send_to(icmp_packet, target).unwrap(); diff --git a/rscanner/src/execute/udp.rs b/rscanner/src/execute/udp.rs index 8d445e6..1d4f613 100644 --- a/rscanner/src/execute/udp.rs +++ b/rscanner/src/execute/udp.rs @@ -5,20 +5,15 @@ use std::time::Duration; use async_trait::async_trait; use hashbrown::{HashMap, HashSet}; use pnet::packet::{ - icmp::{ - destination_unreachable::IcmpCodes, - echo_reply::EchoReplyPacket - }, + icmp::{destination_unreachable::IcmpCodes, echo_reply::EchoReplyPacket}, ipv4::Ipv4Packet, Packet, - udp::UdpPacket + udp::UdpPacket, }; use tokio::{ net::UdpSocket, - sync::{ - Mutex, OnceCell - }, - time::Instant + sync::{Mutex, OnceCell}, + time::Instant, }; use crate::setting::command::ScanOpts; @@ -45,12 +40,12 @@ async fn get_udp_unreachable_addrs() -> &'static Arc>> { } async fn add_udp_unreachable_addrs(addr: IpAddr) { - let mut handle = get_udp_unreachable_addrs().await.lock().unwrap(); + let mut handle = get_udp_unreachable_addrs().await.lock().await; handle.insert(addr); } async fn ip_udp_send_interval_millis(ip: IpAddr) -> u128 { - let udp_last_send = get_upd_last_ip_send().await.lock().unwrap(); + let udp_last_send = get_upd_last_ip_send().await.lock().await; let last_send = udp_last_send.get(&ip); if let Some(last_send) = last_send { let elapsed = last_send.elapsed().as_millis(); @@ -113,17 +108,17 @@ impl SocketScanner for UdpSocketScanner { let ip = socket.ip(); let sleep_millis = ip_udp_send_interval_millis(ip).await; tokio::time::sleep(Duration::from_millis(sleep_millis as u64)).await; - let mut last_send = get_upd_last_ip_send().await.lock().unwrap(); + let mut last_send = get_upd_last_ip_send().await.lock().await; last_send.insert(ip, Instant::now()); Ok(()) } async fn after_scan() -> anyhow::Result<()> { let udp_unreachable_addrs = { - let handle = get_udp_unreachable_addrs().await.lock().unwrap(); + let handle = get_udp_unreachable_addrs().await.lock().await; (*handle).clone() }; - let socket_manager = get_socket_manager().await.lock().unwrap(); + let socket_manager = get_socket_manager().await.lock().await; for socket in &*socket_manager { if udp_unreachable_addrs.contains(&socket.ip()) { println!("rscan|udp|{socket}|"); diff --git a/rscanner/src/interfaces.rs b/rscanner/src/interfaces.rs index 58b26bf..1aedab0 100644 --- a/rscanner/src/interfaces.rs +++ b/rscanner/src/interfaces.rs @@ -17,11 +17,10 @@ pub fn get_interface_ipv4(interface: &NetworkInterface) -> Option { } pub fn get_interface_by_name(interface_name: &str) -> Result { - Ok(pnet::datalink::interfaces() + pnet::datalink::interfaces() .into_iter() .find(|interface| interface.name == interface_name) - .context(OptionEmptySnafu)? - } + .context(OptionEmptySnafu) } pub fn interface_normal_running(interface: &NetworkInterface) -> bool { diff --git a/rscanner/src/lib.rs b/rscanner/src/lib.rs index e35b892..914469b 100644 --- a/rscanner/src/lib.rs +++ b/rscanner/src/lib.rs @@ -1,7 +1,7 @@ +pub mod consts; pub mod err; pub mod execute; pub mod interfaces; pub mod monitor; pub mod performance; pub mod setting; -pub mod consts; diff --git a/rscanner/src/monitor.rs b/rscanner/src/monitor.rs index e8bc2df..d8a5828 100644 --- a/rscanner/src/monitor.rs +++ b/rscanner/src/monitor.rs @@ -2,10 +2,7 @@ use std::net::IpAddr; use std::sync::Arc; use hashbrown::HashSet; -use tokio::sync::{ - Mutex, - OnceCell -}; +use tokio::sync::{Mutex, OnceCell}; static RECEIVE_PACKETS: OnceCell>>> = OnceCell::const_new(); @@ -16,11 +13,7 @@ pub async fn receive_packets_handle() -> &'static Arc>> { } pub async fn is_addr_received(addr: &IpAddr) -> bool { - receive_packets_handle() - .await - .lock() - .await - .contains(addr) + receive_packets_handle().await.lock().await.contains(addr) } pub async fn add_receive_ipaddr(addr: IpAddr) { diff --git a/rscanner/src/setting/parse.rs b/rscanner/src/setting/parse.rs index 1d6a52d..69207c6 100644 --- a/rscanner/src/setting/parse.rs +++ b/rscanner/src/setting/parse.rs @@ -3,8 +3,9 @@ use std::str::FromStr; use std::sync::Arc; use ipnetwork::Ipv4Network; +use snafu::{IntoError, NoneError, ResultExt}; -use crate::err::{APPError, PortFormatSnafu}; +use crate::err::{PortFormatSnafu, PortParseSnafu, Result}; /// parse input like /// single ip `192.168.1.1` @@ -37,18 +38,20 @@ pub fn parse_ports(input: &str) -> anyhow::Result>> { Ok(Arc::new(ports)) } -pub fn parse_ports_range(input: &str) -> anyhow::Result> { - let range = input - .split('-') - .map(str::parse) - .collect::, std::num::ParseIntError>>()?; +pub fn parse_ports_range(input: &str) -> Result> { + let range = input.split('-').collect::>(); if let [start, end] = range.as_slice() { - let result = (*start..*end).collect::>(); - return Ok(result); + let parsed_start = start.parse::().context(PortParseSnafu { + value: start.to_owned(), + })?; + let parsed_end = end.parse::().context(PortParseSnafu { + value: end.to_owned(), + })?; + let result = (parsed_start..parsed_end).collect::>(); + Ok(result) + } else { + Err(PortFormatSnafu { value: input }.into_error(NoneError)) } - PortFormatSnafu { - value: input - }.into() } #[cfg(test)] @@ -80,14 +83,14 @@ mod tests { for i in 0..4 { values.push(Ipv4Addr::from(i + min)); } - assert_eq!(result, values); + assert_eq!(result, Arc::new(values)); } #[test] fn ip_single() { assert_eq!( parse_hosts("192.168.1.1").unwrap(), - vec![Ipv4Addr::from_str("192.168.1.1").unwrap()] + Arc::new(vec![Ipv4Addr::from_str("192.168.1.1").unwrap()]) ); } } From d4ca62b4c4cc15d9eb5b6096e17d43373ec4a186 Mon Sep 17 00:00:00 2001 From: gongzhengyang Date: Tue, 12 Dec 2023 17:19:35 +0800 Subject: [PATCH 3/3] feat: improve scan ip speed --- .gitignore | 1 + Cargo.toml | 2 + rscanner/Cargo.toml | 2 +- rscanner/src/execute/icmp/common.rs | 2 +- rscanner/src/execute/icmp/interface.rs | 101 +++++++++++++------------ rscanner/src/execute/icmp/receive.rs | 22 +++--- rscanner/src/execute/icmp/send.rs | 28 +++---- rscanner/src/execute/udp.rs | 4 +- rscanner/src/monitor.rs | 35 ++++++++- rscanner/src/setting/command.rs | 4 + src/main.rs | 9 ++- 11 files changed, 127 insertions(+), 83 deletions(-) diff --git a/.gitignore b/.gitignore index 7e38809..af498ec 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,4 @@ .idea/ Cargo.lock tmp +temp.rs diff --git a/Cargo.toml b/Cargo.toml index dd36f4d..6505c3f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,8 @@ rscanner = { path = "rscanner" } tokio = { version = "1", features = ["full", "tracing"] } tracing = "0.1" tracing-subscriber = "0.3" +cached = "0.46" +chrono = "0.4" [workspace] members = [".", "rscanner"] diff --git a/rscanner/Cargo.toml b/rscanner/Cargo.toml index e4a5e6a..7102a5a 100644 --- a/rscanner/Cargo.toml +++ b/rscanner/Cargo.toml @@ -14,8 +14,8 @@ repository = "https://github.com/gongzhengyang/rscan" [dependencies] anyhow = "1.0.71" async-trait = "0.1" +cached = "0.46" clap = { version = "4.2", features = ["derive"] } -hashbrown = "0.14" ipnetwork = "0.20" itertools = "0.12" pnet = "0.34" diff --git a/rscanner/src/execute/icmp/common.rs b/rscanner/src/execute/icmp/common.rs index 5890a99..4c200d9 100644 --- a/rscanner/src/execute/icmp/common.rs +++ b/rscanner/src/execute/icmp/common.rs @@ -19,7 +19,7 @@ pub fn get_transport_channel() -> anyhow::Result<(TransportSender, TransportRece Ok(pnet_transport::transport_channel(4096, channel_type)?) } -pub fn modify_icmp_packet(icmp_packet: &mut MutableEchoRequestPacket) { +pub fn set_icmp_send_packet(icmp_packet: &mut MutableEchoRequestPacket) { icmp_packet.set_icmp_type(IcmpTypes::EchoRequest); icmp_packet.set_icmp_code(IcmpCodes::NoCode); icmp_packet.set_identifier(rand::random::()); diff --git a/rscanner/src/execute/icmp/interface.rs b/rscanner/src/execute/icmp/interface.rs index 2cf1ff1..65b5f86 100644 --- a/rscanner/src/execute/icmp/interface.rs +++ b/rscanner/src/execute/icmp/interface.rs @@ -1,65 +1,70 @@ use std::net::Ipv4Addr; -use crate::err::OptionEmptySnafu; -use pnet::datalink::{Channel, MacAddr}; -use pnet::packet::ethernet::EtherTypes; -use pnet::packet::ethernet::MutableEthernetPacket; +use cached::proc_macro::cached; +use pnet::datalink::{Channel, MacAddr, NetworkInterface}; use pnet::packet::{ - icmp::echo_request::MutableEchoRequestPacket, ip::IpNextHeaderProtocols, + ethernet::{EtherTypes, MutableEthernetPacket}, + icmp::echo_request::MutableEchoRequestPacket, + ip::IpNextHeaderProtocols, ipv4::MutableIpv4Packet, + MutablePacket, Packet, }; -use pnet::packet::{MutablePacket, Packet}; -use snafu::OptionExt; +use snafu::{OptionExt, ResultExt}; + +use crate::err::{CommonIoSnafu, OptionEmptySnafu, Result}; +use crate::interfaces::{self, interface_normal_running}; use super::common; -use crate::err::Result; -use crate::interfaces; -use crate::interfaces::interface_normal_running; + +#[cached(time = 60)] +pub fn running_interface_with_ip() -> Vec<(NetworkInterface, Ipv4Addr)> { + pnet::datalink::interfaces() + .into_iter() + .filter(interface_normal_running) + .filter_map(|interface| interfaces::get_interface_ipv4(&interface).map(|x| (interface, x))) + .collect() +} pub fn send_with_interface(target_ip: Ipv4Addr) -> Result<()> { tracing::debug!("{target_ip} send by specific interface"); - for interface in pnet::datalink::interfaces() { - if !interface_normal_running(&interface) { - continue; - } - if let Some(source_ip) = interfaces::get_interface_ipv4(&interface) { - let mut header = [0u8; common::ICMP_LEN]; - let mut icmp_packet = - MutableEchoRequestPacket::new(&mut header).context(OptionEmptySnafu)?; - common::modify_icmp_packet(&mut icmp_packet); + for (interface, source_ip) in running_interface_with_ip() { + let mut header = [0u8; common::ICMP_LEN]; + let mut icmp_packet = + MutableEchoRequestPacket::new(&mut header).context(OptionEmptySnafu)?; + common::set_icmp_send_packet(&mut icmp_packet); - let mut ipv4_header = [0u8; common::IPV4_LEN]; - let mut ipv4_packet = MutableIpv4Packet::new(&mut ipv4_header).unwrap(); - ipv4_packet.set_version(4); - ipv4_packet.set_header_length(5); - ipv4_packet.set_dscp(4); - ipv4_packet.set_ecn(1); - ipv4_packet.set_ttl(64); - ipv4_packet.set_next_level_protocol(IpNextHeaderProtocols::Icmp); - ipv4_packet.set_source(source_ip); - ipv4_packet.set_destination(target_ip); - ipv4_packet.set_total_length(common::IPV4_LEN as u16); - ipv4_packet.set_payload(icmp_packet.packet_mut()); - ipv4_packet.set_checksum(pnet::packet::ipv4::checksum(&ipv4_packet.to_immutable())); + let mut ipv4_header = [0u8; common::IPV4_LEN]; + let mut ipv4_packet = MutableIpv4Packet::new(&mut ipv4_header).context(OptionEmptySnafu)?; + ipv4_packet.set_version(4); + ipv4_packet.set_header_length(5); + ipv4_packet.set_dscp(4); + ipv4_packet.set_ecn(1); + ipv4_packet.set_ttl(64); + ipv4_packet.set_next_level_protocol(IpNextHeaderProtocols::Icmp); + ipv4_packet.set_source(source_ip); + ipv4_packet.set_destination(target_ip); + ipv4_packet.set_total_length(common::IPV4_LEN as u16); + ipv4_packet.set_payload(icmp_packet.packet_mut()); + ipv4_packet.set_checksum(pnet::packet::ipv4::checksum(&ipv4_packet.to_immutable())); - let mut ethernet_buffer = [0u8; common::ETHERNET_LEN]; - let mut ethernet_packet = MutableEthernetPacket::new(&mut ethernet_buffer).unwrap(); + let mut ethernet_buffer = [0u8; common::ETHERNET_LEN]; + let mut ethernet_packet = + MutableEthernetPacket::new(&mut ethernet_buffer).context(OptionEmptySnafu)?; - ethernet_packet.set_destination(MacAddr::broadcast()); - ethernet_packet.set_source(interface.mac.unwrap()); - ethernet_packet.set_ethertype(EtherTypes::Ipv4); - ethernet_packet.set_payload(ipv4_packet.packet_mut()); + ethernet_packet.set_destination(MacAddr::broadcast()); + ethernet_packet.set_source(interface.mac.context(OptionEmptySnafu)?); + ethernet_packet.set_ethertype(EtherTypes::Ipv4); + ethernet_packet.set_payload(ipv4_packet.packet_mut()); - let (mut sender, _) = match pnet::datalink::channel(&interface, Default::default()) { - Ok(Channel::Ethernet(tx, rx)) => (tx, rx), - Ok(_) => panic!("Unknown channel type"), - Err(e) => panic!("Error happened {}", e), - }; - sender - .send_to(ethernet_packet.packet(), Some(interface)) - .unwrap() - .unwrap(); - } + let (mut sender, _) = match pnet::datalink::channel(&interface, Default::default()) { + Ok(Channel::Ethernet(tx, rx)) => (tx, rx), + Ok(_) => panic!("Unknown channel type"), + Err(e) => panic!("Error happened {}", e), + }; + sender + .send_to(ethernet_packet.packet(), Some(interface)) + .context(OptionEmptySnafu)? + .context(CommonIoSnafu)?; } Ok(()) } diff --git a/rscanner/src/execute/icmp/receive.rs b/rscanner/src/execute/icmp/receive.rs index df5e951..419a015 100644 --- a/rscanner/src/execute/icmp/receive.rs +++ b/rscanner/src/execute/icmp/receive.rs @@ -1,24 +1,26 @@ use std::time::Duration; -use pnet::packet::icmp::{echo_reply::EchoReplyPacket, IcmpTypes}; -use pnet::packet::Packet; +use pnet::packet::{ + icmp::{echo_reply::EchoReplyPacket, IcmpTypes}, + Packet, +}; -use super::common; use crate::monitor; +use super::common; + pub async fn receive_packets() -> anyhow::Result<()> { let (_, mut rx) = common::get_transport_channel()?; let mut iter = pnet_transport::icmp_packet_iter(&mut rx); loop { - tokio::time::sleep(Duration::from_micros(10)).await; - if let Ok(Some((packet, addr))) = iter.next_with_timeout(Duration::from_millis(1)) { - if monitor::is_addr_received(&addr).await { - continue; - } + tokio::time::sleep(Duration::from_millis(1)).await; + while let Ok(Some((packet, addr))) = iter.next_with_timeout(Duration::from_millis(1)) { if let Some(reply_packet) = EchoReplyPacket::new(packet.packet()) { if reply_packet.get_icmp_type() == IcmpTypes::EchoReply { - println!("rscan|icmp|{addr}|"); - monitor::add_receive_ipaddr(addr).await; + tracing::debug!("receive {addr}"); + if monitor::add_receive_ipaddr(addr).await { + println!("rscan|icmp|{addr}|"); + } } } } diff --git a/rscanner/src/execute/icmp/send.rs b/rscanner/src/execute/icmp/send.rs index 1c3f30e..cb4c62a 100644 --- a/rscanner/src/execute/icmp/send.rs +++ b/rscanner/src/execute/icmp/send.rs @@ -12,7 +12,7 @@ use super::common; use super::interface; use super::receive; -static R: AtomicU64 = AtomicU64::new(0); +static SEND: AtomicU64 = AtomicU64::new(0); pub async fn scan(scan_opts: ScanOpts) -> anyhow::Result<()> { let mut interval = tokio::time::interval(Duration::from_secs(scan_opts.retry_interval)); @@ -31,12 +31,10 @@ pub async fn scan(scan_opts: ScanOpts) -> anyhow::Result<()> { interval.tick().await; } }); - let result = - tokio::time::timeout(Duration::from_secs(timeout), receive::receive_packets()).await; - match result { + match tokio::time::timeout(Duration::from_secs(timeout), receive::receive_packets()).await { Err(_) => { tracing::info!("receive packets thread over because timeout"); - let send_count = R.load(Ordering::Relaxed); + let send_count = SEND.load(Ordering::Relaxed); let total_received = monitor::receive_packets_handle().await.lock().await.len(); println!("send {send_count} ips, receive packets from {total_received} ips"); } @@ -50,21 +48,19 @@ pub async fn scan(scan_opts: ScanOpts) -> anyhow::Result<()> { pub async fn icmp_ips_chunks(hosts: Vec) -> anyhow::Result<()> { let (mut tx, _) = common::get_transport_channel()?; for host in hosts { - let mut header = [0u8; common::ICMP_LEN]; - let mut icmp_packet = MutableEchoRequestPacket::new(&mut header).unwrap(); - common::modify_icmp_packet(&mut icmp_packet); - tracing::debug!("build icmp success for {host}"); let target = IpAddr::from(host); if monitor::is_addr_received(&target).await { + tracing::debug!("skip target {target} because received"); continue; } - tx.send_to(icmp_packet, target).unwrap_or_else(|_| { - interface::send_with_interface(host).unwrap(); - 0 - }); - // tx.send_to(icmp_packet, target).unwrap(); - R.fetch_add(1, Ordering::Relaxed); - tracing::debug!("sending packets count {}", R.load(Ordering::Relaxed)); + let mut header = [0u8; common::ICMP_LEN]; + let mut icmp_packet = MutableEchoRequestPacket::new(&mut header).unwrap(); + common::set_icmp_send_packet(&mut icmp_packet); + tracing::debug!("build icmp success for {host}"); + if tx.send_to(icmp_packet, target).is_err() { + interface::send_with_interface(host).unwrap_or_default(); + } + SEND.fetch_add(1, Ordering::Relaxed); } Ok(()) } diff --git a/rscanner/src/execute/udp.rs b/rscanner/src/execute/udp.rs index 1d4f613..c68ff6d 100644 --- a/rscanner/src/execute/udp.rs +++ b/rscanner/src/execute/udp.rs @@ -1,14 +1,14 @@ +use std::collections::{HashMap, HashSet}; use std::net::{IpAddr, SocketAddr}; use std::sync::Arc; use std::time::Duration; use async_trait::async_trait; -use hashbrown::{HashMap, HashSet}; use pnet::packet::{ icmp::{destination_unreachable::IcmpCodes, echo_reply::EchoReplyPacket}, ipv4::Ipv4Packet, - Packet, udp::UdpPacket, + Packet, }; use tokio::{ net::UdpSocket, diff --git a/rscanner/src/monitor.rs b/rscanner/src/monitor.rs index d8a5828..1872d2c 100644 --- a/rscanner/src/monitor.rs +++ b/rscanner/src/monitor.rs @@ -1,9 +1,14 @@ +use std::collections::HashSet; use std::net::IpAddr; +use std::path::Path; use std::sync::Arc; -use hashbrown::HashSet; +use itertools::Itertools; +use snafu::ResultExt; use tokio::sync::{Mutex, OnceCell}; +use crate::err::{CommonIoSnafu, Result}; + static RECEIVE_PACKETS: OnceCell>>> = OnceCell::const_new(); pub async fn receive_packets_handle() -> &'static Arc>> { @@ -16,7 +21,31 @@ pub async fn is_addr_received(addr: &IpAddr) -> bool { receive_packets_handle().await.lock().await.contains(addr) } -pub async fn add_receive_ipaddr(addr: IpAddr) { +pub async fn add_receive_ipaddr(addr: IpAddr) -> bool { let mut receive_handle = receive_packets_handle().await.lock().await; - receive_handle.insert(addr); + receive_handle.insert(addr) +} + +pub async fn save_receive_addrs(path: &str) -> Result<()> { + let cache = receive_packets_handle().await; + let addrs = cache + .lock() + .await + .iter() + .map(|x| x.to_string()) + .collect_vec(); + let path = Path::new(path); + if let Some(parent) = path.parent() { + if !path.exists() { + tokio::fs::create_dir_all(parent) + .await + .context(CommonIoSnafu)?; + } + } + let results = addrs.join("\n"); + tracing::debug!("receive data :{results}"); + tokio::fs::write(path, results) + .await + .context(CommonIoSnafu)?; + Ok(()) } diff --git a/rscanner/src/setting/command.rs b/rscanner/src/setting/command.rs index a7a8906..caf8e56 100644 --- a/rscanner/src/setting/command.rs +++ b/rscanner/src/setting/command.rs @@ -59,6 +59,10 @@ pub struct ScanOpts { /// The seconds retry interval when retry is set bigger than 0 #[arg(long, default_value_t = 3)] pub retry_interval: u64, + + /// save ipaddr results into file at `tmp/2023-12-12--16:25:55.txt` if the args is None + #[arg(long)] + pub filepath: Option, } impl ScanOpts { diff --git a/src/main.rs b/src/main.rs index 878dedf..04995f0 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,18 +4,22 @@ use clap::Parser; use rscanner::execute; use rscanner::execute::common::SocketScanner; -use rscanner::setting::command::{ScanType, ScanOpts}; +use rscanner::setting::command::{ScanOpts, ScanType}; #[tokio::main] async fn main() { // used for tokio task manage - console_subscriber::init(); + // console_subscriber::init(); + tracing_subscriber::fmt::init(); let scan_opts = ScanOpts::parse(); #[cfg(unix)] rscanner::performance::improve_limits().unwrap(); let timeout = scan_opts.timeout; tracing::info!("waiting for {} seconds", timeout); + let result_save_path = scan_opts.filepath.clone().unwrap_or_else(|| { + format!("tmp/{}.txt", chrono::Local::now().format("%Y-%m-%d--%H:%M:%S")) + }); match scan_opts.execute { ScanType::Icmp => { tracing::info!("execute icmp scan"); @@ -65,4 +69,5 @@ async fn main() { } } tokio::time::sleep(Duration::from_secs(timeout + 1)).await; + rscanner::monitor::save_receive_addrs(&result_save_path).await.unwrap(); }