net_agent/
agent.rs

1use std::sync::atomic::{AtomicBool, Ordering};
2use std::sync::Arc;
3
4use pcap::Capture;
5
6use crate::config::Config;
7use crate::core::poller::Poller;
8use crate::packet_handler::PacketHandler;
9
10#[derive(Clone)]
11pub struct Agent {
12    config: Config,
13    running: Arc<AtomicBool>,
14}
15
16impl Agent {
17    pub fn new(config: Config, running: Arc<AtomicBool>) -> Self {
18        Self { 
19            config,
20            running,
21        }
22    }
23
24    pub fn run(&self) -> u64 {
25        log::info!("Starting capturing on {}", self.config.get_device_name());
26        let capture = Capture::from_device(self.config.get_device_name());
27        if capture.is_err() {
28            log::error!("Couldn't open a capture handle for a device: {}\nProvide a valid device", capture.err().unwrap());
29            return 0;
30        };
31        let capture = capture.unwrap();
32        let capture = capture.buffer_size(self.config.get_buffer_size()).open();
33        if capture.is_err() {
34            log::error!("Couldn't activates an inactive capture: {}", capture.err().unwrap());
35            return 0;
36        };
37
38        let poller = Poller::builder()
39            .with_capture(capture.unwrap())
40            .with_handler(PacketHandler {
41                directory: self.config.get_output_directory().to_string(),
42            })
43            .with_running(self.running.clone());
44        let poller = match self.config.get_number_packages() {
45            Some(packet_cnt) => poller.with_packet_cnt(packet_cnt),
46            _ => poller,
47        };
48        poller.build().poll()
49    }
50}
51
52impl Drop for Agent {
53    fn drop(&mut self) {
54        self.running.store(false, Ordering::SeqCst);
55    }
56}