pistol/
utils.rs

1use num_cpus;
2use pnet::datalink::MacAddr;
3use rand::Rng;
4use std::net::IpAddr;
5use std::net::Ipv4Addr;
6use std::net::Ipv6Addr;
7use std::time::Duration;
8use threadpool::ThreadPool;
9use tracing::warn;
10
11use crate::DEFAULT_TIMEOUT;
12use crate::SYSTEM_NET_CACHE;
13use crate::error::PistolError;
14
15const MAX_THREADS: usize = 1000;
16
17pub fn time_sec_to_string(cost: Duration) -> String {
18    if cost.as_secs_f64() > 1.0 {
19        format!("{:.3}s", cost.as_secs_f64())
20    } else {
21        format!("{:.3} ms", cost.as_secs_f64() * 1000.0)
22    }
23}
24
25pub fn num_threads_check(num_threads: usize) -> usize {
26    let mut num_threads = num_threads;
27    if num_threads > MAX_THREADS {
28        warn!(
29            "system try to create too many threads (current threads num: {}, fixed threads num: {}))",
30            num_threads, MAX_THREADS
31        );
32        num_threads = MAX_THREADS;
33    }
34    num_threads
35}
36
37pub fn neigh_cache_update(addr: IpAddr, mac: MacAddr) -> Result<(), PistolError> {
38    // release the lock when leaving the function
39    let mut snc = match SYSTEM_NET_CACHE.lock() {
40        Ok(snc) => snc,
41        Err(e) => {
42            return Err(PistolError::TryLockGlobalVarFailed {
43                var_name: String::from("SYSTEM_NET_CACHE"),
44                e: e.to_string(),
45            });
46        }
47    };
48    Ok(snc.update_neighbor_cache(addr, mac))
49}
50
51/// Returns the random port
52pub fn random_port() -> u16 {
53    let mut rng = rand::rng();
54    rng.random_range(10000..=65535)
55}
56
57/// Returns the random ipv4 addr
58pub fn random_ipv4_addr() -> Ipv4Addr {
59    let mut rng = rand::rng();
60    let x0 = rng.random_range(0..=255);
61    let x1 = rng.random_range(0..=255);
62    let x2 = rng.random_range(0..=255);
63    let x3 = rng.random_range(0..=255);
64    Ipv4Addr::new(x0, x1, x2, x3)
65}
66
67/// Returns the random ipv6 addr
68pub fn random_ipv6_addr() -> Ipv6Addr {
69    let mut rng = rand::rng();
70    let x0 = rng.random_range(0..=65535);
71    let x1 = rng.random_range(0..=65535);
72    let x2 = rng.random_range(0..=65535);
73    let x3 = rng.random_range(0..=65535);
74    let x4 = rng.random_range(0..=65535);
75    let x5 = rng.random_range(0..=65535);
76    let x6 = rng.random_range(0..=65535);
77    let x7 = rng.random_range(0..=65535);
78    Ipv6Addr::new(x0, x1, x2, x3, x4, x5, x6, x7)
79}
80
81/// Returns the random port with range
82pub fn random_port_range(start: u16, end: u16) -> u16 {
83    let mut rng = rand::rng();
84    rng.random_range(start..=end)
85}
86
87/// Returns the number of CPUs in the machine
88pub fn get_cpu_num() -> usize {
89    num_cpus::get()
90}
91
92pub fn get_threads_pool(num_threads: usize) -> ThreadPool {
93    let pool = if num_threads > 0 {
94        ThreadPool::new(num_threads)
95    } else {
96        let cpus = get_cpu_num();
97        ThreadPool::new(cpus)
98    };
99    pool
100}
101
102pub fn get_default_timeout() -> Duration {
103    Duration::from_secs_f64(DEFAULT_TIMEOUT)
104}
105
106pub struct PistolHex {
107    pub hex: String, // hex => dec
108}
109
110impl PistolHex {
111    pub fn new_hex(hex_str: &str) -> PistolHex {
112        let hex_str_len = hex_str.len();
113        let after_fix = if hex_str_len % 2 == 1 {
114            format!("0{}", hex_str)
115        } else {
116            hex_str.to_string()
117        };
118        PistolHex { hex: after_fix }
119    }
120    pub fn be_vec_to_u32(input: &[u8]) -> Result<u32, PistolError> {
121        if input.len() <= 4 {
122            let padding_size = 4 - input.len();
123            let mut after_padding = vec![0; padding_size];
124            for &i in input {
125                after_padding.push(i as u32);
126            }
127            let a = after_padding[0] << 24;
128            let b = after_padding[1] << 16;
129            let c = after_padding[2] << 8;
130            let d = after_padding[3];
131            Ok(a + b + c + d)
132        } else {
133            let v = format!("{:?}", input);
134            Err(PistolError::InputTooLoog { v })
135        }
136    }
137    pub fn decode_as_u32(&self) -> Result<u32, PistolError> {
138        match hex::decode(&self.hex) {
139            Ok(d) => Self::be_vec_to_u32(&d),
140            Err(e) => Err(e.into()),
141        }
142    }
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148    use pnet::datalink::interfaces;
149    #[test]
150    fn test_convert() {
151        let v: Vec<u8> = vec![1, 1, 1, 1];
152        let r = PistolHex::be_vec_to_u32(&v).unwrap();
153        assert_eq!(r, 16843009);
154
155        let s = "51E80C";
156        let h = PistolHex::new_hex(s);
157        let r = h.decode_as_u32().unwrap();
158        assert_eq!(r, 5367820);
159
160        let s = "1C";
161        let h = PistolHex::new_hex(s);
162        let r = h.decode_as_u32().unwrap();
163        assert_eq!(r, 28);
164
165        let s = "A";
166        let h = PistolHex::new_hex(s);
167        let r = h.decode_as_u32().unwrap();
168        assert_eq!(r, 10);
169    }
170    #[test]
171    fn test_get_cpus() {
172        let cpus = get_cpu_num();
173        println!("{}", cpus);
174    }
175    #[test]
176    fn interface_loopback() {
177        for interface in interfaces() {
178            if interface.is_loopback() {
179                println!("{} is loopback interface", interface);
180            }
181        }
182    }
183    #[test]
184    fn interface_list() {
185        for interface in interfaces() {
186            println!("list interface: {}, {:?}", interface.name, interface.ips);
187        }
188    }
189}
pistol/utils.rs

pistol/
utils.rs
