From c6939c1fc52f388a3aea693329ef3b6d0c79f133 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 15:12:51 +0900 Subject: [PATCH] style: `rustfmt` --- src/cli.rs | 10 +++--- src/logger.rs | 10 ++++-- src/main.rs | 60 ++++++++++++++++---------------- src/net/capture.rs | 10 +++--- src/net/live_fluereflow.rs | 32 +++++++---------- src/net/offline_fluereflows.rs | 19 +++++----- src/net/online_fluereflow.rs | 36 +++++++++---------- src/net/packet_pcap.rs | 6 ++-- src/utils/fluere_csv_exporter.rs | 2 +- 9 files changed, 89 insertions(+), 96 deletions(-) diff --git a/src/cli.rs b/src/cli.rs index a683ad0..7d764f8 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -254,7 +254,7 @@ pub async fn handle_mode(mode: &str, args: &ArgMatches) -> (Args, u8) { exit(0); } - let arg_data = match mode { + let arg_data = match mode { "online" | "live" => parse_online_live_args(args, mode), "offline" => parse_offline_args(args), "pcap" => parse_pcap_args(args), @@ -326,10 +326,10 @@ fn parse_offline_args(args: &clap::ArgMatches) -> Args { .parse::() .unwrap(); // let verbose = args - // .get_one::("verbose") - // .unwrap() - // .parse::() - // .unwrap(); + // .get_one::("verbose") + // .unwrap() + // .parse::() + // .unwrap(); Args::new( None, diff --git a/src/logger.rs b/src/logger.rs index d8adf84..c133035 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -18,7 +18,12 @@ pub struct Logger { } impl Logger { - pub fn new(file_path: Option, severity: Option, write_to_std: Option, write_to_file: bool) -> Self { + pub fn new( + file_path: Option, + severity: Option, + write_to_std: Option, + write_to_file: bool, + ) -> Self { let mut path = file_path; if path.is_none() { path = Some(PathBuf::from( @@ -45,7 +50,6 @@ impl Logger { if path.as_ref().unwrap().parent().is_some() { std::fs::create_dir_all(path.as_ref().unwrap().parent().unwrap()).unwrap(); } - if write_to_file { file = Some(File::create(path.as_ref().unwrap()).unwrap()); @@ -73,7 +77,7 @@ impl Log for Logger { // Y M S, H:M:S Timezone let timestamp = Local::now().format("%Y-%m-%d %H:%M:%S %z").to_string(); let formatted_message = format!("[{}] [{}]: {}", timestamp, record.level(), record.args()); - + if self.write_to_std.as_ref().is_some() && record.level() <= self.severity { match self.write_to_std.as_ref().unwrap() { Logstdout::Stdout => { diff --git a/src/main.rs b/src/main.rs index 78d5fc6..c84fab5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,18 +10,22 @@ pub mod plugin; pub mod types; pub mod utils; -use std::{fmt::Display, process::exit}; use std::fs::File; +use std::{fmt::Display, process::exit}; use crate::logger::{Logger, Logstdout}; use crate::net::capture::DeviceError; // use env_logger;::{init, Logger}; -use log::{Level, LevelFilter, info, debug}; - +use log::{debug, info, Level, LevelFilter}; // FEAT:MAYBE: seprate `std` as feature flag for fluere and log crate -static LOGGER: Logger = Logger{write_to_file: false, file: None, write_to_std: Some(Logstdout::Stdout), severity: Level::Info}; +static LOGGER: Logger = Logger { + write_to_file: false, + file: None, + write_to_std: Some(Logstdout::Stdout), + severity: Level::Info, +}; #[derive(Debug)] enum FluereError { @@ -82,17 +86,16 @@ impl Display for Mode { } } - - fn from_verbose(level: u8) -> LevelFilter { - match level { - 0 => LevelFilter::Error, - 1 => LevelFilter::Warn, - 2 => LevelFilter::Info, - 3 => LevelFilter::Debug, - 4 => LevelFilter::Trace, - _ => unreachable!(), - } +fn from_verbose(level: u8) -> LevelFilter { + match level { + 0 => LevelFilter::Error, + 1 => LevelFilter::Warn, + 2 => LevelFilter::Info, + 3 => LevelFilter::Debug, + 4 => LevelFilter::Trace, + _ => unreachable!(), } +} struct Fluere { interface: String, @@ -128,34 +131,31 @@ impl Fluere { #[tokio::main] async fn main() { let args = cli::cli_template().get_matches(); - - // let mode = match args.subcommand() { - // Some((mode, _sub_args)) => mode, - // None => { - // log::error!("No mode selected. Use --help for more information."); - // exit(1); - // } + + // let mode = match args.subcommand() { + // Some((mode, _sub_args)) => mode, + // None => { + // log::error!("No mode selected. Use --help for more information."); + // exit(1); + // } // }; - if let Some((mode, sub_args)) = args.subcommand() { let mode_type: Mode = Mode::from(mode); debug!("Mode: {}", mode_type); let parems = cli::handle_mode(mode, sub_args).await; let _log_stdout = Logstdout::Stdout; - let _log_file :Option = None; + let _log_file: Option = None; let _log_level = Level::Info; - let logger = Logger::new(None,Some(Level::Trace), Some(Logstdout::Stdout),false); - + let logger = Logger::new(None, Some(Level::Trace), Some(Logstdout::Stdout), false); // (Args, u8) let filter = from_verbose(parems.1); - let _ = log::set_boxed_logger(Box::new(logger)) - .map(|()| log::set_max_level(filter)); - - debug!("Fluere started"); - + let _ = log::set_boxed_logger(Box::new(logger)).map(|()| log::set_max_level(filter)); + + debug!("Fluere started"); + match mode_type { Mode::Online => net::online_fluereflow::packet_capture(parems.0).await, Mode::Offline => net::fluereflow_fileparse(parems.0).await, diff --git a/src/net/capture.rs b/src/net/capture.rs index ed168ce..5394ac6 100644 --- a/src/net/capture.rs +++ b/src/net/capture.rs @@ -1,9 +1,9 @@ -use std::{fmt, time::Instant, borrow::Cow}; +use std::{borrow::Cow, fmt, time::Instant}; use crate::net::NetError; +use log::{debug, info}; use pcap::{Active, Address, Capture, Device, Error as PcapError}; -use log::{info, debug}; #[derive(Debug)] pub enum DeviceError { @@ -30,7 +30,7 @@ pub struct CaptureDevice { impl CaptureDevice { pub fn new(device: Device) -> Result { let capture = initialize_capture(device.clone())?; - let name: Cow<'static, str> = Cow::Owned(device.name); + let name: Cow<'static, str> = Cow::Owned(device.name); let desc: Cow<'static, str> = Cow::Owned(device.desc.unwrap_or("".to_string())); Ok(CaptureDevice { @@ -82,10 +82,10 @@ pub fn find_device(identifier: &str) -> Result { fn initialize_capture(device: Device) -> Result, PcapError> { info!("Opening capture session for device {}", device.name); - Ok(Capture::from_device(device)? + Capture::from_device(device)? .promisc(true) .snaplen(1024) .timeout(60000) .immediate_mode(true) - .open()?) + .open() } diff --git a/src/net/live_fluereflow.rs b/src/net/live_fluereflow.rs index 8ae6692..a125b19 100644 --- a/src/net/live_fluereflow.rs +++ b/src/net/live_fluereflow.rs @@ -1,13 +1,6 @@ // This file contains the implementation of the live packet capture functionality. // It uses the pcap library to capture packets from a network interface and the fluereflow library to convert the packets into NetFlow data. // The data is then displayed in a terminal user interface using the ratatui library. -use std::{ - collections::HashMap, - fs, - io, - sync::Arc, - time::{Duration, Instant, SystemTime}, -}; use crate::{ net::{ find_device, @@ -19,6 +12,12 @@ use crate::{ types::{Args, UDFlowKey}, utils::{cur_time_file, fluere_exporter}, }; +use std::{ + collections::HashMap, + fs, io, + sync::Arc, + time::{Duration, Instant, SystemTime}, +}; use fluere_config::Config; use fluere_plugin::PluginManager; @@ -29,6 +28,7 @@ use crossterm::{ execute, terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen}, }; +use log::{debug, trace}; use ratatui::{ backend::CrosstermBackend, layout::{Constraint, Direction, Layout}, @@ -36,13 +36,7 @@ use ratatui::{ widgets::{Block, Borders, Gauge, List, ListItem, Paragraph}, Frame, Terminal, }; -use tokio::{ - sync::Mutex, - task, - time::sleep, -}; -use log::{info, debug, trace}; - +use tokio::{sync::Mutex, task}; const MAX_RECENT_FLOWS: usize = 50; @@ -75,7 +69,7 @@ pub async fn online_packet_capture(arg: Args) { let duration = arg.parameters.duration.unwrap(); let interval = arg.parameters.interval.unwrap(); let flow_timeout = arg.parameters.timeout.unwrap(); - let sleep_windows = arg.parameters.sleep_windows.unwrap(); + let _sleep_windows = arg.parameters.sleep_windows.unwrap(); let config = Config::new(); let plugin_manager = PluginManager::new().expect("Failed to create plugin manager"); let plugin_worker = plugin_manager.start_worker(); @@ -261,10 +255,10 @@ pub async fn online_packet_capture(arg: Args) { }; update_flow(flow, is_reverse, update_key); - trace!( - "{} flow updated", - if is_reverse { "reverse" } else { "forward" } - ); + trace!( + "{} flow updated", + if is_reverse { "reverse" } else { "forward" } + ); if flags.fin == 1 || flags.rst == 1 { trace!("flow finished"); diff --git a/src/net/offline_fluereflows.rs b/src/net/offline_fluereflows.rs index ea1a642..e200b5c 100644 --- a/src/net/offline_fluereflows.rs +++ b/src/net/offline_fluereflows.rs @@ -11,9 +11,9 @@ use crate::{ }; use fluereflow::FluereRecord; +use log::{debug, info, trace}; use pcap::Capture; use tokio::task; -use log::{info, debug, trace}; pub async fn fluereflow_fileparse(arg: Args) { let csv_file = arg.files.csv.unwrap(); @@ -63,15 +63,14 @@ pub async fn fluereflow_fileparse(arg: Args) { if flags.syn > 0 { active_flow.insert(key_value, flowdata); - trace!("flow established"); - + trace!("flow established"); } else { continue; } } else { active_flow.insert(key_value, flowdata); - trace!("flow established"); + trace!("flow established"); } false @@ -115,10 +114,10 @@ pub async fn fluereflow_fileparse(arg: Args) { }; update_flow(flow, is_reverse, update_key); - trace!( - "{} flow updated", - if is_reverse { "reverse" } else { "forward" } - ); + trace!( + "{} flow updated", + if is_reverse { "reverse" } else { "forward" } + ); if flags.fin == 1 || flags.rst == 1 { trace!("flow finished"); @@ -130,7 +129,7 @@ pub async fn fluereflow_fileparse(arg: Args) { info!("Captured in {:?}", start.elapsed()); let ac_flow_cnt = active_flow.len(); let ended_flow_cnt = records.len(); - + for (_key, flow) in active_flow.clone().iter() { records.push(*flow); } @@ -142,5 +141,5 @@ pub async fn fluereflow_fileparse(arg: Args) { info!("Export {} result: {:?}", file_path, result); println!("Active flow {:?}", ac_flow_cnt); - println!("Ended flow {:?}", ended_flow_cnt); + println!("Ended flow {:?}", ended_flow_cnt); } diff --git a/src/net/online_fluereflow.rs b/src/net/online_fluereflow.rs index 0bbb4d7..4936232 100644 --- a/src/net/online_fluereflow.rs +++ b/src/net/online_fluereflow.rs @@ -1,15 +1,12 @@ // This file contains the implementation of the online packet capture functionality.online // It uses the pcap library to capture packets from a network interface and the fluereflow library to convert the packets into NetFlow data. // The data is then exported to a CSV file. -extern crate csv; -use fluere_config::Config; -use fluere_plugin::PluginManager; -use fluereflow::FluereRecord; - -use tokio::task; - -use log::{info, debug, trace}; +use std::{ + collections::HashMap, + fs, + time::{Duration, Instant}, +}; use crate::{ net::{ @@ -23,11 +20,13 @@ use crate::{ utils::{cur_time_file, fluere_exporter}, }; -use std::{ - collections::HashMap, - fs, - time::{Duration, Instant}, -}; +use fluere_config::Config; +use fluere_plugin::PluginManager; +use fluereflow::FluereRecord; + +use tokio::task; + +use log::{debug, info, trace}; // This function captures packets from a network interface and converts them into NetFlow data. // It takes the command line arguments as input, which specify the network interface to capture from and other parameters. @@ -156,7 +155,6 @@ pub async fn packet_capture(arg: Args) { "{} flow updated", if is_reverse { "reverse" } else { "forward" } ); - if flags.fin == 1 || flags.rst == 1 { trace!("flow finished"); @@ -172,11 +170,11 @@ pub async fn packet_capture(arg: Args) { packet_count += 1; // slow down the loop for windows to avoid random shutdown // if packet_count % sleep_windows == 0 && cfg!(target_os = "windows") { - // if verbose >= 3 { - // println!("Slow down the loop for windows"); - // } - // sleep(Duration::from_millis(0)).await; - // } + // if verbose >= 3 { + // println!("Slow down the loop for windows"); + // } + // sleep(Duration::from_millis(0)).await; + // } // Export flows if the interval has been reached if last_export.elapsed() >= Duration::from_millis(interval) && interval != 0 { diff --git a/src/net/packet_pcap.rs b/src/net/packet_pcap.rs index 750f180..82138de 100644 --- a/src/net/packet_pcap.rs +++ b/src/net/packet_pcap.rs @@ -6,16 +6,14 @@ use crate::net::CaptureDevice; use crate::types::Args; use crate::utils::cur_time_file; -use tokio::time::sleep; -use log::{info, debug, trace}; - +use log::{debug, trace}; pub async fn pcap_capture(args: Args) { let pcap_file = args.files.pcap.unwrap(); let interface_name = args.interface.expect("interface not found"); let duration = args.parameters.duration.unwrap(); let _interval = args.parameters.interval.unwrap(); - let sleep_windows = args.parameters.sleep_windows.unwrap(); + let _sleep_windows = args.parameters.sleep_windows.unwrap(); let interface = find_device(interface_name.as_str()).unwrap(); let mut cap_device = CaptureDevice::new(interface.clone()).unwrap(); diff --git a/src/utils/fluere_csv_exporter.rs b/src/utils/fluere_csv_exporter.rs index 0f86bf2..6ff0ff9 100644 --- a/src/utils/fluere_csv_exporter.rs +++ b/src/utils/fluere_csv_exporter.rs @@ -1,5 +1,5 @@ use fluereflow::FluereRecord; -use log::{ debug, trace}; +use log::{debug, trace}; use std::fs::File; pub fn fluere_exporter(records: Vec, file: File) {