From 1e82fd381c3c8a530700b1d5c4293516da6bb139 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 8 Jan 2024 13:52:39 +0900 Subject: [PATCH] feat: migrate capture device handling to `capture.rs` --- src/main.rs | 16 +++++++++------- src/net/live_fluereflow.rs | 19 +++++++------------ src/net/mod.rs | 5 +++++ src/net/online_fluereflow.rs | 25 +++++++++++++------------ 4 files changed, 34 insertions(+), 31 deletions(-) diff --git a/src/main.rs b/src/main.rs index 81e4afe..79ef476 100644 --- a/src/main.rs +++ b/src/main.rs @@ -11,6 +11,8 @@ pub mod utils; use clap::{Arg, ArgAction, Command}; use pnet::datalink; +use crate::net::list_devices; + use std::process::exit; // This function sets up the command line interface for the application using the clap library. @@ -266,11 +268,11 @@ async fn main() { println!("Online mode"); utils::get_local_ip(); if args.get_flag("list") { - println!("List of interfaces"); - for iface in interfaces { - println!("[{}]: {}", iface.index, iface.name); + let interfaces = list_devices().unwrap(); + println!("Found {} devices", interfaces.len()); + for (i, interface) in interfaces.iter().enumerate() { + println!("[{}]: {}", i, interface.name); } - exit(0); } let use_mac = args.get_flag("useMACaddress"); @@ -332,11 +334,11 @@ async fn main() { Some(("live", args)) => { println!("Live mode"); if args.get_flag("list") { - println!("List of interfaces"); + let interfaces = list_devices().unwrap(); + println!("Found {} devices", interfaces.len()); for (i, interface) in interfaces.iter().enumerate() { println!("[{}]: {}", i, interface.name); - } - + } exit(0); } let use_mac = args.get_flag("useMACaddress"); diff --git a/src/net/live_fluereflow.rs b/src/net/live_fluereflow.rs index 8e70222..14e89b3 100644 --- a/src/net/live_fluereflow.rs +++ b/src/net/live_fluereflow.rs @@ -13,7 +13,6 @@ use fluere_config::Config; use fluere_plugin::PluginManager; use fluereflow::FluereRecord; -use pcap::Capture; use ratatui::{ backend::CrosstermBackend, layout::{Constraint, Direction, Layout}, @@ -25,13 +24,13 @@ use tokio::sync::Mutex; use tokio::task; use tokio::time::sleep; -use super::interface::get_interface; - use crate::{ net::{ flows::update_flow, parser::{microseconds_to_timestamp, parse_fluereflow, parse_keys, parse_microseconds}, types::TcpFlags, + CaptureDevice, + find_device, }, types::{Args, UDFlowKey}, utils::{cur_time_file, fluere_exporter}, @@ -84,15 +83,10 @@ pub async fn online_packet_capture(arg: Args) { .load_plugins(&config) .await .expect("Failed to load plugins"); - let interface = get_interface(interface_name.as_str()); - let mut cap = Capture::from_device(interface) - .unwrap() - .promisc(true) - .timeout(60000) - //.buffer_size(100000000) - .immediate_mode(true) - .open() - .unwrap(); + + let interface = find_device(interface_name.as_str()).unwrap(); + let cap_device = CaptureDevice::new(interface.clone()).unwrap(); + let mut cap = cap_device.capture; let file_dir = "./output"; match fs::create_dir_all(<&str>::clone(&file_dir)) { @@ -178,6 +172,7 @@ pub async fn online_packet_capture(arg: Args) { }); tokio::spawn(listen_for_exit_keys()); + loop { match cap.next_packet() { Err(_) => { diff --git a/src/net/mod.rs b/src/net/mod.rs index 0e62b07..ae0e099 100644 --- a/src/net/mod.rs +++ b/src/net/mod.rs @@ -2,6 +2,7 @@ pub mod errors; //mod fluereflow; mod flows; mod interface; +mod capture; pub mod live_fluereflow; mod offline_fluereflows; pub mod online_fluereflow; @@ -14,4 +15,8 @@ pub use interface::list_interface_names; pub use interface::list_interfaces; pub use offline_fluereflows::fluereflow_fileparse; pub use packet_pcap::pcap_capture; +pub use capture::CaptureDevice; +pub use capture::CaptureError; +pub use capture::find_device; +pub use capture::list_devices; //pub use types::FluereRecord; diff --git a/src/net/online_fluereflow.rs b/src/net/online_fluereflow.rs index 4f0676c..1730843 100644 --- a/src/net/online_fluereflow.rs +++ b/src/net/online_fluereflow.rs @@ -3,8 +3,6 @@ // The data is then exported to a CSV file. extern crate csv; -use pcap::Capture; - use fluere_config::Config; use fluere_plugin::PluginManager; use fluereflow::FluereRecord; @@ -12,13 +10,13 @@ use fluereflow::FluereRecord; use tokio::task; use tokio::time::sleep; -use super::interface::get_interface; - use crate::{ net::{ flows::update_flow, parser::{parse_fluereflow, parse_keys, parse_microseconds}, types::{Key, TcpFlags}, + CaptureDevice, + find_device, }, types::{Args, UDFlowKey}, utils::{cur_time_file, fluere_exporter}, @@ -51,14 +49,17 @@ pub async fn packet_capture(arg: Args) { .await .expect("Failed to load plugins"); - let interface = get_interface(interface_name.as_str()); - let mut cap = Capture::from_device(interface) - .unwrap() - .promisc(true) - //.buffer_size(100000000) - .immediate_mode(true) - .open() - .unwrap(); + let interface = find_device(interface_name.as_str()).unwrap(); + let cap_device = CaptureDevice::new(interface.clone()).unwrap(); + let mut cap = cap_device.capture; + // let mut cp_device + // let mut cap = Capture::from_device(interface) + // .unwrap() + // .promisc(true) + // //.buffer_size(100000000) + // .immediate_mode(true) + // .open() + // .unwrap(); let file_dir = "./output"; match fs::create_dir_all(<&str>::clone(&file_dir)) {