Merge pull request #5 from gongzhengyang/feature-global-improve

Feature global improve

.gitignore

@@ -2,3 +2,4 @@
 .idea/
 Cargo.lock
 tmp
+temp.rs

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"]

rscanner/Cargo.toml

@@ -14,15 +14,15 @@ 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.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"

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

rscanner/src/consts.rs

@@ -0,0 +1 @@
+

rscanner/src/err.rs

@@ -1,33 +1,37 @@
-use serde::Serialize;
-use thiserror::Error;
+use snafu::prelude::*;
+use snafu::Location;
+use std::num::ParseIntError;
 
-#[derive(Error, Debug, Serialize, PartialEq)]
+#[derive(Snafu, Debug)]
+#[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 value parse failed {}", source))]
+    PortParse {
+        location: Location,
+        source: ParseIntError,
+        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,
-}
 
-#[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(())
-    }
+    #[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<T> = std::result::Result<T, APPError>;

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<Ipv4Addr>,
 ) -> 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,15 @@ 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);

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?;

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::<u16>());

rscanner/src/execute/icmp/interface.rs

@@ -1,61 +1,70 @@
 use std::net::Ipv4Addr;
 
-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, ResultExt};
 
-use crate::interfaces;
-use crate::interfaces::interface_normal_running;
+use crate::err::{CommonIoSnafu, OptionEmptySnafu, Result};
+use crate::interfaces::{self, interface_normal_running};
 
 use super::common;
 
-pub fn send_with_interface(target_ip: Ipv4Addr) {
+#[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).unwrap();
-            common::modify_icmp_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 ethernet_buffer = [0u8; common::ETHERNET_LEN];
-            let mut ethernet_packet = MutableEthernetPacket::new(&mut ethernet_buffer).unwrap();
-
-            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());
-
-            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();
-        }
+    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).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).context(OptionEmptySnafu)?;
+
+        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))
+            .context(OptionEmptySnafu)?
+            .context(CommonIoSnafu)?;
     }
+    Ok(())
 }

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}|");
+                    }
                 }
             }
         }