From ca2c1bc4db5c974649c908c08dead99b1e5f10c6 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 5 Feb 2024 09:45:25 +0900 Subject: [PATCH 01/44] refactor: separate cli function from main --- src/cli.rs | 239 +++++++++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 240 +--------------------------------------------------- 2 files changed, 241 insertions(+), 238 deletions(-) create mode 100644 src/cli.rs diff --git a/src/cli.rs b/src/cli.rs new file mode 100644 index 0000000..49ecf35 --- /dev/null +++ b/src/cli.rs @@ -0,0 +1,239 @@ +use clap::{Arg, ArgAction, Command}; + +// This function sets up the command line interface for the application using the clap library. +// It defines the available commands and their arguments. +pub fn cli_template() -> Command { + Command::new("fluere") + .version("0.6.2") + .author("Skuld Norniern. ") + .about("Netflow Capture Tool") + .subcommand_required(true) + .subcommand( + Command::new("online") + .about("Capture netflow online") + .arg( + Arg::new("csv") + .help("Title of the exported csv file") + .short('c') + .long("csv") + .default_value("output"), + ) + .arg( + Arg::new("list") + .help("List of network interfaces") + .short('l') + .long("list") + .action(ArgAction::SetTrue), + ) + .arg( + Arg::new("interface") + .help("Select network interface to use [Required]") + .short('i') + .long("interface") + //.required(true), + ) + .arg( + Arg::new("duration") + .help("Set capture duration, in milliseconds (0: infinite)") + .default_value("0") + .short('d') + .long("duration"), + ) + .arg( + Arg::new("timeout") + .help("Set flow timeout, in milliseconds (0: infinite)") + .default_value("600000") + .short('t') + .long("timeout"), + ) + .arg( + Arg::new("useMACaddress") + .help("Set use MAC address on Key value [default: false]") + .short('M') + .long("useMAC") + .action(ArgAction::SetTrue), + ) + .arg( + Arg::new("interval") + .help("Set export interval, in milliseconds") + .default_value("1800000") + .short('I') + .long("interval"), + ) + .arg( + Arg::new("sleep_windows") + .help("Set inverval of thread pause for (only)MS Windows per n packet (need it for stopping random stop on Windows)") + .default_value("10") + .short('s') + .long("sleep"), + ) + .arg( + Arg::new("verbose") + .help("Set verbosity level") + .default_value("1") + .short('v') + .long("verbose"), // 0: quiet, 1: normal,2: extended, 3: verbose + ), + ) + .subcommand( + Command::new("offline") + .about("Convet pcap files to netflow") + .arg( + Arg::new("file") + .help("Name of the input pcap file [Required]") + .short('f') + .long("file") + .required(true), + ) + .arg( + Arg::new("csv") + .help("Title of the exported csv file") + .short('c') + .long("csv") + .default_value("output"), + ) + .arg( + Arg::new("timeout") + .help("Set flow timeout, in milliseconds (0: infinite)") + .default_value("600000") + .short('t') + .long("timeout"), + ) + .arg( + Arg::new("useMACaddress") + .help("Set use MAC address on Key value [default: false]") + .short('M') + .long("useMAC") + .action(ArgAction::SetTrue), + ) + .arg( + Arg::new("verbose") + .help("Set verbosity level") + .default_value("1") + .short('v') + .long("verbose"), // 0: quiet, 1: normal,2: extended, 3: verbose + ), + ) + .subcommand( + Command::new("live") + .about("Capture netflow online with live TUI feedback") + .arg( + Arg::new("csv") + .help("Title of the exported csv file") + .short('c') + .long("csv") + .default_value("output"), + ) + .arg( + Arg::new("list") + .help("List of network interfaces") + .short('l') + .long("list") + .action(ArgAction::SetTrue), + ) + .arg( + Arg::new("interface") + .help("Select network interface to use [Required]") + .short('i') + .long("interface") + //.required(true), + ) + .arg( + Arg::new("duration") + .help("Set capture duration, in milliseconds (0: infinite)") + .default_value("0") + .short('d') + .long("duration"), + ) + .arg( + Arg::new("timeout") + .help("Set flow timeout, in milliseconds (0: infinite)") + .default_value("600000") + .short('t') + .long("timeout"), + ) + .arg( + Arg::new("useMACaddress") + .help("Set use MAC address on Key value [default: false]") + .short('M') + .long("useMAC") + .action(ArgAction::SetTrue), + ) + .arg( + Arg::new("interval") + .help("Set export interval, in milliseconds") + .default_value("1800000") + .short('I') + .long("interval"), + ) + .arg( + Arg::new("sleep_windows") + .help("Set inverval of thread pause for (only)MS Windows per n packet (need it for stopping random stop on Windows)") + .default_value("10") + .short('s') + .long("sleep"), + ) + .arg( + Arg::new("verbose") + .help("Set verbosity level") + .default_value("1") + .short('v') + .long("verbose"), // 0: quiet, 1: normal,2: extended, 3: verbose + ), + ) + .subcommand( + Command::new("pcap") + .about("Collect packet and save to .pcap file") + .arg( + Arg::new("pcap") + .help("Name of the output pcap files title [Required]") + .short('p') + .long("pcap") + //.required(true), + ) + .arg( + Arg::new("interface") + .help("Select network interface to use [Required]") + .short('i') + .long("interface") + //.required(true), + ) + .arg( + Arg::new("duration") + .help("Set capture duration, in milliseconds (0: infinite)") + .default_value("0") + .short('d') + .long("duration"), + ) + .arg( + Arg::new("interval") + .help("Set export interval, in milliseconds") + .default_value("1800000") + .short('I') + .long("interval"), + ) + .arg( + Arg::new("sleep_windows") + .help("Set inverval of thread pause for (only)MS Windows per n packet (need it for stopping random stop on Windows)") + .default_value("10") + .short('s') + .long("sleep"), + ) + .arg( + Arg::new("list") + .help("List of network interfaces") + .short('l') + .long("list") + .action(ArgAction::SetTrue), + ) + .arg( + Arg::new("verbose") + .help("Set verbosity level") + .default_value("1") + .short('v') + .long("verbose"), // 0: quiet, 1: normal,2: extended, 3: verbose + ), + ) +} + + diff --git a/src/main.rs b/src/main.rs index e47dc9d..91e1a05 100644 --- a/src/main.rs +++ b/src/main.rs @@ -7,255 +7,19 @@ pub mod net; pub mod plugin; pub mod types; pub mod utils; +pub mod cli; -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. -// It defines the available commands and their arguments. -fn cli() -> Command { - Command::new("fluere") - .version("0.6.2") - .author("Skuld Norniern. ") - .about("Netflow Capture Tool") - .subcommand_required(true) - .subcommand( - Command::new("online") - .about("Capture netflow online") - .arg( - Arg::new("csv") - .help("Title of the exported csv file") - .short('c') - .long("csv") - .default_value("output"), - ) - .arg( - Arg::new("list") - .help("List of network interfaces") - .short('l') - .long("list") - .action(ArgAction::SetTrue), - ) - .arg( - Arg::new("interface") - .help("Select network interface to use [Required]") - .short('i') - .long("interface") - //.required(true), - ) - .arg( - Arg::new("duration") - .help("Set capture duration, in milliseconds (0: infinite)") - .default_value("0") - .short('d') - .long("duration"), - ) - .arg( - Arg::new("timeout") - .help("Set flow timeout, in milliseconds (0: infinite)") - .default_value("600000") - .short('t') - .long("timeout"), - ) - .arg( - Arg::new("useMACaddress") - .help("Set use MAC address on Key value [default: false]") - .short('M') - .long("useMAC") - .action(ArgAction::SetTrue), - ) - .arg( - Arg::new("interval") - .help("Set export interval, in milliseconds") - .default_value("1800000") - .short('I') - .long("interval"), - ) - .arg( - Arg::new("sleep_windows") - .help("Set inverval of thread pause for (only)MS Windows per n packet (need it for stopping random stop on Windows)") - .default_value("10") - .short('s') - .long("sleep"), - ) - .arg( - Arg::new("verbose") - .help("Set verbosity level") - .default_value("1") - .short('v') - .long("verbose"), // 0: quiet, 1: normal,2: extended, 3: verbose - ), - ) - .subcommand( - Command::new("offline") - .about("Convet pcap files to netflow") - .arg( - Arg::new("file") - .help("Name of the input pcap file [Required]") - .short('f') - .long("file") - .required(true), - ) - .arg( - Arg::new("csv") - .help("Title of the exported csv file") - .short('c') - .long("csv") - .default_value("output"), - ) - .arg( - Arg::new("timeout") - .help("Set flow timeout, in milliseconds (0: infinite)") - .default_value("600000") - .short('t') - .long("timeout"), - ) - .arg( - Arg::new("useMACaddress") - .help("Set use MAC address on Key value [default: false]") - .short('M') - .long("useMAC") - .action(ArgAction::SetTrue), - ) - .arg( - Arg::new("verbose") - .help("Set verbosity level") - .default_value("1") - .short('v') - .long("verbose"), // 0: quiet, 1: normal,2: extended, 3: verbose - ), - ) - .subcommand( - Command::new("live") - .about("Capture netflow online with live TUI feedback") - .arg( - Arg::new("csv") - .help("Title of the exported csv file") - .short('c') - .long("csv") - .default_value("output"), - ) - .arg( - Arg::new("list") - .help("List of network interfaces") - .short('l') - .long("list") - .action(ArgAction::SetTrue), - ) - .arg( - Arg::new("interface") - .help("Select network interface to use [Required]") - .short('i') - .long("interface") - //.required(true), - ) - .arg( - Arg::new("duration") - .help("Set capture duration, in milliseconds (0: infinite)") - .default_value("0") - .short('d') - .long("duration"), - ) - .arg( - Arg::new("timeout") - .help("Set flow timeout, in milliseconds (0: infinite)") - .default_value("600000") - .short('t') - .long("timeout"), - ) - .arg( - Arg::new("useMACaddress") - .help("Set use MAC address on Key value [default: false]") - .short('M') - .long("useMAC") - .action(ArgAction::SetTrue), - ) - .arg( - Arg::new("interval") - .help("Set export interval, in milliseconds") - .default_value("1800000") - .short('I') - .long("interval"), - ) - .arg( - Arg::new("sleep_windows") - .help("Set inverval of thread pause for (only)MS Windows per n packet (need it for stopping random stop on Windows)") - .default_value("10") - .short('s') - .long("sleep"), - ) - .arg( - Arg::new("verbose") - .help("Set verbosity level") - .default_value("1") - .short('v') - .long("verbose"), // 0: quiet, 1: normal,2: extended, 3: verbose - ), - ) - .subcommand( - Command::new("pcap") - .about("Collect packet and save to .pcap file") - .arg( - Arg::new("pcap") - .help("Name of the output pcap files title [Required]") - .short('p') - .long("pcap") - //.required(true), - ) - .arg( - Arg::new("interface") - .help("Select network interface to use [Required]") - .short('i') - .long("interface") - //.required(true), - ) - .arg( - Arg::new("duration") - .help("Set capture duration, in milliseconds (0: infinite)") - .default_value("0") - .short('d') - .long("duration"), - ) - .arg( - Arg::new("interval") - .help("Set export interval, in milliseconds") - .default_value("1800000") - .short('I') - .long("interval"), - ) - .arg( - Arg::new("sleep_windows") - .help("Set inverval of thread pause for (only)MS Windows per n packet (need it for stopping random stop on Windows)") - .default_value("10") - .short('s') - .long("sleep"), - ) - .arg( - Arg::new("list") - .help("List of network interfaces") - .short('l') - .long("list") - .action(ArgAction::SetTrue), - ) - .arg( - Arg::new("verbose") - .help("Set verbosity level") - .default_value("1") - .short('v') - .long("verbose"), // 0: quiet, 1: normal,2: extended, 3: verbose - ), - ) -} - // This is the main function of the application. // It gets the command line arguments, parses them, and calls the appropriate functions based on the arguments. #[tokio::main] async fn main() { - let args = cli().get_matches(); + let args = cli::cli_template().get_matches(); let interfaces = datalink::interfaces(); //let _plugins = scan_plugins("plugins"); //println!("Plugins: {:?}", plugins); //match generate_config() { From 647e34c8910df0d7bc9d970a23d3d83ced6b0c0e Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 5 Feb 2024 11:06:40 +0900 Subject: [PATCH 02/44] build: move to `log` from `env_logger` --- Cargo.toml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 362f1e3..515166f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,8 +35,7 @@ fluereflow = { version = "0.3.2", path = "./fluereflow" } ratatui = { version = "0.25.0", features = ["all-widgets"] } crossterm = "0.27" dirs = "5.0.1" -logs = "0.7.1" -env_logger = "0.10.0" +log = "0.4.20" [workspace] members = [ From e3f2d8f35deda16bbf78dedc7b54141c9b61f920 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 5 Feb 2024 11:35:22 +0900 Subject: [PATCH 03/44] feat: implement simple logger --- src/logger.rs | 78 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) create mode 100644 src/logger.rs diff --git a/src/logger.rs b/src/logger.rs new file mode 100644 index 0000000..baab109 --- /dev/null +++ b/src/logger.rs @@ -0,0 +1,78 @@ +use std::fs::File; +use std::io::{ Write, stdout, stderr}; +use std::path::PathBuf; + +use chrono::Local; // Import the Local struct from the chrono crate +use log::{Level, Record, Metadata, Log }; + +pub enum Logstdout{ + Stdout, + StdErr, +} + +pub struct Logger { + write_to_file: bool, + write_to_std: Option, + severity: Level, + file: Option, +} + +impl Logger { + pub fn new(write_to_file: bool, file_path: Option) -> Self { + let mut path = file_path; + if path.is_none() { + path = Some(PathBuf::from( + #[cfg(target_os = "linux")] + "/var/log/fluere/fluere.log", + #[cfg(target_os = "windows")] + "C:\\Program Files\\fluere\\fluere.log", + #[cfg(target_os = "macos")] + "/Library/Logs/fluere/fluere.log", + #[cfg(target_os = "bsd")] + "/var/log/fluere/fluere.log", + #[cfg(not(any(target_os = "linux", target_os = "windows", target_os = "macos", target_os = "bsd")))] + "/var/log/fluere/fluere.log", + )); + } + let mut file = None; + if write_to_file { + file = Some(File::create(path.as_ref().unwrap()).unwrap()); + } + Logger {write_to_file: true, write_to_std: None, severity: Level::Info , file} + } + + // pub fn log(&mut self, severity: Level, message: &str) { + // let timestamp = Local::now(); // Get the current timestamp using Local::now() + // let log_message = format!("{:?} {}: {}", timestamp, severity, message); // Format the timestamp and append it to the log message + // } +} + +impl Log for Logger { + fn enabled(&self, _metadata: &Metadata) -> bool { + true + } + + fn log(&self, record: &Record) { + let timestamp = Local::now(); + + if self.write_to_std.as_ref().is_some() && record.level() <= self.severity{ + match self.write_to_std.as_ref().unwrap() { + Logstdout::Stdout => { + writeln!(stdout(), "[{}]: {} {}",timestamp, record.level(), record.args()).unwrap(); + } + Logstdout::StdErr => { + writeln!(stderr(), "[{}]: {} {}",timestamp, record.level(), record.args()).unwrap(); + } + } + } + + if self.write_to_file { + + writeln!(self.file.as_ref().unwrap(), "[{}]: {} {}",timestamp, record.level(), record.args()).unwrap(); + } + } + + fn flush(&self) { + + } +} From 899ce01984679ba0292b49d100130c9c99b2377e Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Tue, 13 Feb 2024 10:47:35 +0900 Subject: [PATCH 04/44] style: `rustfmt` --- src/cli.rs | 2 - src/logger.rs | 62 +++++++++++++++------ src/main.rs | 49 +++++++++++++++- src/net/capture.rs | 2 +- src/net/packet_pcap.rs | 4 +- src/utils/log.rs | 124 ++++++++++++++++++++++++++++++----------- 6 files changed, 189 insertions(+), 54 deletions(-) diff --git a/src/cli.rs b/src/cli.rs index 49ecf35..591bb29 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -235,5 +235,3 @@ pub fn cli_template() -> Command { ), ) } - - diff --git a/src/logger.rs b/src/logger.rs index baab109..cb14ba4 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -1,11 +1,11 @@ use std::fs::File; -use std::io::{ Write, stdout, stderr}; +use std::io::{stderr, stdout, Write}; use std::path::PathBuf; use chrono::Local; // Import the Local struct from the chrono crate -use log::{Level, Record, Metadata, Log }; +use log::{Level, Log, Metadata, Record}; -pub enum Logstdout{ +pub enum Logstdout { Stdout, StdErr, } @@ -21,7 +21,7 @@ impl Logger { pub fn new(write_to_file: bool, file_path: Option) -> Self { let mut path = file_path; if path.is_none() { - path = Some(PathBuf::from( + path = Some(PathBuf::from( #[cfg(target_os = "linux")] "/var/log/fluere/fluere.log", #[cfg(target_os = "windows")] @@ -30,20 +30,30 @@ impl Logger { "/Library/Logs/fluere/fluere.log", #[cfg(target_os = "bsd")] "/var/log/fluere/fluere.log", - #[cfg(not(any(target_os = "linux", target_os = "windows", target_os = "macos", target_os = "bsd")))] + #[cfg(not(any( + target_os = "linux", + target_os = "windows", + target_os = "macos", + target_os = "bsd" + )))] "/var/log/fluere/fluere.log", )); } let mut file = None; if write_to_file { file = Some(File::create(path.as_ref().unwrap()).unwrap()); - } - Logger {write_to_file: true, write_to_std: None, severity: Level::Info , file} + } + Logger { + write_to_file: true, + write_to_std: None, + severity: Level::Info, + file, + } } // pub fn log(&mut self, severity: Level, message: &str) { - // let timestamp = Local::now(); // Get the current timestamp using Local::now() - // let log_message = format!("{:?} {}: {}", timestamp, severity, message); // Format the timestamp and append it to the log message + // let timestamp = Local::now(); // Get the current timestamp using Local::now() + // let log_message = format!("{:?} {}: {}", timestamp, severity, message); // Format the timestamp and append it to the log message // } } @@ -55,24 +65,42 @@ impl Log for Logger { fn log(&self, record: &Record) { let timestamp = Local::now(); - if self.write_to_std.as_ref().is_some() && record.level() <= self.severity{ + if self.write_to_std.as_ref().is_some() && record.level() <= self.severity { match self.write_to_std.as_ref().unwrap() { Logstdout::Stdout => { - writeln!(stdout(), "[{}]: {} {}",timestamp, record.level(), record.args()).unwrap(); + writeln!( + stdout(), + "[{}]: {}: {}", + timestamp, + record.level(), + record.args() + ) + .unwrap(); } Logstdout::StdErr => { - writeln!(stderr(), "[{}]: {} {}",timestamp, record.level(), record.args()).unwrap(); + writeln!( + stderr(), + "[{}]: {}: {}", + timestamp, + record.level(), + record.args() + ) + .unwrap(); } } } if self.write_to_file { - - writeln!(self.file.as_ref().unwrap(), "[{}]: {} {}",timestamp, record.level(), record.args()).unwrap(); + writeln!( + self.file.as_ref().unwrap(), + "[{}]: {}: {}", + timestamp, + record.level(), + record.args() + ) + .unwrap(); } } - fn flush(&self) { - - } + fn flush(&self) {} } diff --git a/src/main.rs b/src/main.rs index 91e1a05..983a937 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,18 +3,65 @@ // It also supports live capture and conversion of NetFlow data. // This file contains the main function which parses the command line arguments and calls the appropriate functions based on the arguments. +pub mod cli; +pub mod logger; pub mod net; pub mod plugin; pub mod types; pub mod utils; -pub mod cli; use pnet::datalink; +// use env_logger::{init, Logger}; +use log::Level; +use crate::logger::Logger; use crate::net::list_devices; +use std::fmt::Display; use std::process::exit; +enum Mode { + Offline, + Online, + Live, + Pcap, +} +impl Display for Mode { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Mode::Offline => write!(f, "Offline"), + Mode::Online => write!(f, "Online"), + Mode::Live => write!(f, "Live"), + Mode::Pcap => write!(f, "Pcap"), + } + } +} + +struct Fluere { + interface: String, + args: types::Args, + mode: Mode, + logger: Logger, + verbose: Level, +} +impl Fluere { + fn new( + interface: String, + args: types::Args, + mode: Mode, + logger: Logger, + verbose: Level, + ) -> Fluere { + Fluere { + interface, + args, + mode, + logger, + verbose, + } + } +} + // This is the main function of the application. // It gets the command line arguments, parses them, and calls the appropriate functions based on the arguments. #[tokio::main] diff --git a/src/net/capture.rs b/src/net/capture.rs index abae835..6f1feca 100644 --- a/src/net/capture.rs +++ b/src/net/capture.rs @@ -1,7 +1,7 @@ use pcap::{Active, Address, Capture, Device, Error as PcapError}; -use std::time::Instant; use std::fmt; +use std::time::Instant; #[derive(Debug)] pub enum DeviceError { diff --git a/src/net/packet_pcap.rs b/src/net/packet_pcap.rs index f2c9aa2..b727191 100644 --- a/src/net/packet_pcap.rs +++ b/src/net/packet_pcap.rs @@ -4,10 +4,10 @@ use std::fs; use tokio::time::sleep; +use crate::net::find_device; +use crate::net::CaptureDevice; use crate::types::Args; use crate::utils::cur_time_file; -use crate::net::CaptureDevice; -use crate::net::find_device; use std::time::{Duration, Instant}; diff --git a/src/utils/log.rs b/src/utils/log.rs index 77666b1..8f3c759 100644 --- a/src/utils/log.rs +++ b/src/utils/log.rs @@ -1,44 +1,106 @@ use std::fs::File; -use std::io::Write; -use std::path::Path; +use std::io::{stderr, stdout, Write}; +use std::path::PathBuf; + use chrono::Local; // Import the Local struct from the chrono crate +use log::{Level, Log, Metadata, Record}; -pub struct Log { - file: File, +pub enum Logstdout { + Stdout, + StdErr, } -impl Log { - #[cfg(target_os = "linux")] - pub fn new() -> Self { - let path = "/var/log/fluere/fluere.log"; - let file = File::create(&path).expect("Failed to create log file"); - Log { file } - } +pub struct Logger { + write_to_file: bool, + write_to_std: Option, + severity: Level, + file: Option, +} - #[cfg(target_os = "windows")] - pub fn new() -> Self { - let path = "C:\\Program Files\\fluere\\fluere.log"; - let file = File::create(&path).expect("Failed to create log file"); - Log { file } +impl Logger { + pub fn new(write_to_file: bool, file_path: Option) -> Self { + let mut path = file_path; + if path.is_none() { + path = Some(PathBuf::from( + #[cfg(target_os = "linux")] + "/var/log/fluere/fluere.log", + #[cfg(target_os = "windows")] + "C:\\Program Files\\fluere\\fluere.log", + #[cfg(target_os = "macos")] + "/Library/Logs/fluere/fluere.log", + #[cfg(target_os = "bsd")] + "/var/log/fluere/fluere.log", + #[cfg(not(any( + target_os = "linux", + target_os = "windows", + target_os = "macos", + target_os = "bsd" + )))] + "/var/log/fluere/fluere.log", + )); + } + let mut file = None; + if write_to_file { + file = Some(File::create(path.as_ref().unwrap()).unwrap()); + } + Logger { + write_to_file: true, + write_to_std: None, + severity: Level::Info, + file, + } } - #[cfg(target_os = "macos")] - pub fn new() -> Self { - let path = "/Library/Logs/fluere/fluere.log"; - let file = File::create(&path).expect("Failed to create log file"); - Log { file } - } + // pub fn log(&mut self, severity: Level, message: &str) { + // let timestamp = Local::now(); // Get the current timestamp using Local::now() + // let log_message = format!("{:?} {}: {}", timestamp, severity, message); // Format the timestamp and append it to the log message + // } +} - #[cfg(target_os = "bsd")] - pub fn new() -> Self { - let path = "/var/log/fluere/fluere.log"; - let file = File::create(&path).expect("Failed to create log file"); - Log { file } +impl Log for Logger { + fn enabled(&self, _metadata: &Metadata) -> bool { + true } - pub fn log(&mut self, severity: Severity, message: &str) { - let timestamp = Local::now(); // Get the current timestamp using Local::now() - let log_message = format!("{:?} {}: {}", timestamp, severity, message); // Format the timestamp and append it to the log message - self.file.write_all(log_message.as_bytes()).expect("Failed to write to log file"); + fn log(&self, record: &Record) { + let timestamp = Local::now(); + + if self.write_to_std.as_ref().is_some() && record.level() <= self.severity { + match self.write_to_std.as_ref().unwrap() { + Logstdout::Stdout => { + write!( + stdout(), + "[{}]: {} {}", + timestamp, + record.level(), + record.args() + ) + .unwrap(); + } + Logstdout::StdErr => { + write!( + stderr(), + "[{}]: {} {}", + timestamp, + record.level(), + record.args() + ) + .unwrap(); + } + } + } + + if self.write_to_file { + writeln!( + self.file.as_ref().unwrap(), + "[{}]: {} {}", + timestamp, + record.level(), + record.args() + ) + .unwrap(); + } } + + fn flush(&self) {} } From e112f25ac6b3a285d8754804d553844d8257fbce Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 09:30:23 +0900 Subject: [PATCH 05/44] build: bump `pcap-rs` --- Cargo.toml | 6 +++-- src/net/capture.rs | 43 +++++++++++++++++------------------- src/net/live_fluereflow.rs | 4 ++-- src/net/online_fluereflow.rs | 4 ++-- src/net/packet_pcap.rs | 4 ++-- 5 files changed, 30 insertions(+), 31 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 515166f..bf6b31a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,7 +18,9 @@ pnet_macros_support = "0.34.0" pnet_macros = "0.34.0" # using custom forked version of pcap-rs for fixing audits -pcap = { version = "1.1.0", git = "https://github.com/SkuldNorniern/pcap", rev = "40f1163" } +# pcap = { version = "1.1.0", git = "https://github.com/SkuldNorniern/pcap", rev = "40f1163" } + +pcap = "1.2.0" chrono = { version = "0.4.29", default-features = false, features = ["clock"] } libc = "0.2" @@ -35,7 +37,7 @@ fluereflow = { version = "0.3.2", path = "./fluereflow" } ratatui = { version = "0.25.0", features = ["all-widgets"] } crossterm = "0.27" dirs = "5.0.1" -log = "0.4.20" +log = { version = "0.4.21", features = ["std"]} [workspace] members = [ diff --git a/src/net/capture.rs b/src/net/capture.rs index 6f1feca..a7b424a 100644 --- a/src/net/capture.rs +++ b/src/net/capture.rs @@ -1,27 +1,19 @@ -use pcap::{Active, Address, Capture, Device, Error as PcapError}; +use std::{fmt, time::Instant}; + +use crate::net::NetError; -use std::fmt; -use std::time::Instant; +use pcap::{Active, Address, Capture, Device, Error as PcapError}; #[derive(Debug)] pub enum DeviceError { - Cap(PcapError), DeviceNotFound(String), - InitializationError(), InvalidDeviceIndex(usize), } -impl From for DeviceError { - fn from(err: PcapError) -> Self { - DeviceError::Cap(err) - } -} impl fmt::Display for DeviceError { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self { - DeviceError::Cap(err) => err.fmt(f), DeviceError::DeviceNotFound(err) => write!(f, "Device not found: {}", err), - DeviceError::InitializationError() => write!(f, "Initialization error"), DeviceError::InvalidDeviceIndex(err) => write!(f, "Invalid device index: {}", err), } } @@ -35,7 +27,7 @@ pub struct CaptureDevice { } impl CaptureDevice { - pub fn new(device: Device) -> Result { + pub fn new(device: Device) -> Result { let capture = initialize_capture(device.clone())?; Ok(CaptureDevice { @@ -47,15 +39,17 @@ impl CaptureDevice { } } -pub fn list_devices() -> Result, DeviceError> { - Device::list().map_err(DeviceError::Cap) +impl Drop for CaptureDevice { + fn drop(&mut self) { + println!("Closing capture session for device {}", self.name); + // self.capture.; + } } - -pub fn find_device(identifier: &str) -> Result { +pub fn find_device(identifier: &str) -> Result { let start = Instant::now(); println!("Requested Device: {}", identifier); - let devices = list_devices()?; + let devices = Device::list()?; if let Ok(index) = identifier.parse::() { if let Some(device) = devices.get(index) { @@ -63,7 +57,9 @@ pub fn find_device(identifier: &str) -> Result { println!("Device {} captured in {:?}", device.name, duration); return Ok(device.clone()); } else { - return Err(DeviceError::InvalidDeviceIndex(index)); + return Err(NetError::DeviceError(DeviceError::InvalidDeviceIndex( + index, + ))); } } @@ -75,12 +71,13 @@ pub fn find_device(identifier: &str) -> Result { } } - Err(DeviceError::DeviceNotFound(identifier.to_string())) + Err(NetError::DeviceError(DeviceError::DeviceNotFound( + identifier.to_string(), + ))) } -fn initialize_capture(device: Device) -> Result, DeviceError> { - Ok(Capture::from_device(device) - .map_err(DeviceError::Cap)? +fn initialize_capture(device: Device) -> Result, PcapError> { + Ok(Capture::from_device(device)? .promisc(true) .snaplen(1024) .timeout(60000) diff --git a/src/net/live_fluereflow.rs b/src/net/live_fluereflow.rs index 841ca0b..182062f 100644 --- a/src/net/live_fluereflow.rs +++ b/src/net/live_fluereflow.rs @@ -85,8 +85,8 @@ pub async fn online_packet_capture(arg: Args) { .expect("Failed to load plugins"); 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 cap_device = CaptureDevice::new(interface.clone()).unwrap(); + let cap = &mut cap_device.capture; let file_dir = "./output"; match fs::create_dir_all(<&str>::clone(&file_dir)) { diff --git a/src/net/online_fluereflow.rs b/src/net/online_fluereflow.rs index 88b0b31..da4cfc1 100644 --- a/src/net/online_fluereflow.rs +++ b/src/net/online_fluereflow.rs @@ -50,8 +50,8 @@ pub async fn packet_capture(arg: Args) { .expect("Failed to load plugins"); 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 cap_device = CaptureDevice::new(interface.clone()).unwrap(); + let cap = &mut cap_device.capture; // let mut cp_device // let mut cap = Capture::from_device(interface) // .unwrap() diff --git a/src/net/packet_pcap.rs b/src/net/packet_pcap.rs index b727191..d4446bf 100644 --- a/src/net/packet_pcap.rs +++ b/src/net/packet_pcap.rs @@ -20,8 +20,8 @@ pub async fn pcap_capture(args: Args) { let verbose = args.verbose.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 cap_device = CaptureDevice::new(interface.clone()).unwrap(); + let cap = &mut cap_device.capture; let file_dir = "./output"; let mut packet_count = 0; From e997986e095703332207c7092930ad3679e76a3d Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 09:47:41 +0900 Subject: [PATCH 06/44] refactor: separate cli as a file --- src/cli.rs | 146 +++++++++++++++++++++++++++++- src/main.rs | 254 ++++++++++++++++------------------------------------ 2 files changed, 222 insertions(+), 178 deletions(-) diff --git a/src/cli.rs b/src/cli.rs index 591bb29..145eaff 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -1,4 +1,9 @@ -use clap::{Arg, ArgAction, Command}; +use std::process::exit; + +use crate::types::{Args, Files, Parameters}; + +use clap::{Arg, ArgAction, ArgMatches, Command}; +use pcap::Device; // This function sets up the command line interface for the application using the clap library. // It defines the available commands and their arguments. @@ -235,3 +240,142 @@ pub fn cli_template() -> Command { ), ) } + +pub async fn handle_mode(mode: &str, args: &ArgMatches) -> Args { + let _verbose = args + .get_one::("verbose") + .map_or(0, |v| v.parse::().unwrap_or(0)); + if args.get_flag("list") { + println!("List of network interfaces"); + println!("--------------------------"); + for (i, device) in Device::list().unwrap().iter().enumerate() { + println!("[{}] {}", i, device.name); + } + exit(0); + } + + match mode { + "online" | "live" => parse_online_live_args(args, mode), + "offline" => parse_offline_args(args), + "pcap" => parse_pcap_args(args), + _ => unreachable!(), + } +} + +fn parse_online_live_args(args: &clap::ArgMatches, _mode: &str) -> Args { + let use_mac = args.get_flag("useMACaddress"); + let csv = args + .get_one::("csv") + .expect("CSV file not specified") + .to_string(); + let interface = args + .get_one::("interface") + .expect("Network interface not specified") + .to_string(); + let timeout = args + .get_one::("timeout") + .unwrap() + .parse::() + .unwrap(); + let duration = args + .get_one::("duration") + .unwrap() + .parse::() + .unwrap(); + let interval = args + .get_one::("interval") + .unwrap() + .parse::() + .unwrap(); + let sleep_windows = args + .get_one::("sleep_windows") + .unwrap() + .parse::() + .unwrap(); + let verbose = args + .get_one::("verbose") + .unwrap() + .parse::() + .unwrap(); + + Args::new( + Some(interface), + Files::new(Some(csv), None, None), + Parameters::new( + Some(use_mac), + Some(timeout), + Some(duration), + Some(interval), + Some(sleep_windows), + ), + Some(verbose), + ) +} +fn parse_offline_args(args: &clap::ArgMatches) -> Args { + let use_mac = args.get_flag("useMACaddress"); + let file = args + .get_one::("file") + .expect("File not specified") + .to_string(); + let csv = args.get_one::("csv").unwrap().to_string(); + let timeout = args + .get_one::("timeout") + .unwrap() + .parse::() + .unwrap(); + let verbose = args + .get_one::("verbose") + .unwrap() + .parse::() + .unwrap(); + + Args::new( + None, + Files::new(Some(csv), Some(file), None), + Parameters::new(Some(use_mac), Some(timeout), None, None, None), + Some(verbose), + ) +} +fn parse_pcap_args(args: &clap::ArgMatches) -> Args { + let pcap = args + .get_one::("pcap") + .expect("Output PCAP file name not specified") + .to_string(); + let interface = args + .get_one::("interface") + .expect("Network interface not specified") + .to_string(); + let duration = args + .get_one::("duration") + .unwrap() + .parse::() + .unwrap(); + let interval = args + .get_one::("interval") + .unwrap() + .parse::() + .unwrap(); + let sleep_windows = args + .get_one::("sleep_windows") + .unwrap() + .parse::() + .unwrap(); + let verbose = args + .get_one::("verbose") + .unwrap() + .parse::() + .unwrap(); + + Args::new( + Some(interface), + Files::new(None, None, Some(pcap)), + Parameters::new( + None, + None, + Some(duration), + Some(interval), + Some(sleep_windows), + ), + Some(verbose), + ) +} diff --git a/src/main.rs b/src/main.rs index 983a937..f01dd59 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,15 +10,47 @@ pub mod plugin; pub mod types; pub mod utils; -use pnet::datalink; -// use env_logger::{init, Logger}; -use log::Level; +use std::{fmt::Display, process::exit}; +use std::fs::File; -use crate::logger::Logger; -use crate::net::list_devices; +use crate::logger::{Logger, Logstdout}; +use crate::net::capture::DeviceError; +// use env_logger;::{init, Logger}; -use std::fmt::Display; -use std::process::exit; +use log::{Level, Log, info, warn, error, debug, trace}; + + +// 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}; + +#[derive(Debug)] +enum FluereError { + InterfaceNotFound, + DeviceError(DeviceError), + ArgumentParseError(String), + ModeNotSupported(String), + NetworkError(String), +} + +impl std::fmt::Display for FluereError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + FluereError::InterfaceNotFound => write!(f, "Network interface not found."), + FluereError::ArgumentParseError(msg) => write!(f, "Argument parsing error: {}", msg), + FluereError::ModeNotSupported(mode) => write!(f, "Mode not supported: {}", mode), + FluereError::NetworkError(msg) => write!(f, "Network error: {}", msg), + FluereError::DeviceError(err) => err.fmt(f), + } + } +} + +impl std::error::Error for FluereError {} + +impl From for FluereError { + fn from(err: DeviceError) -> Self { + FluereError::DeviceError(err) + } +} enum Mode { Offline, @@ -60,185 +92,53 @@ impl Fluere { verbose, } } + async fn online_fluereflow(&self, params: types::Args) { + net::online_fluereflow::packet_capture(params).await; + info!("Online fluereflow mode completed"); + } } - // This is the main function of the application. // It gets the command line arguments, parses them, and calls the appropriate functions based on the arguments. #[tokio::main] async fn main() { let args = cli::cli_template().get_matches(); - let interfaces = datalink::interfaces(); //let _plugins = scan_plugins("plugins"); - //println!("Plugins: {:?}", plugins); - //match generate_config() { - // Ok(_) => println!("Config file generated"), - // Err(e) => println!("Error: {e}"), - //} - //let mut interface = "None"; - match args.subcommand() { - Some(("online", args)) => { - println!("Online mode"); - utils::get_local_ip(); - if args.get_flag("list") { - 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"); - let csv = args.get_one::("csv").expect("default"); - let interface = args - .get_one::("interface") - .ok_or("Required Interface") - .unwrap(); - - let timeout = args.get_one::("timeout").unwrap(); - let timeout: u64 = timeout.parse().unwrap(); - let duration = args.get_one::("duration").expect("default"); - let duration: u64 = duration.parse().unwrap(); - let interval = args.get_one::("interval").expect("default"); - let interval: u64 = interval.parse().unwrap(); - let sleep_windows = args.get_one::("sleep_windows").expect("default"); - let sleep_windows: u64 = sleep_windows.parse().unwrap(); - let verbose = args.get_one::("verbose").expect("default"); - let verbose: u8 = verbose.parse().unwrap(); - - let args: types::Args = types::Args::new( - Some(interface.to_string()), - types::Files::new(Some(csv.to_string()), None, None), - types::Parameters::new( - Some(use_mac), - Some(timeout), - Some(duration), - Some(interval), - Some(sleep_windows), - ), - Some(verbose), - ); - if verbose >= 1 { - println!("Interface {} selected", interface); - } //net::packet_capture(interface); - net::online_fluereflow::packet_capture(args).await; - //net::netflow(_interface); - } - Some(("offline", args)) => { - println!("Offline mode"); - let use_mac = args.get_flag("useMACaddress"); - let file = args.get_one::("file").unwrap(); - let csv = args.get_one::("csv").expect("default"); - let timeout = args.get_one::("timeout").unwrap(); - let timeout: u64 = timeout.parse().unwrap(); - let verbose = args.get_one::("verbose").expect("default"); - let verbose: u8 = verbose.parse().unwrap(); - - let args: types::Args = types::Args::new( - None, - types::Files::new(Some(csv.to_string()), Some(file.to_string()), None), - types::Parameters::new(Some(use_mac), Some(timeout), None, None, None), - Some(verbose), - ); - - net::fluereflow_fileparse(args).await; - //net::netflow(_file, _csv); - } - Some(("live", args)) => { - println!("Live mode"); - if args.get_flag("list") { - let interfaces = list_devices().unwrap(); - println!("Found {} devices", interfaces.len()); - for (i, interface) in interfaces.iter().enumerate() { - println!("[{}]: {}", i, interface.name); + let log_stdout = Logstdout::Stdout; + let log_file :Option = None; + let log_level = Level::Info; + let logger = Logger::new(None,Some(Level::Trace), Some(Logstdout::Stdout),false); + + let _ = log::set_boxed_logger(Box::new(logger)) + .map(|()| log::set_max_level(log::LevelFilter::Info)); + // let mode = match args.subcommand() { + // Some((mode, _sub_args)) => mode, + // None => { + // log::error!("No mode selected. Use --help for more information."); + // exit(1); + // } + // }; + + info!("Fluere started"); + if let Some((mode, sub_args)) = args.subcommand() { + match mode { + "online" | "offline" | "live" | "pcap" => { + log::debug!("Mode: {}", mode); + let parems = cli::handle_mode(mode, sub_args).await; + + match mode { + "online" => net::online_fluereflow::packet_capture(parems).await, + "offline" => net::fluereflow_fileparse(parems).await, + "live" => net::live_fluereflow::packet_capture(parems) + .await + .expect("Error on live mode"), + "pcap" => net::pcap_capture(parems).await, + _ => unreachable!(), } - exit(0); } - let use_mac = args.get_flag("useMACaddress"); - let csv = args.get_one::("csv").expect("default"); - let interface = args - .get_one::("interface") - .ok_or("Required Interface") - .unwrap(); - let timeout = args.get_one::("timeout").unwrap(); - let timeout: u64 = timeout.parse().unwrap(); - let duration = args.get_one::("duration").expect("default"); - let duration: u64 = duration.parse().unwrap(); - let interval = args.get_one::("interval").expect("default"); - let interval: u64 = interval.parse().unwrap(); - let sleep_windows = args.get_one::("sleep_windows").expect("default"); - let sleep_windows: u64 = sleep_windows.parse().unwrap(); - let verbose = args.get_one::("verbose").expect("default"); - let verbose: u8 = verbose.parse().unwrap(); - - let args: types::Args = types::Args::new( - Some(interface.to_string()), - types::Files::new(Some(csv.to_string()), None, None), - types::Parameters::new( - Some(use_mac), - Some(timeout), - Some(duration), - Some(interval), - Some(sleep_windows), - ), - Some(verbose), - ); - if verbose >= 1 { - println!("Interface {} selected", interface); - } //net::packet_capture(interface); - net::live_fluereflow::packet_capture(args) - .await - .expect("Error on live mode"); - //net::netflow(_interface); - } - Some(("pcap", args)) => { - println!("Pcap mode"); - if args.get_flag("list") { - println!("List of interfaces"); - for (i, interface) in interfaces.iter().enumerate() { - println!("[{}]: {}", i, interface.name); - } - - exit(0); - } - - let pcap = args - .get_one::("pcap") - .ok_or("Required output pcap file name") - .unwrap(); - let interface = args - .get_one::("interface") - .ok_or("Required Interface") - .unwrap(); - let duration = args.get_one::("duration").expect("default"); - let duration: u64 = duration.parse().unwrap(); - let interval = args.get_one::("interval").expect("default"); - let interval: u64 = interval.parse().unwrap(); - let sleep_windows = args.get_one::("sleep_windows").expect("default"); - let sleep_windows: u64 = sleep_windows.parse().unwrap(); - let verbose = args.get_one::("verbose").expect("default"); - let verbose: u8 = verbose.parse().unwrap(); - - let args: types::Args = types::Args::new( - Some(interface.to_string()), - types::Files::new(None, None, Some(pcap.to_string())), - types::Parameters::new( - None, - None, - Some(duration), - Some(interval), - Some(sleep_windows), - ), - Some(verbose), - ); - if verbose >= 1 { - println!("Interface {interface} selected"); - } - - net::pcap_capture(args).await; - } - _ => { - println!("No mode selected"); - exit(0); + // Match occures from the CLI side, which make this unreachable + _ => unreachable!() } + } else { + exit(0); } } From 88fdcd4d5b0c9f1dcde6f6a80e0dcc9073a1f469 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 09:48:39 +0900 Subject: [PATCH 07/44] feat: implement logger to `main.rs` --- src/logger.rs | 55 ++++++++++++++++++++------------------------------- src/main.rs | 14 +++++++++++++ 2 files changed, 35 insertions(+), 34 deletions(-) diff --git a/src/logger.rs b/src/logger.rs index cb14ba4..3762ebd 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -1,5 +1,5 @@ use std::fs::File; -use std::io::{stderr, stdout, Write}; +use std::io::Write; use std::path::PathBuf; use chrono::Local; // Import the Local struct from the chrono crate @@ -11,14 +11,14 @@ pub enum Logstdout { } pub struct Logger { - write_to_file: bool, - write_to_std: Option, - severity: Level, - file: Option, + pub write_to_file: bool, + pub write_to_std: Option, + pub severity: Level, + pub file: Option, } impl Logger { - pub fn new(write_to_file: bool, file_path: Option) -> 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( @@ -40,13 +40,20 @@ impl Logger { )); } let mut file = None; + + // check if there is a file at the path and create it if it doesn't exist + 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()); } Logger { - write_to_file: true, - write_to_std: None, - severity: Level::Info, + write_to_file: false, + write_to_std, + severity: severity.unwrap_or(Level::Info), file, } } @@ -64,41 +71,21 @@ impl Log for Logger { fn log(&self, record: &Record) { let timestamp = Local::now(); - + 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 => { - writeln!( - stdout(), - "[{}]: {}: {}", - timestamp, - record.level(), - record.args() - ) - .unwrap(); + println!("{}", formatted_message); } Logstdout::StdErr => { - writeln!( - stderr(), - "[{}]: {}: {}", - timestamp, - record.level(), - record.args() - ) - .unwrap(); + eprintln!("{}", formatted_message); } } } if self.write_to_file { - writeln!( - self.file.as_ref().unwrap(), - "[{}]: {}: {}", - timestamp, - record.level(), - record.args() - ) - .unwrap(); + writeln!(self.file.as_ref().unwrap(), "{}", formatted_message).unwrap(); } } diff --git a/src/main.rs b/src/main.rs index f01dd59..869697b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -58,6 +58,19 @@ enum Mode { Live, Pcap, } +impl From<&str> for Mode { + fn from(s: &str) -> Self { + match s { + "offline" => Mode::Offline, + "online" => Mode::Online, + "live" => Mode::Live, + "pcap" => Mode::Pcap, + + // Match occures from the CLI side, which make this unreachable + _ => unreachable!(), + } + } +} impl Display for Mode { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { @@ -75,6 +88,7 @@ struct Fluere { mode: Mode, logger: Logger, verbose: Level, + // interfaces: Vec, } impl Fluere { fn new( From 355993f465e51a65b9d7a63b79bd1590b3cf00ef Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 09:49:17 +0900 Subject: [PATCH 08/44] feat: overhaul interface, error --- src/net/interface.rs | 7 ++----- src/net/mod.rs | 36 +++++++++++++++++++++++++++++++++--- 2 files changed, 35 insertions(+), 8 deletions(-) diff --git a/src/net/interface.rs b/src/net/interface.rs index 6cac518..1f56a39 100644 --- a/src/net/interface.rs +++ b/src/net/interface.rs @@ -1,7 +1,7 @@ -extern crate pcap; +use std::time::Instant; + use pcap::Device; use pnet::datalink::{self, NetworkInterface}; -use std::time::Instant; pub fn get_interface(device_name: &str) -> Device { let start = Instant::now(); @@ -62,9 +62,6 @@ pub fn get_default_interface_name(interfaces: &[NetworkInterface]) -> String { .clone() }*/ -pub fn list_interfaces() -> Vec { - datalink::interfaces() -} pub fn list_interface_names() { let interfaces = Device::list(); /*let mut interface_names: Vec = Vec::new(); diff --git a/src/net/mod.rs b/src/net/mod.rs index c15a9e0..6a94278 100644 --- a/src/net/mod.rs +++ b/src/net/mod.rs @@ -1,6 +1,6 @@ +//mod fluereflow +pub mod capture; pub mod errors; -//mod fluereflow; -mod capture; mod flows; // mod interface; pub mod live_fluereflow; @@ -12,7 +12,6 @@ pub mod types; //pub use flows::packet_capture; pub use capture::find_device; -pub use capture::list_devices; pub use capture::CaptureDevice; pub use capture::DeviceError; // pub use interface::list_interface_names; @@ -20,3 +19,34 @@ pub use capture::DeviceError; pub use offline_fluereflows::fluereflow_fileparse; pub use packet_pcap::pcap_capture; //pub use types::FluereRecord; + +use std::fmt::{Display, Formatter, Result as FmtResult}; + +use pcap::Error; + +#[derive(Debug)] +pub enum NetError { + DeviceError(DeviceError), + PcapError(Error), +} + +impl From for NetError { + fn from(err: DeviceError) -> Self { + NetError::DeviceError(err) + } +} + +impl From for NetError { + fn from(err: Error) -> Self { + NetError::PcapError(err) + } +} + +impl Display for NetError { + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { + match self { + NetError::DeviceError(err) => err.fmt(f), + NetError::PcapError(err) => err.fmt(f), + } + } +} From 3d82693e46bd3505bf90c8afff380ed0e0835628 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 14:46:44 +0900 Subject: [PATCH 09/44] refator: refactor cli to match the log level of the log crate --- src/cli.rs | 62 ++++++++++++++++++++++++++++-------------------------- 1 file changed, 32 insertions(+), 30 deletions(-) diff --git a/src/cli.rs b/src/cli.rs index 145eaff..a683ad0 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -75,9 +75,9 @@ pub fn cli_template() -> Command { .arg( Arg::new("verbose") .help("Set verbosity level") - .default_value("1") + .default_value("2") .short('v') - .long("verbose"), // 0: quiet, 1: normal,2: extended, 3: verbose + .long("verbose"), // 0: Error, 1: Warning, 2: Info, 3: Debug 4: Trace ), ) .subcommand( @@ -114,9 +114,9 @@ pub fn cli_template() -> Command { .arg( Arg::new("verbose") .help("Set verbosity level") - .default_value("1") + .default_value("2") .short('v') - .long("verbose"), // 0: quiet, 1: normal,2: extended, 3: verbose + .long("verbose"), // 0: Error, 1: Warning, 2: Info, 3: Debug 4: Trace ), ) .subcommand( @@ -181,9 +181,9 @@ pub fn cli_template() -> Command { .arg( Arg::new("verbose") .help("Set verbosity level") - .default_value("1") + .default_value("2") .short('v') - .long("verbose"), // 0: quiet, 1: normal,2: extended, 3: verbose + .long("verbose"), // 0: Error, 1: Warning, 2: Info, 3: Debug 4: Trace ), ) .subcommand( @@ -234,15 +234,15 @@ pub fn cli_template() -> Command { .arg( Arg::new("verbose") .help("Set verbosity level") - .default_value("1") + .default_value("2") .short('v') - .long("verbose"), // 0: quiet, 1: normal,2: extended, 3: verbose + .long("verbose"), // 0: Error, 1: Warn, 2: Info, 3: Debug, 4: Trace ), ) } -pub async fn handle_mode(mode: &str, args: &ArgMatches) -> Args { - let _verbose = args +pub async fn handle_mode(mode: &str, args: &ArgMatches) -> (Args, u8) { + let verbose = args .get_one::("verbose") .map_or(0, |v| v.parse::().unwrap_or(0)); if args.get_flag("list") { @@ -254,12 +254,14 @@ pub async fn handle_mode(mode: &str, args: &ArgMatches) -> Args { exit(0); } - 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), _ => unreachable!(), - } + }; + + (arg_data, verbose) } fn parse_online_live_args(args: &clap::ArgMatches, _mode: &str) -> Args { @@ -292,11 +294,11 @@ fn parse_online_live_args(args: &clap::ArgMatches, _mode: &str) -> Args { .unwrap() .parse::() .unwrap(); - let verbose = args - .get_one::("verbose") - .unwrap() - .parse::() - .unwrap(); + // let verbose = args + // .get_one::("verbose") + // .unwrap() + // .parse::() + // .unwrap(); Args::new( Some(interface), @@ -308,7 +310,7 @@ fn parse_online_live_args(args: &clap::ArgMatches, _mode: &str) -> Args { Some(interval), Some(sleep_windows), ), - Some(verbose), + // Some(verbose), ) } fn parse_offline_args(args: &clap::ArgMatches) -> Args { @@ -323,17 +325,17 @@ fn parse_offline_args(args: &clap::ArgMatches) -> Args { .unwrap() .parse::() .unwrap(); - let verbose = args - .get_one::("verbose") - .unwrap() - .parse::() - .unwrap(); + // let verbose = args + // .get_one::("verbose") + // .unwrap() + // .parse::() + // .unwrap(); Args::new( None, Files::new(Some(csv), Some(file), None), Parameters::new(Some(use_mac), Some(timeout), None, None, None), - Some(verbose), + // Some(verbose), ) } fn parse_pcap_args(args: &clap::ArgMatches) -> Args { @@ -360,11 +362,11 @@ fn parse_pcap_args(args: &clap::ArgMatches) -> Args { .unwrap() .parse::() .unwrap(); - let verbose = args - .get_one::("verbose") - .unwrap() - .parse::() - .unwrap(); + // let verbose = args + // .get_one::("verbose") + // .unwrap() + // .parse::() + // .unwrap(); Args::new( Some(interface), @@ -376,6 +378,6 @@ fn parse_pcap_args(args: &clap::ArgMatches) -> Args { Some(interval), Some(sleep_windows), ), - Some(verbose), + // Some(verbose), ) } From 2e91ca4da772b6c5a22743b3b463a715bf503f95 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 14:48:43 +0900 Subject: [PATCH 10/44] feat: set time format for logger --- src/logger.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/logger.rs b/src/logger.rs index 3762ebd..d8adf84 100644 --- a/src/logger.rs +++ b/src/logger.rs @@ -70,8 +70,9 @@ impl Log for Logger { } fn log(&self, record: &Record) { - let timestamp = Local::now(); - let formatted_message = format!("{} [{}]: {}", timestamp, record.level(), record.args()); + // 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() { From cdd806e09b9a69e2c27e9a3df8505d2fb9e4e953 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 14:49:09 +0900 Subject: [PATCH 11/44] fix: remove `String`, introduce `Cow` --- src/net/capture.rs | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/src/net/capture.rs b/src/net/capture.rs index a7b424a..ed168ce 100644 --- a/src/net/capture.rs +++ b/src/net/capture.rs @@ -1,8 +1,9 @@ -use std::{fmt, time::Instant}; +use std::{fmt, time::Instant, borrow::Cow}; use crate::net::NetError; use pcap::{Active, Address, Capture, Device, Error as PcapError}; +use log::{info, debug}; #[derive(Debug)] pub enum DeviceError { @@ -20,8 +21,8 @@ impl fmt::Display for DeviceError { } pub struct CaptureDevice { - pub name: String, - pub desc: String, + pub name: Cow<'static, str>, + pub desc: Cow<'static, str>, pub address: Vec
, pub capture: Capture, } @@ -29,10 +30,12 @@ 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 desc: Cow<'static, str> = Cow::Owned(device.desc.unwrap_or("".to_string())); Ok(CaptureDevice { - name: device.name, - desc: device.desc.as_deref().unwrap_or("").to_string(), + name, + desc, address: device.addresses, capture, }) @@ -41,20 +44,21 @@ impl CaptureDevice { impl Drop for CaptureDevice { fn drop(&mut self) { - println!("Closing capture session for device {}", self.name); + info!("Closing capture session for device {}", self.name); + // println!("Closing capture session for device {}", self.name); // self.capture.; } } pub fn find_device(identifier: &str) -> Result { let start = Instant::now(); - println!("Requested Device: {}", identifier); + debug!("Looking for device: {}", identifier); let devices = Device::list()?; if let Ok(index) = identifier.parse::() { if let Some(device) = devices.get(index) { let duration = start.elapsed(); - println!("Device {} captured in {:?}", device.name, duration); + info!("Device {} captured in {:?}", device.name, duration); return Ok(device.clone()); } else { return Err(NetError::DeviceError(DeviceError::InvalidDeviceIndex( @@ -66,7 +70,7 @@ pub fn find_device(identifier: &str) -> Result { for device in devices { if device.name == identifier { let duration = start.elapsed(); - println!("{} device captured in {:?}", identifier, duration); + info!("Device {} captured in {:?}", device.name, duration); return Ok(device); } } @@ -77,6 +81,7 @@ 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)? .promisc(true) .snaplen(1024) From ee592cb29a916ff0d647ca2b0143d847f4a0b2c0 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 14:49:48 +0900 Subject: [PATCH 12/44] feat: introduce logger usage to `live` collect --- src/net/live_fluereflow.rs | 111 +++++++++++++++---------------------- 1 file changed, 44 insertions(+), 67 deletions(-) diff --git a/src/net/live_fluereflow.rs b/src/net/live_fluereflow.rs index 182062f..8ae6692 100644 --- a/src/net/live_fluereflow.rs +++ b/src/net/live_fluereflow.rs @@ -1,18 +1,34 @@ // 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. -extern crate csv; - -use crossterm::{ - event::{self, DisableMouseCapture, EnableMouseCapture, KeyCode, KeyEvent}, - execute, - terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen}, +use std::{ + collections::HashMap, + fs, + io, + sync::Arc, + time::{Duration, Instant, SystemTime}, +}; +use crate::{ + net::{ + find_device, + flows::update_flow, + parser::{microseconds_to_timestamp, parse_fluereflow, parse_keys, parse_microseconds}, + types::TcpFlags, + CaptureDevice, + }, + types::{Args, UDFlowKey}, + utils::{cur_time_file, fluere_exporter}, }; use fluere_config::Config; use fluere_plugin::PluginManager; use fluereflow::FluereRecord; +use crossterm::{ + event::{self, DisableMouseCapture, EnableMouseCapture, KeyCode, KeyEvent}, + execute, + terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen}, +}; use ratatui::{ backend::CrosstermBackend, layout::{Constraint, Direction, Layout}, @@ -20,28 +36,13 @@ use ratatui::{ widgets::{Block, Borders, Gauge, List, ListItem, Paragraph}, Frame, Terminal, }; -use tokio::sync::Mutex; -use tokio::task; -use tokio::time::sleep; - -use crate::{ - net::{ - find_device, - flows::update_flow, - parser::{microseconds_to_timestamp, parse_fluereflow, parse_keys, parse_microseconds}, - types::TcpFlags, - CaptureDevice, - }, - types::{Args, UDFlowKey}, - utils::{cur_time_file, fluere_exporter}, +use tokio::{ + sync::Mutex, + task, + time::sleep, }; +use log::{info, debug, trace}; -use std::{ - collections::HashMap, - fs, io, - sync::Arc, - time::{Duration, Instant, SystemTime}, -}; const MAX_RECENT_FLOWS: usize = 50; @@ -49,9 +50,10 @@ const MAX_RECENT_FLOWS: usize = 50; // It takes the command line arguments as input and calls the online_packet_capture function. // It returns a Result indicating whether the operation was successful. pub async fn packet_capture(arg: Args) -> Result<(), io::Error> { - println!("TUI"); + debug!("Starting Terminal User Interface"); online_packet_capture(arg).await; + debug!("Terminal User Interface Stopped"); Ok(()) } #[derive(Debug, Clone)] @@ -74,7 +76,6 @@ pub async fn online_packet_capture(arg: Args) { let interval = arg.parameters.interval.unwrap(); let flow_timeout = arg.parameters.timeout.unwrap(); let sleep_windows = arg.parameters.sleep_windows.unwrap(); - let verbose = arg.verbose.unwrap(); let config = Config::new(); let plugin_manager = PluginManager::new().expect("Failed to create plugin manager"); let plugin_worker = plugin_manager.start_worker(); @@ -90,11 +91,7 @@ pub async fn online_packet_capture(arg: Args) { let file_dir = "./output"; match fs::create_dir_all(<&str>::clone(&file_dir)) { - Ok(_) => { - if verbose >= 1 { - println!("Created directory: {}", file_dir) - } - } + Ok(_) => debug!("Created directory: {}", file_dir), Err(error) => panic!("Problem creating directory: {:?}", error), }; @@ -179,9 +176,7 @@ pub async fn online_packet_capture(arg: Args) { continue; } Ok(packet) => { - if verbose >= 3 { - println!("received packet"); - } + trace!("received packet"); let (mut key_value, mut reverse_key) = match parse_keys(packet.clone()) { Ok(keys) => keys, @@ -209,9 +204,7 @@ pub async fn online_packet_capture(arg: Args) { if flowdata.prot == 6 { if flags.syn > 0 { active_flow_guard.insert(key_value, flowdata); - if verbose >= 2 { - println!("flow established"); - } + trace!("flow established"); let mut recent_flows_guard = recent_flows.lock().await; recent_flows_guard.push(FlowSummary { src: key_value.src_ip.to_string(), @@ -228,9 +221,7 @@ pub async fn online_packet_capture(arg: Args) { } } else { active_flow_guard.insert(key_value, flowdata); - if verbose >= 2 { - println!("flow established"); - } + trace!("flow established"); let mut recent_flows_guard = recent_flows.lock().await; recent_flows_guard.push(FlowSummary { src: key_value.src_ip.to_string(), @@ -270,17 +261,13 @@ pub async fn online_packet_capture(arg: Args) { }; update_flow(flow, is_reverse, update_key); - if verbose >= 2 { - println!( + trace!( "{} flow updated", if is_reverse { "reverse" } else { "forward" } ); - } if flags.fin == 1 || flags.rst == 1 { - if verbose >= 2 { - println!("flow finished"); - } + trace!("flow finished"); plugin_manager.process_flow_data(*flow).await.unwrap(); records.push(*flow); active_flow_guard.remove(flow_key); @@ -289,12 +276,10 @@ pub async fn online_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 packet_count % sleep_windows == 0 && cfg!(target_os = "windows") { + // println!("Slow down the loop for windows"); + // sleep(Duration::from_millis(0)).await; + // } // Export flows if the interval has been reached let mut last_export_guard = last_export.lock().await; @@ -304,9 +289,7 @@ pub async fn online_packet_capture(arg: Args) { packet_count = 0; for (key, flow) in active_flow_guard.iter() { if flow_timeout > 0 && flow.last < (time - (flow_timeout * 1000)) { - if verbose >= 2 { - println!("flow expired"); - } + trace!("flow expired"); plugin_manager.process_flow_data(*flow).await.unwrap(); records.push(*flow); expired_flows.push(*key); @@ -314,7 +297,7 @@ pub async fn online_packet_capture(arg: Args) { } active_flow_guard.retain(|key, _| !expired_flows.contains(key)); let cloned_records = records.clone(); - println!("{} flows exported", records.len()); + debug!("{} flows exported", records.len()); records.clear(); let tasks = task::spawn(async { fluere_exporter(cloned_records, file); @@ -338,9 +321,7 @@ pub async fn online_packet_capture(arg: Args) { let mut expired_flows = vec![]; for (key, flow) in active_flow_guard.iter() { if flow.last < (time - (flow_timeout * 1000)) { - if verbose >= 2 { - println!("flow expired"); - } + trace!("flow expired"); plugin_manager.process_flow_data(*flow).await.unwrap(); records.push(*flow); expired_flows.push(*key); @@ -352,9 +333,7 @@ pub async fn online_packet_capture(arg: Args) { } } } - if verbose >= 1 { - println!("Captured in {:?}", start.elapsed()); - } + debug!("Captured in {:?}", start.elapsed()); let active_flow_guard = active_flow.lock().await; for (_key, flow) in active_flow_guard.iter() { @@ -367,9 +346,7 @@ pub async fn online_packet_capture(arg: Args) { }); let result = tasks.await; - if verbose >= 1 { - println!("Exporting task excutation result: {:?}", result); - } + debug!("Exporting task excutation result: {:?}", result); let _ = draw_task.await; disable_raw_mode().expect("failed to disable raw mode"); let mut terminal_guard = terminal.lock().await; From 3fb5b477063321a9c66fcc210f705d07d5dd4228 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 14:50:08 +0900 Subject: [PATCH 13/44] feat: introduce logger to `offline` conversion --- src/net/offline_fluereflows.rs | 52 ++++++++++++++-------------------- 1 file changed, 21 insertions(+), 31 deletions(-) diff --git a/src/net/offline_fluereflows.rs b/src/net/offline_fluereflows.rs index e852678..ea1a642 100644 --- a/src/net/offline_fluereflows.rs +++ b/src/net/offline_fluereflows.rs @@ -1,8 +1,4 @@ -extern crate csv; - -use fluereflow::FluereRecord; -use pcap::Capture; -use tokio::task; +use std::{collections::HashMap, fs, time::Instant}; use crate::{ net::{ @@ -14,23 +10,23 @@ use crate::{ utils::{cur_time_file, fluere_exporter}, }; -use std::{collections::HashMap, fs, time::Instant}; +use fluereflow::FluereRecord; +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(); let file_name = arg.files.file.unwrap(); let use_mac = arg.parameters.use_mac.unwrap(); let _flow_timeout = arg.parameters.timeout.unwrap(); - let verbose = arg.verbose.unwrap(); let mut cap = Capture::from_file(file_name).unwrap(); let file_dir = "./output"; match fs::create_dir_all(<&str>::clone(&file_dir)) { Ok(_) => { - if verbose >= 1 { - println!("Created directory: {}", file_dir) - } + debug!("Created directory: {}", file_dir) } Err(error) => panic!("Problem creating directory: {:?}", error), }; @@ -66,17 +62,16 @@ pub async fn fluereflow_fileparse(arg: Args) { if flowdata.prot == 6 { if flags.syn > 0 { active_flow.insert(key_value, flowdata); - if verbose >= 2 { - println!("flow established"); - } + + trace!("flow established"); + } else { continue; } } else { active_flow.insert(key_value, flowdata); - if verbose >= 2 { - println!("flow established"); - } + + trace!("flow established"); } false @@ -120,27 +115,22 @@ pub async fn fluereflow_fileparse(arg: Args) { }; update_flow(flow, is_reverse, update_key); - if verbose >= 2 { - println!( + trace!( "{} flow updated", if is_reverse { "reverse" } else { "forward" } ); - } if flags.fin == 1 || flags.rst == 1 { - if verbose >= 2 { - println!("flow finished"); - } + trace!("flow finished"); records.push(*flow); active_flow.remove(flow_key); } } } - if verbose >= 1 { - println!("Captured in {:?}", start.elapsed()); - } - println!("Active flow {:?}", active_flow.len()); - println!("Ended flow {:?}", records.len()); + 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); } @@ -149,8 +139,8 @@ pub async fn fluereflow_fileparse(arg: Args) { }); let result = tasks.await; - if verbose >= 1 { - println!("Export {} result: {:?}", file_path, result); - } - //println!("records {:?}", records); + info!("Export {} result: {:?}", file_path, result); + + println!("Active flow {:?}", ac_flow_cnt); + println!("Ended flow {:?}", ended_flow_cnt); } From fe44fb32a3e9c4a3494265e5b1b78d7108cd9172 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 14:50:26 +0900 Subject: [PATCH 14/44] feat: introduce logger to `online` collect --- src/net/online_fluereflow.rs | 79 +++++++++++------------------------- 1 file changed, 24 insertions(+), 55 deletions(-) diff --git a/src/net/online_fluereflow.rs b/src/net/online_fluereflow.rs index da4cfc1..0bbb4d7 100644 --- a/src/net/online_fluereflow.rs +++ b/src/net/online_fluereflow.rs @@ -8,7 +8,8 @@ use fluere_plugin::PluginManager; use fluereflow::FluereRecord; use tokio::task; -use tokio::time::sleep; + +use log::{info, debug, trace}; use crate::{ net::{ @@ -38,8 +39,7 @@ pub async fn 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 verbose = arg.verbose.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(); @@ -52,21 +52,11 @@ pub async fn packet_capture(arg: Args) { let interface = find_device(interface_name.as_str()).unwrap(); let mut cap_device = CaptureDevice::new(interface.clone()).unwrap(); let cap = &mut 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)) { Ok(_) => { - if verbose >= 1 { - println!("Created directory: {}", file_dir) - } + trace!("Created directory: {}", file_dir) } Err(error) => panic!("Problem creating directory: {:?}", error), }; @@ -88,9 +78,7 @@ pub async fn packet_capture(arg: Args) { continue; } Ok(packet) => { - if verbose >= 3 { - println!("received packet"); - } + trace!("received packet"); let (mut key_value, mut reverse_key) = match parse_keys(packet.clone()) { Ok(keys) => keys, @@ -115,19 +103,14 @@ pub async fn packet_capture(arg: Args) { if flowdata.prot == 6 { if flags.syn > 0 { active_flow.insert(key_value, flowdata); - if verbose >= 2 { - println!("flow established"); - } + trace!("flow established"); } else { continue; } } else { active_flow.insert(key_value, flowdata); - if verbose >= 2 { - println!("flow established"); - } + trace!("flow established"); } - false } Some(_) => true, @@ -169,17 +152,14 @@ pub async fn packet_capture(arg: Args) { }; update_flow(flow, is_reverse, update_key); - if verbose >= 2 { - println!( - "{} flow updated", - if is_reverse { "reverse" } else { "forward" } - ); - } + trace!( + "{} flow updated", + if is_reverse { "reverse" } else { "forward" } + ); + if flags.fin == 1 || flags.rst == 1 { - if verbose >= 2 { - println!("flow finished"); - } + trace!("flow finished"); // plugin_manager.process_flow_data(flow).expect("Failed to process flow data"); plugin_manager.process_flow_data(*flow).await.unwrap(); @@ -191,12 +171,12 @@ 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 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; + // } // Export flows if the interval has been reached if last_export.elapsed() >= Duration::from_millis(interval) && interval != 0 { @@ -204,9 +184,7 @@ pub async fn packet_capture(arg: Args) { packet_count = 0; for (key, flow) in active_flow.iter() { if flow_timeout > 0 && flow.last < (time - (flow_timeout * 1000)) { - if verbose >= 2 { - println!("flow expired"); - } + trace!("flow expired"); plugin_manager.process_flow_data(*flow).await.unwrap(); records.push(*flow); expired_flows.push(*key); @@ -223,9 +201,7 @@ pub async fn packet_capture(arg: Args) { }); let result = tasks.await; - if verbose >= 1 { - println!("Export {} result: {:?}", file_path, result); - } + info!("Export {} result: {:?}", file_path, result); file_path = cur_time_file(csv_file.as_str(), file_dir, ".csv").await; file = fs::File::create(file_path.clone()).unwrap(); last_export = Instant::now(); @@ -236,9 +212,7 @@ pub async fn packet_capture(arg: Args) { let mut expired_flows = vec![]; for (key, flow) in active_flow.iter() { if flow.last < (time - (flow_timeout * 1000)) { - if verbose >= 2 { - println!("flow expired"); - } + trace!("flow expired"); plugin_manager.process_flow_data(*flow).await.unwrap(); records.push(*flow); expired_flows.push(*key); @@ -250,9 +224,7 @@ pub async fn packet_capture(arg: Args) { } } } - if verbose >= 1 { - println!("Captured in {:?}", start.elapsed()); - } + debug!("Captured in {:?}", start.elapsed()); for (_key, flow) in active_flow.iter() { plugin_manager.process_flow_data(*flow).await.unwrap(); records.push(*flow); @@ -264,8 +236,5 @@ pub async fn packet_capture(arg: Args) { plugin_manager.await_completion(plugin_worker).await; let result = tasks.await; - if verbose >= 1 { - println!("Exporting task excutation result: {:?}", result); - } - //println!("records {:?}", records); + info!("Exporting task excutation result: {:?}", result); } From 8d51cfc57e258db897418680ca27823e36b33177 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 14:50:47 +0900 Subject: [PATCH 15/44] feat: introduce logger to `packet` collect --- src/net/packet_pcap.rs | 34 +++++++++++----------------------- 1 file changed, 11 insertions(+), 23 deletions(-) diff --git a/src/net/packet_pcap.rs b/src/net/packet_pcap.rs index d4446bf..750f180 100644 --- a/src/net/packet_pcap.rs +++ b/src/net/packet_pcap.rs @@ -1,15 +1,14 @@ -extern crate chrono; - use std::fs; - -use tokio::time::sleep; +use std::time::{Duration, Instant}; use crate::net::find_device; use crate::net::CaptureDevice; use crate::types::Args; use crate::utils::cur_time_file; -use std::time::{Duration, Instant}; +use tokio::time::sleep; +use log::{info, debug, trace}; + pub async fn pcap_capture(args: Args) { let pcap_file = args.files.pcap.unwrap(); @@ -17,7 +16,6 @@ pub async fn pcap_capture(args: Args) { let duration = args.parameters.duration.unwrap(); let _interval = args.parameters.interval.unwrap(); let sleep_windows = args.parameters.sleep_windows.unwrap(); - let verbose = args.verbose.unwrap(); let interface = find_device(interface_name.as_str()).unwrap(); let mut cap_device = CaptureDevice::new(interface.clone()).unwrap(); @@ -26,11 +24,7 @@ pub async fn pcap_capture(args: Args) { let file_dir = "./output"; let mut packet_count = 0; match fs::create_dir_all(<&str>::clone(&file_dir)) { - Ok(_) => { - if verbose >= 1 { - println!("Created directory: {}", file_dir) - } - } + Ok(_) => debug!("Created directory: {}", file_dir), Err(error) => panic!("Problem creating directory: {:?}", error), }; @@ -42,27 +36,21 @@ pub async fn pcap_capture(args: Args) { let start = Instant::now(); while let Ok(packet) = cap.next_packet() { - if verbose >= 3 { - println!("received packet"); - } + trace!("received packet"); //println!("packet: {:?}", packet); file.write(&packet); 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 packet_count % sleep_windows == 0 && cfg!(target_os = "windows") { + // println!("Slow down the loop for windows"); + // sleep(Duration::from_millis(0)).await; + // } // Check if the duration has been reached if start.elapsed() >= Duration::from_millis(duration) && duration != 0 { break; } } - if verbose >= 1 { - println!("Captured in {:?}", start.elapsed()); - } + debug!("Captured in {:?}", start.elapsed()); } From 88d01390c652ae05ba0770ddeb4931d7b9174c1a Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 14:51:19 +0900 Subject: [PATCH 16/44] refactor: remove `verbose` field from `Argument` --- src/types/argument.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/types/argument.rs b/src/types/argument.rs index e81b251..58ba8af 100644 --- a/src/types/argument.rs +++ b/src/types/argument.rs @@ -3,7 +3,7 @@ pub struct Args { pub interface: Option, pub files: Files, pub parameters: Parameters, - pub verbose: Option, + // pub verbose: Option, } impl Args { @@ -11,13 +11,13 @@ impl Args { interface: Option, files: Files, parameters: Parameters, - verbose: Option, + // verbose: Option, ) -> Self { Self { interface, files, parameters, - verbose, + // verbose, } } } From f758f3c04eaec98f226215b0b288d662c47b01c1 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 14:52:00 +0900 Subject: [PATCH 17/44] feat: introduce logger to csv exporter --- src/utils/fluere_csv_exporter.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/utils/fluere_csv_exporter.rs b/src/utils/fluere_csv_exporter.rs index 927f049..0f86bf2 100644 --- a/src/utils/fluere_csv_exporter.rs +++ b/src/utils/fluere_csv_exporter.rs @@ -1,11 +1,12 @@ use fluereflow::FluereRecord; +use log::{ debug, trace}; use std::fs::File; pub fn fluere_exporter(records: Vec, file: File) { let mut wtr = csv::Writer::from_writer(file); - // println!("Writing {} records", records.len()); - // println!(" record: {:?}", records); + debug!("Writing {} records", records.len()); + trace!(" record: {:?}", records); wtr.write_record([ "source", "destination", @@ -68,4 +69,5 @@ pub fn fluere_exporter(records: Vec, file: File) { ]) .unwrap(); } + debug!("Wrote {} records", records.len()); } From 4555aec36e0dfb58681fd9b7a65136efdaaa927c Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 14:52:31 +0900 Subject: [PATCH 18/44] feat: implement log filter level conversion --- src/main.rs | 65 ++++++++++++++++++++++++++++++++--------------------- 1 file changed, 39 insertions(+), 26 deletions(-) diff --git a/src/main.rs b/src/main.rs index 869697b..78d5fc6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -17,7 +17,7 @@ use crate::logger::{Logger, Logstdout}; use crate::net::capture::DeviceError; // use env_logger;::{init, Logger}; -use log::{Level, Log, info, warn, error, debug, trace}; +use log::{Level, LevelFilter, info, debug}; // FEAT:MAYBE: seprate `std` as feature flag for fluere and log crate @@ -82,6 +82,18 @@ 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!(), + } + } + struct Fluere { interface: String, args: types::Args, @@ -116,13 +128,7 @@ impl Fluere { #[tokio::main] async fn main() { let args = cli::cli_template().get_matches(); - let log_stdout = Logstdout::Stdout; - let log_file :Option = None; - let log_level = Level::Info; - let logger = Logger::new(None,Some(Level::Trace), Some(Logstdout::Stdout),false); - let _ = log::set_boxed_logger(Box::new(logger)) - .map(|()| log::set_max_level(log::LevelFilter::Info)); // let mode = match args.subcommand() { // Some((mode, _sub_args)) => mode, // None => { @@ -131,26 +137,33 @@ async fn main() { // } // }; - info!("Fluere started"); + if let Some((mode, sub_args)) = args.subcommand() { - match mode { - "online" | "offline" | "live" | "pcap" => { - log::debug!("Mode: {}", mode); - let parems = cli::handle_mode(mode, sub_args).await; - - match mode { - "online" => net::online_fluereflow::packet_capture(parems).await, - "offline" => net::fluereflow_fileparse(parems).await, - "live" => net::live_fluereflow::packet_capture(parems) - .await - .expect("Error on live mode"), - "pcap" => net::pcap_capture(parems).await, - _ => unreachable!(), - } - } - - // Match occures from the CLI side, which make this unreachable - _ => unreachable!() + 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_level = Level::Info; + 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"); + + match mode_type { + Mode::Online => net::online_fluereflow::packet_capture(parems.0).await, + Mode::Offline => net::fluereflow_fileparse(parems.0).await, + Mode::Live => net::live_fluereflow::packet_capture(parems.0) + .await + .expect("Error on live mode"), + Mode::Pcap => net::pcap_capture(parems.0).await, + _ => unreachable!(), } } else { exit(0); From 45323a99d2bb6a98b7cd9dabc1ab20cfc32ff283 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 15:12:51 +0900 Subject: [PATCH 19/44] 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) { From 334a4f70a63aadae1a4f4f083bbc79410259d423 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 16:37:51 +0900 Subject: [PATCH 20/44] feat(plugin): implement `log` as a feature --- Cargo.toml | 2 +- fluere-plugin/Cargo.toml | 6 +++ fluere-plugin/src/downloader.rs | 23 +++++++++- fluere-plugin/src/lib.rs | 76 +++++++++++++++++++++++++++------ 4 files changed, 91 insertions(+), 16 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index bf6b31a..bb08c49 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,7 +30,7 @@ snafu = "0.8.0" serde = { version = "1.0.193", features = ["derive"] } toml = "0.8.8" -fluere_plugin = { version = "0.1.1", path = "./fluere-plugin" } +fluere_plugin = { version = "0.1.1", path = "./fluere-plugin", features = ["log"] } fluere-config = { version = "0.1.3", path = "./fluere-config" } fluereflow = { version = "0.3.2", path = "./fluereflow" } diff --git a/fluere-plugin/Cargo.toml b/fluere-plugin/Cargo.toml index 4579ca8..c1acc83 100644 --- a/fluere-plugin/Cargo.toml +++ b/fluere-plugin/Cargo.toml @@ -23,5 +23,11 @@ dirs = "5.0.1" mlua = { version = "0.9.2", features = ["lua54", "vendored","async","send"] } inksac = "0.4.0" +log = { version = "0.4.21", features = ["std"], optional = true } + [lib] name = "fluere_plugin" + + +[features] +log = ["dep:log"] \ No newline at end of file diff --git a/fluere-plugin/src/downloader.rs b/fluere-plugin/src/downloader.rs index 3833a19..db218a5 100644 --- a/fluere-plugin/src/downloader.rs +++ b/fluere-plugin/src/downloader.rs @@ -12,6 +12,9 @@ use std::time::Duration; use crate::util::home_cache_path; +#[cfg(feature = "log")] +use log::{debug, info, trace, warn, error}; + #[derive(Debug)] pub enum DownloadError { Io(std::io::Error), @@ -81,7 +84,11 @@ pub fn download_plugin_from_github(repo_name: &str) -> Result<(), DownloadError> if has_local_changes(&repo)? { println!("{}: You have uncommitted changes. Updating will overwrite these changes. Continue? [y/N] (auto skip in 5 seconds)","Warning".styled(warn_style)); if !user_confirms()? { + #[cfg(feature = "log")] + info!("Update skipped for {}", repo_name.styled(highlight_style)); + #[cfg(not(feature = "log"))] println!("Update skipped for {}", repo_name.styled(highlight_style)); + return Ok(()); } } @@ -96,14 +103,26 @@ pub fn download_plugin_from_github(repo_name: &str) -> Result<(), DownloadError> repo.checkout_tree(fetch_commit_obj.as_object(), Some(&mut checkout_builder))?; repo.set_head_detached(fetch_commit)?; + #[cfg(feature = "log")] + info!( + "Successfully updated to the latest version for {}", + repo_name.styled(highlight_style) + ); + #[cfg(not(feature = "log"))] println!( "Successfully updated to the latest version for {}", repo_name.styled(highlight_style) ); } else { + #[cfg(feature = "log")] + info!("Update skipped for {}", repo_name.styled(highlight_style)); + #[cfg(not(feature = "log"))] println!("Update skipped for {}", repo_name.styled(highlight_style)); } } else { + #[cfg(feature = "log")] + info!("{} is up to date.", repo_name.styled(highlight_style)); + #[cfg(not(feature = "log"))] println!("{} is up to date.", repo_name.styled(highlight_style)); } @@ -129,7 +148,9 @@ fn user_confirms() -> Result { match receiver.recv_timeout(Duration::from_secs(5)) { Ok(result) => Ok(result), Err(_) => { - print!("Timeout. "); + #[cfg(feature = "log")] + debug!("Timeout. "); + Ok(false) } } diff --git a/fluere-plugin/src/lib.rs b/fluere-plugin/src/lib.rs index 9278ece..cc01244 100644 --- a/fluere-plugin/src/lib.rs +++ b/fluere-plugin/src/lib.rs @@ -1,8 +1,3 @@ -use fluere_config::Config; -use fluereflow::FluereRecord; -use mlua::{Lua, Result}; -use tokio::sync::{mpsc, Mutex}; - use std::collections::HashSet; use std::sync::Arc; @@ -12,6 +7,14 @@ mod util; use downloader::download_plugin_from_github; use util::home_cache_path; +use fluere_config::Config; +use fluereflow::FluereRecord; +use mlua::{Lua, Result}; +use tokio::sync::{mpsc, Mutex}; + +#[cfg(feature = "log")] +use log::{debug, info, trace, warn, error}; + pub struct PluginManager { lua: Arc>, sender: mpsc::Sender, @@ -34,6 +37,9 @@ impl PluginManager { } pub async fn load_plugins(&self, config: &Config) -> Result<()> { + #[cfg(feature = "log")] + debug!("Loading plugins"); + let plugins_clone = self.plugins.clone(); let mut plugins_guard = plugins_clone.lock().await; for (name, plugin_config) in &config.plugins { @@ -50,7 +56,10 @@ impl PluginManager { let lua_clone = self.lua.clone(); let lua_guard = lua_clone.lock().await; let lua = &*lua_guard; - // println!("lua path: {}", path); + + #[cfg(feature = "log")] + debug!("Lua path: {}", path); + let lua_plugin_path = format!("package.path = package.path .. \";{}/?.lua\"", path); let _ = lua.load(lua_plugin_path).exec(); @@ -61,7 +70,9 @@ impl PluginManager { let argument_table = lua.create_table()?; - // println!("extra argument details{:?}", plugin_config.extra_arguments); + #[cfg(feature = "log")] + debug!("extra argument details{:?}", plugin_config.extra_arguments); + for (key, value) in plugin_config.extra_arguments.clone().unwrap().iter() { @@ -89,11 +100,22 @@ impl PluginManager { //lua_guard.load(&code).exec().expect(format!("Error on plugin: {}", name).as_str()); plugins_guard.insert(name.clone()); + #[cfg(feature = "log")] + info!("Loaded plugin {}", name); + #[cfg(not(feature = "log"))] println!("Loaded plugin {}", name); } Err(err) => { - println!("Failed to read plugin: {}", name); - println!("Error: {}", err); + #[cfg(feature = "log")] + { + warn!("Failed to read plugin: {}", name); + error!("Error: {}", err); + } + #[cfg(not(feature = "log"))] + { + println!("Failed to read plugin: {}", name); + println!("Error: {}", err); + } continue; } }; @@ -147,18 +169,37 @@ impl PluginManager { Ok(()) }).expect(format!("Error on plugin: {}", name).as_str());*/ plugins_guard.insert(name.clone()); + #[cfg(feature = "log")] + info!("Loaded plugin {}", name); + #[cfg(not(feature = "log"))] println!("Loaded plugin {}", name); } Err(eri) => { - println!("Failed to read plugin: {}", name); - println!("Error: {}", eri); + #[cfg(feature = "log")] + { + warn!("Failed to read plugin: {}", name); + error!("Error: {}", eri); + } + #[cfg(not(feature = "log"))] + { + println!("Failed to read plugin: {}", name); + println!("Error: {}", eri); + } continue; } } } Err(eri) => { - println!("Unable to download plugin: {}", name); - println!("Error: {}", eri); + #[cfg(feature = "log")] + { + warn!("Unable to download plugin: {}", name); + error!("Error: {}", eri); + } + #[cfg(not(feature = "log"))] + { + println!("Unable to download plugin: {}", name); + println!("Error: {}", eri); + } } } } @@ -233,10 +274,14 @@ impl PluginManager { func.call::, ()>(lua_table.clone()) .unwrap_or_else(|_| panic!("Error on plugin: {}", plugin_name)); } else { - println!( + #[cfg(feature = "log")] + error!( "'process_data' function not found in plugin: {}", plugin_name ); + + #[cfg(not(feature = "log"))] + println!("'process_data' function not found in plugin: {}", plugin_name); } } }) @@ -272,6 +317,9 @@ impl PluginManager { func.call::<(), ()>(()) .unwrap_or_else(|_| panic!("Error on plugin: {}", plugin_name)); } else { + #[cfg(feature = "log")] + warn!("cleanup function not found in plugin: {}", plugin_name); + #[cfg(not(feature = "log"))] println!("cleanup function not found in plugin: {}", plugin_name); } } From a864cf578632959f707597000e52195c75ed1844 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 16:39:41 +0900 Subject: [PATCH 21/44] style(plugin): `rustfmt` --- fluere-plugin/src/downloader.rs | 6 +++--- fluere-plugin/src/lib.rs | 11 +++++++---- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/fluere-plugin/src/downloader.rs b/fluere-plugin/src/downloader.rs index db218a5..ad3b1e6 100644 --- a/fluere-plugin/src/downloader.rs +++ b/fluere-plugin/src/downloader.rs @@ -13,7 +13,7 @@ use std::time::Duration; use crate::util::home_cache_path; #[cfg(feature = "log")] -use log::{debug, info, trace, warn, error}; +use log::{debug, info}; #[derive(Debug)] pub enum DownloadError { @@ -88,7 +88,7 @@ pub fn download_plugin_from_github(repo_name: &str) -> Result<(), DownloadError> info!("Update skipped for {}", repo_name.styled(highlight_style)); #[cfg(not(feature = "log"))] println!("Update skipped for {}", repo_name.styled(highlight_style)); - + return Ok(()); } } @@ -115,7 +115,7 @@ pub fn download_plugin_from_github(repo_name: &str) -> Result<(), DownloadError> ); } else { #[cfg(feature = "log")] - info!("Update skipped for {}", repo_name.styled(highlight_style)); + info!("Update skipped for {}", repo_name.styled(highlight_style)); #[cfg(not(feature = "log"))] println!("Update skipped for {}", repo_name.styled(highlight_style)); } diff --git a/fluere-plugin/src/lib.rs b/fluere-plugin/src/lib.rs index cc01244..39e545c 100644 --- a/fluere-plugin/src/lib.rs +++ b/fluere-plugin/src/lib.rs @@ -13,7 +13,7 @@ use mlua::{Lua, Result}; use tokio::sync::{mpsc, Mutex}; #[cfg(feature = "log")] -use log::{debug, info, trace, warn, error}; +use log::{debug, error, info, warn}; pub struct PluginManager { lua: Arc>, @@ -108,8 +108,8 @@ impl PluginManager { Err(err) => { #[cfg(feature = "log")] { - warn!("Failed to read plugin: {}", name); - error!("Error: {}", err); + warn!("Failed to read plugin: {}", name); + error!("Error: {}", err); } #[cfg(not(feature = "log"))] { @@ -281,7 +281,10 @@ impl PluginManager { ); #[cfg(not(feature = "log"))] - println!("'process_data' function not found in plugin: {}", plugin_name); + println!( + "'process_data' function not found in plugin: {}", + plugin_name + ); } } }) From 32d1571738beaccbeab2cde04a9314bb7c610031 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 16:39:58 +0900 Subject: [PATCH 22/44] build(plugin): bump fluere-plugin version --- Cargo.toml | 2 +- fluere-plugin/Cargo.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index bb08c49..2c323b1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,7 +30,7 @@ snafu = "0.8.0" serde = { version = "1.0.193", features = ["derive"] } toml = "0.8.8" -fluere_plugin = { version = "0.1.1", path = "./fluere-plugin", features = ["log"] } +fluere_plugin = { version = "0.2.0", path = "./fluere-plugin", features = ["log"] } fluere-config = { version = "0.1.3", path = "./fluere-config" } fluereflow = { version = "0.3.2", path = "./fluereflow" } diff --git a/fluere-plugin/Cargo.toml b/fluere-plugin/Cargo.toml index c1acc83..667cea2 100644 --- a/fluere-plugin/Cargo.toml +++ b/fluere-plugin/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "fluere_plugin" -version = "0.1.1" +version = "0.2.0" authors = ["Skuld Norniern "] edition = "2021" description = "Plugin API for Fluere." From 315f4b679dbadd966351f04321f2bd90f26f9d0f Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 16:48:05 +0900 Subject: [PATCH 23/44] feat(config): implement `log` as feature --- Cargo.toml | 2 +- fluere-config/Cargo.toml | 4 ++++ fluere-config/src/init.rs | 33 ++++++++++++++++++++++++++++----- 3 files changed, 33 insertions(+), 6 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 2c323b1..1f65aea 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,7 +31,7 @@ serde = { version = "1.0.193", features = ["derive"] } toml = "0.8.8" fluere_plugin = { version = "0.2.0", path = "./fluere-plugin", features = ["log"] } -fluere-config = { version = "0.1.3", path = "./fluere-config" } +fluere-config = { version = "0.1.3", path = "./fluere-config", features = ["log"] } fluereflow = { version = "0.3.2", path = "./fluereflow" } ratatui = { version = "0.25.0", features = ["all-widgets"] } diff --git a/fluere-config/Cargo.toml b/fluere-config/Cargo.toml index d736838..3d9b27d 100644 --- a/fluere-config/Cargo.toml +++ b/fluere-config/Cargo.toml @@ -13,7 +13,11 @@ repository = "https://github.com/SkuldNorniern/fluere" dirs = "5.0.1" serde = { version = "1.0.193", features = ["derive"]} toml = "0.8.8" +log = { version = "0.4.21", features = ["std"], optional = true } [lib] name = "fluere_config" + +[features] +log = ["dep:log"] \ No newline at end of file diff --git a/fluere-config/src/init.rs b/fluere-config/src/init.rs index a654d02..648218a 100644 --- a/fluere-config/src/init.rs +++ b/fluere-config/src/init.rs @@ -1,20 +1,35 @@ -use dirs::config_dir; +use std::{default::Default, env, fs, path::Path, path::PathBuf}; use crate::Config; -use std::{default::Default, env, fs, path::Path, path::PathBuf}; +use dirs::config_dir; + +#[cfg(feature = "log")] +use log::{debug, error, info, warn}; impl Config { pub fn new() -> Self { let path_base = home_config_path(); let path_file = path_base.join(Path::new("fluere.toml")); - println!("path_file: {:?}", path_file); + + #[cfg(feature = "log")] + debug!("Using config file from: {:?}", path_file); + #[cfg(not(feature = "log"))] + println!("Using config file from: {:?}", path_file); if !path_base.exists() { match fs::create_dir_all(&path_base) { - Ok(_) => (), + Ok(_) => { + #[cfg(feature = "log")] + debug!("Created directory at {:?}", path_base); + () + } Err(e) => { + #[cfg(feature = "log")] + error!("Failed to create directory at {:?}: {}", path_base, e); + #[cfg(not(feature = "log"))] eprintln!("Failed to create directory at {:?}: {}", path_base, e); + return Config::default(); } } @@ -25,8 +40,16 @@ impl Config { } match Self::load(path_file.to_str().unwrap().to_string()) { - Ok(config) => config, + Ok(config) => { + #[cfg(feature = "log")] + debug!("Loaded configuration from: {:?}", path_file); + + config + } Err(_) => { + #[cfg(feature = "log")] + warn!("failed to load configuration, using default config"); + #[cfg(not(feature = "log"))] println!("failed to load configuration, using default config"); Config::default() } From f92cbb6968c9f033bb458a93ddc64d3be93c5038 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 16:49:39 +0900 Subject: [PATCH 24/44] build(config): bump version of fluere-config --- Cargo.toml | 2 +- fluere-config/Cargo.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 1f65aea..1dd2129 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,7 +31,7 @@ serde = { version = "1.0.193", features = ["derive"] } toml = "0.8.8" fluere_plugin = { version = "0.2.0", path = "./fluere-plugin", features = ["log"] } -fluere-config = { version = "0.1.3", path = "./fluere-config", features = ["log"] } +fluere-config = { version = "0.2.0", path = "./fluere-config", features = ["log"] } fluereflow = { version = "0.3.2", path = "./fluereflow" } ratatui = { version = "0.25.0", features = ["all-widgets"] } diff --git a/fluere-config/Cargo.toml b/fluere-config/Cargo.toml index 3d9b27d..4779fa4 100644 --- a/fluere-config/Cargo.toml +++ b/fluere-config/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "fluere-config" -version = "0.1.3" +version = "0.2.0" edition = "2021" readme = "README.md" description = "Configuration manager for Fluere." From 1807564d699bd198f74e68e7a3f27068f72649c6 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Mon, 11 Mar 2024 17:07:08 +0900 Subject: [PATCH 25/44] style: `rustfmt` --- fluere-config/Cargo.toml | 2 +- fluere-plugin/Cargo.toml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/fluere-config/Cargo.toml b/fluere-config/Cargo.toml index 4779fa4..8808c78 100644 --- a/fluere-config/Cargo.toml +++ b/fluere-config/Cargo.toml @@ -20,4 +20,4 @@ name = "fluere_config" [features] -log = ["dep:log"] \ No newline at end of file +log = ["dep:log"] diff --git a/fluere-plugin/Cargo.toml b/fluere-plugin/Cargo.toml index 667cea2..a6839e9 100644 --- a/fluere-plugin/Cargo.toml +++ b/fluere-plugin/Cargo.toml @@ -16,7 +16,7 @@ repository = "https://github.com/SkuldNorniern/fluere" [dependencies] git2 = "0.18.1" tokio = { version = "1.32", features = ["full","macros", "rt-multi-thread"] } -fluere-config = { version = "0.1.2", path = "../fluere-config" } +fluere-config = { version = "0.2.0", path = "../fluere-config" } #fluere-plugin-trait = { path = "../fluere-plugin-trait" } fluereflow = { version = "0.3.2", path = "../fluereflow" } dirs = "5.0.1" @@ -30,4 +30,4 @@ name = "fluere_plugin" [features] -log = ["dep:log"] \ No newline at end of file +log = ["dep:log"] From 88d509f33fadd90839b7244680f972de03c99446 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Tue, 12 Mar 2024 09:49:43 +0900 Subject: [PATCH 26/44] refactor: commented unused `packet_count` --- src/net/live_fluereflow.rs | 6 +++--- src/net/online_fluereflow.rs | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/net/live_fluereflow.rs b/src/net/live_fluereflow.rs index a125b19..549f3fc 100644 --- a/src/net/live_fluereflow.rs +++ b/src/net/live_fluereflow.rs @@ -106,7 +106,7 @@ pub async fn online_packet_capture(arg: Args) { let recent_flows: Arc>> = Arc::new(Mutex::new(Vec::new())); let active_flow = Arc::new(Mutex::new(HashMap::new())); - let mut packet_count = 0; + // let mut packet_count = 0; enable_raw_mode().expect("Unable to enable raw mode"); let mut stdout = io::stdout(); @@ -268,7 +268,7 @@ pub async fn online_packet_capture(arg: Args) { } } - packet_count += 1; + // packet_count += 1; // slow down the loop for windows to avoid random shutdown // if packet_count % sleep_windows == 0 && cfg!(target_os = "windows") { // println!("Slow down the loop for windows"); @@ -280,7 +280,7 @@ pub async fn online_packet_capture(arg: Args) { let mut last_export_unix_time_guard = last_export_unix_time.lock().await; if last_export_guard.elapsed() >= Duration::from_millis(interval) && interval != 0 { let mut expired_flows = vec![]; - packet_count = 0; + // packet_count = 0; for (key, flow) in active_flow_guard.iter() { if flow_timeout > 0 && flow.last < (time - (flow_timeout * 1000)) { trace!("flow expired"); diff --git a/src/net/online_fluereflow.rs b/src/net/online_fluereflow.rs index 4936232..f2c8f01 100644 --- a/src/net/online_fluereflow.rs +++ b/src/net/online_fluereflow.rs @@ -69,7 +69,7 @@ pub async fn packet_capture(arg: Args) { let mut records: Vec = Vec::new(); let mut active_flow: HashMap = HashMap::new(); - let mut packet_count = 0; + // let mut packet_count = 0; loop { match cap.next_packet() { @@ -167,7 +167,7 @@ pub async fn packet_capture(arg: Args) { } } - packet_count += 1; + // 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 { @@ -179,7 +179,7 @@ pub async fn packet_capture(arg: Args) { // Export flows if the interval has been reached if last_export.elapsed() >= Duration::from_millis(interval) && interval != 0 { let mut expired_flows = vec![]; - packet_count = 0; + // packet_count = 0; for (key, flow) in active_flow.iter() { if flow_timeout > 0 && flow.last < (time - (flow_timeout * 1000)) { trace!("flow expired"); From b0d7f18fd7973d55aeaa536f72b87dfbd6e48d52 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Tue, 12 Mar 2024 10:24:18 +0900 Subject: [PATCH 27/44] style: `rustfmt` --- fluere-config/src/init.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fluere-config/src/init.rs b/fluere-config/src/init.rs index 648218a..86be016 100644 --- a/fluere-config/src/init.rs +++ b/fluere-config/src/init.rs @@ -5,7 +5,7 @@ use crate::Config; use dirs::config_dir; #[cfg(feature = "log")] -use log::{debug, error, info, warn}; +use log::{debug, error, warn}; impl Config { pub fn new() -> Self { From 7f6815d184e811a60fefe026359f984b693c3392 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Tue, 12 Mar 2024 10:25:10 +0900 Subject: [PATCH 28/44] refactor: change log level of few case --- src/net/capture.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/net/capture.rs b/src/net/capture.rs index 5394ac6..05cc0f1 100644 --- a/src/net/capture.rs +++ b/src/net/capture.rs @@ -58,7 +58,7 @@ pub fn find_device(identifier: &str) -> Result { if let Ok(index) = identifier.parse::() { if let Some(device) = devices.get(index) { let duration = start.elapsed(); - info!("Device {} captured in {:?}", device.name, duration); + debug!("Device {} captured in {:?}", device.name, duration); return Ok(device.clone()); } else { return Err(NetError::DeviceError(DeviceError::InvalidDeviceIndex( @@ -70,7 +70,7 @@ pub fn find_device(identifier: &str) -> Result { for device in devices { if device.name == identifier { let duration = start.elapsed(); - info!("Device {} captured in {:?}", device.name, duration); + debug!("Device {} captured in {:?}", device.name, duration); return Ok(device); } } From 7dbbd21a55c41b2cc1f6e30600e1e608dce4f3dd Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Tue, 12 Mar 2024 10:25:41 +0900 Subject: [PATCH 29/44] fix: remove deprecated function --- src/net/parser/time.rs | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/net/parser/time.rs b/src/net/parser/time.rs index 7bc7dd4..aa80ca3 100644 --- a/src/net/parser/time.rs +++ b/src/net/parser/time.rs @@ -1,14 +1,17 @@ use chrono::prelude::*; +use chrono::DateTime; +#[inline] pub fn parse_microseconds(sec: u64, usec: u64) -> u64 { sec * 1000000 + usec } +#[inline] pub fn microseconds_to_timestamp(usec: u64) -> String { - let naive = NaiveDateTime::from_timestamp_opt(usec as i64, 0); - + let naive = DateTime::from_timestamp(usec as i64, 0).unwrap().naive_utc(); + #[cfg(not(target_os = "windows"))] - let datetime = DateTime::::from_naive_utc_and_offset(naive.unwrap(), Utc); + let datetime = DateTime::::from_naive_utc_and_offset(naive, Utc); #[cfg(target_os = "windows")] let datetime = DateTime::::from_utc(naive.unwrap(), Utc); From f4b6e00fe0d06060ae4cd3b459c0ff39a31017de Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Tue, 12 Mar 2024 10:25:58 +0900 Subject: [PATCH 30/44] refactor(plugin): use drop on task await --- fluere-plugin/src/lib.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/fluere-plugin/src/lib.rs b/fluere-plugin/src/lib.rs index 39e545c..0acda82 100644 --- a/fluere-plugin/src/lib.rs +++ b/fluere-plugin/src/lib.rs @@ -301,7 +301,7 @@ impl PluginManager { pub async fn await_completion(&self, target_worker: Arc>>) { let worker_clone = target_worker.clone(); - let _ = worker_clone.lock().await; + let worker = worker_clone.lock().await; // Cleanup each plugin before exiting let lua_clone = self.lua.clone(); @@ -326,5 +326,9 @@ impl PluginManager { println!("cleanup function not found in plugin: {}", plugin_name); } } + + drop(lua); + drop(plugins); + drop(worker); } } From 16da5dbe9bbcb4b28ac48027fe93197ec44e8ccd Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Tue, 12 Mar 2024 10:26:52 +0900 Subject: [PATCH 31/44] refactor: remove unused struct --- src/main.rs | 55 ++++++++--------------------------------------------- 1 file changed, 8 insertions(+), 47 deletions(-) diff --git a/src/main.rs b/src/main.rs index c84fab5..b022d27 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,7 +6,7 @@ pub mod cli; pub mod logger; pub mod net; -pub mod plugin; +// pub mod plugin; pub mod types; pub mod utils; @@ -17,15 +17,15 @@ use crate::logger::{Logger, Logstdout}; use crate::net::capture::DeviceError; // use env_logger;::{init, Logger}; -use log::{debug, info, Level, LevelFilter}; +use log::{debug, 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 { @@ -97,52 +97,14 @@ fn from_verbose(level: u8) -> LevelFilter { } } -struct Fluere { - interface: String, - args: types::Args, - mode: Mode, - logger: Logger, - verbose: Level, - // interfaces: Vec, -} -impl Fluere { - fn new( - interface: String, - args: types::Args, - mode: Mode, - logger: Logger, - verbose: Level, - ) -> Fluere { - Fluere { - interface, - args, - mode, - logger, - verbose, - } - } - async fn online_fluereflow(&self, params: types::Args) { - net::online_fluereflow::packet_capture(params).await; - info!("Online fluereflow mode completed"); - } -} // This is the main function of the application. // It gets the command line arguments, parses them, and calls the appropriate functions based on the arguments. #[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); - // } - // }; - 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; @@ -163,7 +125,6 @@ async fn main() { .await .expect("Error on live mode"), Mode::Pcap => net::pcap_capture(parems.0).await, - _ => unreachable!(), } } else { exit(0); From 545a3578e29d1ad90df3bb470b6f28e2e860af9d Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Tue, 12 Mar 2024 10:27:16 +0900 Subject: [PATCH 32/44] refactor: re-arrange some codes --- src/net/online_fluereflow.rs | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/net/online_fluereflow.rs b/src/net/online_fluereflow.rs index f2c8f01..901ad8c 100644 --- a/src/net/online_fluereflow.rs +++ b/src/net/online_fluereflow.rs @@ -21,6 +21,10 @@ use crate::{ }; use fluere_config::Config; + +// FEAT:TASK: set plugin as feature +// | Since the plugin manager uses Lua, for edge cases that require minimal feature, +// | setting the plugin as a feature would be beneficial. use fluere_plugin::PluginManager; use fluereflow::FluereRecord; @@ -139,7 +143,10 @@ pub async fn packet_capture(arg: Args) { //println!("time: {:?}", time); let pkt = flowdata.min_pkt; let ttl = flowdata.min_ttl; - //println!("current inputed flow{:?}", active_flow.get(&key_value).unwrap()); + trace!( + "current inputed flow{:?}", + active_flow.get(&key_value).unwrap() + ); let flow_key = if is_reverse { &reverse_key } else { &key_value }; if let Some(flow) = active_flow.get_mut(flow_key) { let update_key = UDFlowKey { @@ -233,6 +240,7 @@ pub async fn packet_capture(arg: Args) { }); plugin_manager.await_completion(plugin_worker).await; + drop(plugin_manager); let result = tasks.await; info!("Exporting task excutation result: {:?}", result); } From 3a37d7542a7241b27780de03454da1568829050e Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Tue, 12 Mar 2024 10:40:47 +0900 Subject: [PATCH 33/44] fix: fix while to loop with match cause of `pcap-rs`'s recommendation --- src/net/packet_pcap.rs | 40 ++++++++++++++++++++++++---------------- 1 file changed, 24 insertions(+), 16 deletions(-) diff --git a/src/net/packet_pcap.rs b/src/net/packet_pcap.rs index 82138de..0f2776c 100644 --- a/src/net/packet_pcap.rs +++ b/src/net/packet_pcap.rs @@ -20,7 +20,7 @@ pub async fn pcap_capture(args: Args) { let cap = &mut cap_device.capture; let file_dir = "./output"; - let mut packet_count = 0; + // let mut packet_count = 0; match fs::create_dir_all(<&str>::clone(&file_dir)) { Ok(_) => debug!("Created directory: {}", file_dir), Err(error) => panic!("Problem creating directory: {:?}", error), @@ -33,21 +33,29 @@ pub async fn pcap_capture(args: Args) { }; let start = Instant::now(); - while let Ok(packet) = cap.next_packet() { - trace!("received packet"); - //println!("packet: {:?}", packet); - file.write(&packet); - - packet_count += 1; - // slow down the loop for windows to avoid random shutdown - // if packet_count % sleep_windows == 0 && cfg!(target_os = "windows") { - // println!("Slow down the loop for windows"); - // sleep(Duration::from_millis(0)).await; - // } - - // Check if the duration has been reached - if start.elapsed() >= Duration::from_millis(duration) && duration != 0 { - break; + + loop { + match cap.next_packet() { + Err(_) => { + continue; + } + Ok(packet) => { + trace!("received packet"); + //println!("packet: {:?}", packet); + file.write(&packet); + + // packet_count += 1; + // slow down the loop for windows to avoid random shutdown + // if packet_count % sleep_windows == 0 && cfg!(target_os = "windows") { + // println!("Slow down the loop for windows"); + // sleep(Duration::from_millis(0)).await; + // } + + // Check if the duration has been reached + if start.elapsed() >= Duration::from_millis(duration) && duration != 0 { + break; + } + } } } debug!("Captured in {:?}", start.elapsed()); From 19a570965727816e9f3851cbd26c41ceec012fe4 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Tue, 12 Mar 2024 13:35:33 +0900 Subject: [PATCH 34/44] feat: add some more debug logs --- src/net/capture.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/net/capture.rs b/src/net/capture.rs index 05cc0f1..bb7917b 100644 --- a/src/net/capture.rs +++ b/src/net/capture.rs @@ -32,7 +32,9 @@ impl CaptureDevice { let capture = initialize_capture(device.clone())?; let name: Cow<'static, str> = Cow::Owned(device.name); let desc: Cow<'static, str> = Cow::Owned(device.desc.unwrap_or("".to_string())); - + debug!("Using device: {}", name); + debug!("Device description: {}", desc); + debug!("Addresses: {:?}", device.addresses); Ok(CaptureDevice { name, desc, @@ -46,7 +48,6 @@ impl Drop for CaptureDevice { fn drop(&mut self) { info!("Closing capture session for device {}", self.name); // println!("Closing capture session for device {}", self.name); - // self.capture.; } } pub fn find_device(identifier: &str) -> Result { @@ -84,6 +85,8 @@ fn initialize_capture(device: Device) -> Result, PcapError> { info!("Opening capture session for device {}", device.name); Capture::from_device(device)? .promisc(true) + + // FEAT:TASK: set snaplen as a Flag from the CLI .snaplen(1024) .timeout(60000) .immediate_mode(true) From 4af371bd85480ae0c65f6bcd6839952baf8b4465 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Tue, 12 Mar 2024 05:01:48 +0000 Subject: [PATCH 35/44] style: `rustfmt` --- src/net/capture.rs | 3 +-- src/net/online_fluereflow.rs | 2 +- src/net/packet_pcap.rs | 2 +- src/net/parser/time.rs | 6 ++++-- 4 files changed, 7 insertions(+), 6 deletions(-) diff --git a/src/net/capture.rs b/src/net/capture.rs index bb7917b..a4d4bab 100644 --- a/src/net/capture.rs +++ b/src/net/capture.rs @@ -85,8 +85,7 @@ fn initialize_capture(device: Device) -> Result, PcapError> { info!("Opening capture session for device {}", device.name); Capture::from_device(device)? .promisc(true) - - // FEAT:TASK: set snaplen as a Flag from the CLI + // FEAT:TASK: set snaplen as a Flag from the CLI .snaplen(1024) .timeout(60000) .immediate_mode(true) diff --git a/src/net/online_fluereflow.rs b/src/net/online_fluereflow.rs index 901ad8c..f5e452d 100644 --- a/src/net/online_fluereflow.rs +++ b/src/net/online_fluereflow.rs @@ -23,7 +23,7 @@ use crate::{ use fluere_config::Config; // FEAT:TASK: set plugin as feature -// | Since the plugin manager uses Lua, for edge cases that require minimal feature, +// | Since the plugin manager uses Lua, for edge cases that require minimal feature, // | setting the plugin as a feature would be beneficial. use fluere_plugin::PluginManager; use fluereflow::FluereRecord; diff --git a/src/net/packet_pcap.rs b/src/net/packet_pcap.rs index 0f2776c..18372cb 100644 --- a/src/net/packet_pcap.rs +++ b/src/net/packet_pcap.rs @@ -33,7 +33,7 @@ pub async fn pcap_capture(args: Args) { }; let start = Instant::now(); - + loop { match cap.next_packet() { Err(_) => { diff --git a/src/net/parser/time.rs b/src/net/parser/time.rs index aa80ca3..0d10846 100644 --- a/src/net/parser/time.rs +++ b/src/net/parser/time.rs @@ -8,8 +8,10 @@ pub fn parse_microseconds(sec: u64, usec: u64) -> u64 { #[inline] pub fn microseconds_to_timestamp(usec: u64) -> String { - let naive = DateTime::from_timestamp(usec as i64, 0).unwrap().naive_utc(); - + let naive = DateTime::from_timestamp(usec as i64, 0) + .unwrap() + .naive_utc(); + #[cfg(not(target_os = "windows"))] let datetime = DateTime::::from_naive_utc_and_offset(naive, Utc); From 096ec800f7498b55c4a6a848304b756f8ae2fb68 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Tue, 12 Mar 2024 15:56:20 +0900 Subject: [PATCH 36/44] feat(plugin): remove use of `String` and use `Cow` --- fluere-plugin/src/lib.rs | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/fluere-plugin/src/lib.rs b/fluere-plugin/src/lib.rs index 0acda82..ef2a3d3 100644 --- a/fluere-plugin/src/lib.rs +++ b/fluere-plugin/src/lib.rs @@ -4,6 +4,8 @@ use std::sync::Arc; mod downloader; mod util; +use std::borrow::Cow; + use downloader::download_plugin_from_github; use util::home_cache_path; @@ -19,7 +21,7 @@ pub struct PluginManager { lua: Arc>, sender: mpsc::Sender, receiver: Arc>>, - plugins: Arc>>, + plugins: Arc>>>, } impl PluginManager { @@ -98,8 +100,7 @@ impl PluginManager { Ok(()) }).expect(format!("Error on plugin: {}", name).as_str(()));*/ //lua_guard.load(&code).exec().expect(format!("Error on plugin: {}", name).as_str()); - - plugins_guard.insert(name.clone()); + let _ = plugins_guard.insert(std::borrow::Cow::Owned(name.clone())); #[cfg(feature = "log")] info!("Loaded plugin {}", name); #[cfg(not(feature = "log"))] @@ -168,7 +169,7 @@ impl PluginManager { Ok(()) }).expect(format!("Error on plugin: {}", name).as_str());*/ - plugins_guard.insert(name.clone()); + let _ = plugins_guard.insert(std::borrow::Cow::Owned(name.clone())); #[cfg(feature = "log")] info!("Loaded plugin {}", name); #[cfg(not(feature = "log"))] @@ -267,7 +268,7 @@ impl PluginManager { for plugin_name in plugins.iter() { let plugin_table: mlua::Table = lua .globals() - .get(plugin_name.as_str()) + .get(plugin_name.as_ref()) .expect("Plugin table not found"); if let Ok(func) = plugin_table.get::<_, mlua::Function>("process_data") { @@ -313,7 +314,7 @@ impl PluginManager { for plugin_name in plugins.iter() { let plugin_table: mlua::Table = lua .globals() - .get(plugin_name.as_str()) + .get(plugin_name.as_ref()) .expect("Plugin table not found"); if let Ok(func) = plugin_table.get::<_, mlua::Function>("cleanup") { @@ -332,3 +333,10 @@ impl PluginManager { drop(worker); } } + +impl Drop for PluginManager { + fn drop(&mut self) { + drop(self.plugins.lock()); + drop(self.lua.lock()); + } +} \ No newline at end of file From 79f809c4f3356e6e8c4a27b8efcd87ff44b8737c Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Wed, 13 Mar 2024 09:29:01 +0900 Subject: [PATCH 37/44] feat: add todo for removing panic! --- src/net/live_fluereflow.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/net/live_fluereflow.rs b/src/net/live_fluereflow.rs index 549f3fc..fde45a8 100644 --- a/src/net/live_fluereflow.rs +++ b/src/net/live_fluereflow.rs @@ -86,6 +86,8 @@ pub async fn online_packet_capture(arg: Args) { let file_dir = "./output"; match fs::create_dir_all(<&str>::clone(&file_dir)) { Ok(_) => debug!("Created directory: {}", file_dir), + + // FIX:ASAP: Remove Panic, return io error instead Err(error) => panic!("Problem creating directory: {:?}", error), }; From 345ed80e5c38d06186a25c42bd2c3c35f23baa7e Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Wed, 13 Mar 2024 09:29:36 +0900 Subject: [PATCH 38/44] feat: apply code review's feedback --- src/net/offline_fluereflows.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/net/offline_fluereflows.rs b/src/net/offline_fluereflows.rs index e200b5c..a520ceb 100644 --- a/src/net/offline_fluereflows.rs +++ b/src/net/offline_fluereflows.rs @@ -140,6 +140,6 @@ pub async fn fluereflow_fileparse(arg: Args) { let result = tasks.await; info!("Export {} result: {:?}", file_path, result); - println!("Active flow {:?}", ac_flow_cnt); - println!("Ended flow {:?}", ended_flow_cnt); + info!("Active flow {:?}", ac_flow_cnt); + info!("Ended flow {:?}", ended_flow_cnt); } From 69efc2fa81a0f7558f1cbd3f928ae373d2850fad Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Wed, 13 Mar 2024 09:29:58 +0900 Subject: [PATCH 39/44] refactor: re-batch imports --- src/net/parser/ports.rs | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/net/parser/ports.rs b/src/net/parser/ports.rs index 0e8dec9..17f2853 100644 --- a/src/net/parser/ports.rs +++ b/src/net/parser/ports.rs @@ -1,8 +1,11 @@ -use pnet::packet::tcp::TcpPacket; -use pnet::packet::udp::UdpPacket; - use crate::net::errors::NetError; +use log::debug; +use pnet::packet::{ + tcp::TcpPacket, + udp::UdpPacket, +}; + pub fn parse_ports(protocol: u8, payload: &[u8]) -> Result<(u16, u16), NetError> { match protocol { 58 => Ok((0, 0)), @@ -20,7 +23,7 @@ pub fn parse_ports(protocol: u8, payload: &[u8]) -> Result<(u16, u16), NetError> 1 => Ok((0, 0)), 0 => Ok((0, 0)), _ => { - println!("Unknown protocol: {}", protocol); + debug!("Unknown protocol: {}", protocol); Err(NetError::UnknownProtocol { protocol: protocol.to_string(), }) From 75f69a848a18daf4e4080a9d880398cf0a382621 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Wed, 13 Mar 2024 09:33:16 +0900 Subject: [PATCH 40/44] fix: remove un-needed `unwrap` --- src/net/parser/time.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/net/parser/time.rs b/src/net/parser/time.rs index 0d10846..79c5ac8 100644 --- a/src/net/parser/time.rs +++ b/src/net/parser/time.rs @@ -16,7 +16,7 @@ pub fn microseconds_to_timestamp(usec: u64) -> String { let datetime = DateTime::::from_naive_utc_and_offset(naive, Utc); #[cfg(target_os = "windows")] - let datetime = DateTime::::from_utc(naive.unwrap(), Utc); + let datetime = DateTime::::from_utc(naive, Utc); datetime.format("%Y-%m-%d_%H-%M-%S UTC").to_string() } From 914d7b31fcc62cf0442a9be31488f780d2fde499 Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Wed, 13 Mar 2024 00:37:08 +0000 Subject: [PATCH 41/44] style: `rustfmt` --- fluere-plugin/src/lib.rs | 5 +++-- src/net/live_fluereflow.rs | 2 +- src/net/parser/ports.rs | 5 +---- 3 files changed, 5 insertions(+), 7 deletions(-) diff --git a/fluere-plugin/src/lib.rs b/fluere-plugin/src/lib.rs index ef2a3d3..c195e13 100644 --- a/fluere-plugin/src/lib.rs +++ b/fluere-plugin/src/lib.rs @@ -169,7 +169,8 @@ impl PluginManager { Ok(()) }).expect(format!("Error on plugin: {}", name).as_str());*/ - let _ = plugins_guard.insert(std::borrow::Cow::Owned(name.clone())); + let _ = plugins_guard + .insert(std::borrow::Cow::Owned(name.clone())); #[cfg(feature = "log")] info!("Loaded plugin {}", name); #[cfg(not(feature = "log"))] @@ -339,4 +340,4 @@ impl Drop for PluginManager { drop(self.plugins.lock()); drop(self.lua.lock()); } -} \ No newline at end of file +} diff --git a/src/net/live_fluereflow.rs b/src/net/live_fluereflow.rs index fde45a8..f99e2d3 100644 --- a/src/net/live_fluereflow.rs +++ b/src/net/live_fluereflow.rs @@ -86,7 +86,7 @@ pub async fn online_packet_capture(arg: Args) { let file_dir = "./output"; match fs::create_dir_all(<&str>::clone(&file_dir)) { Ok(_) => debug!("Created directory: {}", file_dir), - + // FIX:ASAP: Remove Panic, return io error instead Err(error) => panic!("Problem creating directory: {:?}", error), }; diff --git a/src/net/parser/ports.rs b/src/net/parser/ports.rs index 17f2853..29fbf96 100644 --- a/src/net/parser/ports.rs +++ b/src/net/parser/ports.rs @@ -1,10 +1,7 @@ use crate::net::errors::NetError; use log::debug; -use pnet::packet::{ - tcp::TcpPacket, - udp::UdpPacket, -}; +use pnet::packet::{tcp::TcpPacket, udp::UdpPacket}; pub fn parse_ports(protocol: u8, payload: &[u8]) -> Result<(u16, u16), NetError> { match protocol { From 4be4977c61b422816f2351121dae3372ebdad06e Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Wed, 13 Mar 2024 10:27:47 +0900 Subject: [PATCH 42/44] feat: change `net` functions to use `net/mod.rs`'s `NetError` --- src/main.rs | 2 +- src/net/mod.rs | 18 ++++++++++++++++-- src/net/parser/fluereflows.rs | 26 +++++++------------------- src/net/parser/keys.rs | 8 ++++---- src/net/parser/ports.rs | 24 ++++++++++-------------- src/net/parser/tos.rs | 6 +++--- 6 files changed, 41 insertions(+), 43 deletions(-) diff --git a/src/main.rs b/src/main.rs index b022d27..a1d1d0c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -14,7 +14,7 @@ use std::fs::File; use std::{fmt::Display, process::exit}; use crate::logger::{Logger, Logstdout}; -use crate::net::capture::DeviceError; +use crate::net::DeviceError; // use env_logger;::{init, Logger}; use log::{debug, Level, LevelFilter}; diff --git a/src/net/mod.rs b/src/net/mod.rs index 6a94278..713ae12 100644 --- a/src/net/mod.rs +++ b/src/net/mod.rs @@ -1,6 +1,6 @@ //mod fluereflow -pub mod capture; -pub mod errors; +mod capture; +// pub mod errors; mod flows; // mod interface; pub mod live_fluereflow; @@ -28,6 +28,11 @@ use pcap::Error; pub enum NetError { DeviceError(DeviceError), PcapError(Error), + UnknownProtocol(u8), + UnknownEtherType(String), + UnknownDSCP(u8), + InvalidPacket, + EmptyPacket, } impl From for NetError { @@ -47,6 +52,15 @@ impl Display for NetError { match self { NetError::DeviceError(err) => err.fmt(f), NetError::PcapError(err) => err.fmt(f), + NetError::UnknownProtocol(protocol) => { + write!(f, "Unknown protocol: {}", protocol) + }, + NetError::UnknownEtherType(ether_type) => { + write!(f, "Unknown ether type: {}", ether_type) + }, + NetError::UnknownDSCP(dscp) => write!(f, "Unknown dscp: {}", dscp), + NetError::InvalidPacket => write!(f, "Invalid packet"), + NetError::EmptyPacket => write!(f, "Empty packet"), } } } diff --git a/src/net/parser/fluereflows.rs b/src/net/parser/fluereflows.rs index 545f44f..18f8fd1 100644 --- a/src/net/parser/fluereflows.rs +++ b/src/net/parser/fluereflows.rs @@ -1,6 +1,6 @@ use pcap; -use crate::net::errors::NetError; +use crate::net::NetError; use crate::net::parser::{ dscp_to_tos, parse_flags, parse_microseconds, parse_ports, protocol_to_number, }; @@ -130,9 +130,9 @@ pub fn parse_fluereflow(packet: pcap::Packet) -> Result<(usize, [u8; 9], FluereR arp_packet(time, i) } _ => { - return Err(NetError::UnknownProtocol { - protocol: ethernet_packet.get_ethertype().to_string(), - }) + return Err(NetError::UnknownEtherType ( + ethernet_packet.get_ethertype().to_string(), + )) } }; @@ -197,15 +197,8 @@ fn ipv4_packet(time: u64, packet: Ipv4Packet) -> Result<(usize, [u8; 9], FluereR let dst_ip = packet.get_destination(); // ports parsing - let parsed_ports = parse_ports(protocol, packet.payload()); - match parsed_ports { - Ok(_) => {} - Err(e) => { - println!("Unknown protocol {}\n Report to the developer", e); - return Err(e); - } - } - let (src_port, dst_port) = parsed_ports.unwrap(); + let (src_port, dst_port) = parse_ports(protocol, packet.payload())?; + // TCP flags Fin Syn Rst Psh Ack Urg Ece Cwr Ns let flags = parse_flags(protocol, packet.payload()); @@ -255,12 +248,7 @@ fn ipv6_packet(time: u64, packet: Ipv6Packet) -> Result<(usize, [u8; 9], FluereR let dst_ip = packet.get_destination(); // ports parsing - let parsed_ports = parse_ports(protocol, packet.payload()); - match parsed_ports { - Ok(_) => {} - Err(e) => return Err(e), - } - let (src_port, dst_port) = parsed_ports.unwrap(); + let (src_port, dst_port) = parse_ports(protocol, packet.payload())?; // TCP flags Fin Syn Rst Psh Ack Urg Ece Cwr Ns let flags = parse_flags(protocol, packet.payload()); diff --git a/src/net/parser/keys.rs b/src/net/parser/keys.rs index b35fc97..8532181 100644 --- a/src/net/parser/keys.rs +++ b/src/net/parser/keys.rs @@ -8,7 +8,7 @@ use pnet::packet::ipv6::Ipv6Packet; use pnet::packet::udp::UdpPacket; use pnet::packet::Packet; -use crate::net::errors::NetError; +use crate::net::NetError; use crate::net::parser::{parse_ports, protocol_to_number}; use crate::net::types::{Key, MacAddress}; @@ -146,9 +146,9 @@ pub fn parse_keys(packet: pcap::Packet) -> Result<(Key, Key), NetError> { } _ => { - return Err(NetError::UnknownProtocol { - protocol: ethernet_packet.get_ethertype().to_string(), - }) + return Err(NetError::UnknownEtherType ( + ethernet_packet.get_ethertype().to_string() + )) } }; diff --git a/src/net/parser/ports.rs b/src/net/parser/ports.rs index 29fbf96..03f4898 100644 --- a/src/net/parser/ports.rs +++ b/src/net/parser/ports.rs @@ -1,4 +1,4 @@ -use crate::net::errors::NetError; +use crate::net::NetError; use log::debug; use pnet::packet::{tcp::TcpPacket, udp::UdpPacket}; @@ -7,27 +7,23 @@ pub fn parse_ports(protocol: u8, payload: &[u8]) -> Result<(u16, u16), NetError> match protocol { 58 => Ok((0, 0)), 17 => { - let udp = UdpPacket::new(payload).unwrap(); - - Ok((udp.get_source(), udp.get_destination())) + match UdpPacket::new(payload) { + Some(udp) => Ok((udp.get_source(), udp.get_destination())), + None => Err(NetError::InvalidPacket), + } } 6 => { - let tcp = TcpPacket::new(payload).unwrap(); - - Ok((tcp.get_source(), tcp.get_destination())) + match TcpPacket::new(payload){ + Some(tcp) => Ok((tcp.get_source(), tcp.get_destination())), + None => Err(NetError::InvalidPacket), + } } 2 => Ok((0, 0)), 1 => Ok((0, 0)), 0 => Ok((0, 0)), _ => { debug!("Unknown protocol: {}", protocol); - Err(NetError::UnknownProtocol { - protocol: protocol.to_string(), - }) + Err(NetError::UnknownProtocol(protocol)) } } - - //Err(NetError::UnknownProtocol { - // protocol: protocol.to_string(), - //}) } diff --git a/src/net/parser/tos.rs b/src/net/parser/tos.rs index 0d4c2bf..26a006c 100644 --- a/src/net/parser/tos.rs +++ b/src/net/parser/tos.rs @@ -1,6 +1,6 @@ -use crate::net::errors::ParseError; +use crate::net::NetError; -pub fn dscp_to_tos(dscp: u8) -> Result { +pub fn dscp_to_tos(dscp: u8) -> Result { let tos = match dscp { 0 => 0, 8 => 32, @@ -23,7 +23,7 @@ pub fn dscp_to_tos(dscp: u8) -> Result { 46 => 184, 48 => 192, 56 => 224, - _ => return Err(ParseError::UnknownDSCP { dscp }), + _ => return Err(NetError::UnknownDSCP ( dscp )), }; Ok(tos) From 5d982dcac4b13a8658f55aac65aa1eacfcd8ceea Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Wed, 13 Mar 2024 11:56:37 +0900 Subject: [PATCH 43/44] feat: log internal parsing error for later feature search --- src/net/offline_fluereflows.rs | 5 ++++- src/net/online_fluereflow.rs | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/src/net/offline_fluereflows.rs b/src/net/offline_fluereflows.rs index a520ceb..3d9ee4f 100644 --- a/src/net/offline_fluereflows.rs +++ b/src/net/offline_fluereflows.rs @@ -51,7 +51,10 @@ pub async fn fluereflow_fileparse(arg: Args) { } let (doctets, raw_flags, flowdata) = match parse_fluereflow(packet.clone()) { Ok(result) => result, - Err(_) => continue, + Err(e) => { + debug!("{}", e); + continue; + } }; let flags = TcpFlags::new(raw_flags); //pushing packet in to active_flows if it is not present diff --git a/src/net/online_fluereflow.rs b/src/net/online_fluereflow.rs index f5e452d..f040658 100644 --- a/src/net/online_fluereflow.rs +++ b/src/net/online_fluereflow.rs @@ -94,7 +94,10 @@ pub async fn packet_capture(arg: Args) { let (doctets, raw_flags, flowdata) = match parse_fluereflow(packet.clone()) { Ok(result) => result, - Err(_) => continue, + Err(e) => { + debug!("{}", e); + continue; + } }; let flags = TcpFlags::new(raw_flags); From 443fd897097662375d83b2db5a09b065f319b45a Mon Sep 17 00:00:00 2001 From: SkuldNorniern Date: Wed, 13 Mar 2024 02:57:37 +0000 Subject: [PATCH 44/44] style: `rustfmt` --- src/net/mod.rs | 4 ++-- src/net/offline_fluereflows.rs | 6 +++--- src/net/parser/fluereflows.rs | 6 +++--- src/net/parser/keys.rs | 6 +++--- src/net/parser/ports.rs | 20 ++++++++------------ src/net/parser/tos.rs | 2 +- 6 files changed, 20 insertions(+), 24 deletions(-) diff --git a/src/net/mod.rs b/src/net/mod.rs index 713ae12..38f58d1 100644 --- a/src/net/mod.rs +++ b/src/net/mod.rs @@ -54,10 +54,10 @@ impl Display for NetError { NetError::PcapError(err) => err.fmt(f), NetError::UnknownProtocol(protocol) => { write!(f, "Unknown protocol: {}", protocol) - }, + } NetError::UnknownEtherType(ether_type) => { write!(f, "Unknown ether type: {}", ether_type) - }, + } NetError::UnknownDSCP(dscp) => write!(f, "Unknown dscp: {}", dscp), NetError::InvalidPacket => write!(f, "Invalid packet"), NetError::EmptyPacket => write!(f, "Empty packet"), diff --git a/src/net/offline_fluereflows.rs b/src/net/offline_fluereflows.rs index 3d9ee4f..c6d38cc 100644 --- a/src/net/offline_fluereflows.rs +++ b/src/net/offline_fluereflows.rs @@ -52,9 +52,9 @@ pub async fn fluereflow_fileparse(arg: Args) { let (doctets, raw_flags, flowdata) = match parse_fluereflow(packet.clone()) { Ok(result) => result, Err(e) => { - debug!("{}", e); - continue; - } + debug!("{}", e); + continue; + } }; let flags = TcpFlags::new(raw_flags); //pushing packet in to active_flows if it is not present diff --git a/src/net/parser/fluereflows.rs b/src/net/parser/fluereflows.rs index 18f8fd1..94a927a 100644 --- a/src/net/parser/fluereflows.rs +++ b/src/net/parser/fluereflows.rs @@ -1,9 +1,9 @@ use pcap; -use crate::net::NetError; use crate::net::parser::{ dscp_to_tos, parse_flags, parse_microseconds, parse_ports, protocol_to_number, }; +use crate::net::NetError; use fluereflow::FluereRecord; use pnet::packet::{ arp::ArpPacket, @@ -130,7 +130,7 @@ pub fn parse_fluereflow(packet: pcap::Packet) -> Result<(usize, [u8; 9], FluereR arp_packet(time, i) } _ => { - return Err(NetError::UnknownEtherType ( + return Err(NetError::UnknownEtherType( ethernet_packet.get_ethertype().to_string(), )) } @@ -198,7 +198,7 @@ fn ipv4_packet(time: u64, packet: Ipv4Packet) -> Result<(usize, [u8; 9], FluereR // ports parsing let (src_port, dst_port) = parse_ports(protocol, packet.payload())?; - + // TCP flags Fin Syn Rst Psh Ack Urg Ece Cwr Ns let flags = parse_flags(protocol, packet.payload()); diff --git a/src/net/parser/keys.rs b/src/net/parser/keys.rs index 8532181..1c50204 100644 --- a/src/net/parser/keys.rs +++ b/src/net/parser/keys.rs @@ -8,9 +8,9 @@ use pnet::packet::ipv6::Ipv6Packet; use pnet::packet::udp::UdpPacket; use pnet::packet::Packet; -use crate::net::NetError; use crate::net::parser::{parse_ports, protocol_to_number}; use crate::net::types::{Key, MacAddress}; +use crate::net::NetError; use std::net::IpAddr; @@ -146,8 +146,8 @@ pub fn parse_keys(packet: pcap::Packet) -> Result<(Key, Key), NetError> { } _ => { - return Err(NetError::UnknownEtherType ( - ethernet_packet.get_ethertype().to_string() + return Err(NetError::UnknownEtherType( + ethernet_packet.get_ethertype().to_string(), )) } }; diff --git a/src/net/parser/ports.rs b/src/net/parser/ports.rs index 03f4898..a778c1c 100644 --- a/src/net/parser/ports.rs +++ b/src/net/parser/ports.rs @@ -6,18 +6,14 @@ use pnet::packet::{tcp::TcpPacket, udp::UdpPacket}; pub fn parse_ports(protocol: u8, payload: &[u8]) -> Result<(u16, u16), NetError> { match protocol { 58 => Ok((0, 0)), - 17 => { - match UdpPacket::new(payload) { - Some(udp) => Ok((udp.get_source(), udp.get_destination())), - None => Err(NetError::InvalidPacket), - } - } - 6 => { - match TcpPacket::new(payload){ - Some(tcp) => Ok((tcp.get_source(), tcp.get_destination())), - None => Err(NetError::InvalidPacket), - } - } + 17 => match UdpPacket::new(payload) { + Some(udp) => Ok((udp.get_source(), udp.get_destination())), + None => Err(NetError::InvalidPacket), + }, + 6 => match TcpPacket::new(payload) { + Some(tcp) => Ok((tcp.get_source(), tcp.get_destination())), + None => Err(NetError::InvalidPacket), + }, 2 => Ok((0, 0)), 1 => Ok((0, 0)), 0 => Ok((0, 0)), diff --git a/src/net/parser/tos.rs b/src/net/parser/tos.rs index 26a006c..e03e296 100644 --- a/src/net/parser/tos.rs +++ b/src/net/parser/tos.rs @@ -23,7 +23,7 @@ pub fn dscp_to_tos(dscp: u8) -> Result { 46 => 184, 48 => 192, 56 => 224, - _ => return Err(NetError::UnknownDSCP ( dscp )), + _ => return Err(NetError::UnknownDSCP(dscp)), }; Ok(tos)