Merge pull request #23 from matissecallewaert/feature/add-different-o…

…utput-types

🕺 Feature/add different output types

feature-extraction-tool/Cargo.toml

@@ -29,6 +29,7 @@ chrono = "0.4.34"
 dashmap = "5.5.3"
 pcap = "1.3.0"
 pnet = "0.34.0"
+lazy_static = "1.4.0"
 
 [[bin]]
 name = "feature-extraction-tool"

feature-extraction-tool/src/args.rs

@@ -1,4 +1,4 @@
-use clap::{Parser, Subcommand};
+use clap::{Args, Parser, Subcommand};
 
 #[derive(Debug, Parser)]
 #[clap(author, version, about)]
@@ -20,6 +20,10 @@ pub enum Commands {
         /// The maximum lifespan of a flow in seconds
         lifespan: u64,
 
+        /// Output method
+        #[clap(flatten)]
+        export_method: Output,
+
         /// The print interval for open flows in seconds, needs to be smaller than the flow maximum lifespan
         interval: Option<u64>,
     },
@@ -43,9 +47,33 @@ pub enum Commands {
 
         /// The relative path to the pcap file
         path: String,
+
+        /// Output method
+        #[clap(flatten)]
+        export_method: Output,
     },
 }
 
+#[derive(Args, Debug, Clone)]
+pub struct Output {
+    /// Output method
+    #[clap(value_enum)]
+    pub method: ExportMethodType,
+
+    /// File path for output (used if method is File)
+    #[clap(required_if_eq("method", "Csv"))]
+    pub export_path: Option<String>,
+}
+
+#[derive(clap::ValueEnum, Clone, Debug)]
+pub enum ExportMethodType {
+    /// The output will be printed to the console
+    Print,
+
+    /// The output will be written to a CSV file
+    Csv,
+}
+
 #[derive(clap::ValueEnum, Clone, Debug)]
 pub enum GeneratedMachineType {
     /// The pcap file was generated on a Windows machine

feature-extraction-tool/src/flows/cic_flow.rs

@@ -1977,7 +1977,8 @@ mod tests {
         let mut cic_flow = setup_cic_flow();
 
         cic_flow.basic_flow.first_timestamp = chrono::Utc::now();
-        cic_flow.basic_flow.last_timestamp = chrono::Utc::now() + chrono::Duration::try_seconds(5).unwrap();
+        cic_flow.basic_flow.last_timestamp =
+            chrono::Utc::now() + chrono::Duration::try_seconds(5).unwrap();
 
         cic_flow.fwd_pkt_len_tot = 100;
         cic_flow.bwd_pkt_len_tot = 100;
@@ -1990,7 +1991,8 @@ mod tests {
         let mut cic_flow = setup_cic_flow();
 
         cic_flow.basic_flow.first_timestamp = chrono::Utc::now();
-        cic_flow.basic_flow.last_timestamp = chrono::Utc::now() + chrono::Duration::try_seconds(5).unwrap();
+        cic_flow.basic_flow.last_timestamp =
+            chrono::Utc::now() + chrono::Duration::try_seconds(5).unwrap();
 
         cic_flow.basic_flow.fwd_packet_count = 5;
         cic_flow.basic_flow.bwd_packet_count = 5;
@@ -2003,7 +2005,8 @@ mod tests {
         let mut cic_flow = setup_cic_flow();
 
         cic_flow.basic_flow.first_timestamp = chrono::Utc::now();
-        cic_flow.basic_flow.last_timestamp = chrono::Utc::now() + chrono::Duration::try_seconds(5).unwrap();
+        cic_flow.basic_flow.last_timestamp =
+            chrono::Utc::now() + chrono::Duration::try_seconds(5).unwrap();
 
         cic_flow.basic_flow.fwd_packet_count = 5;
 
@@ -2015,7 +2018,8 @@ mod tests {
         let mut cic_flow = setup_cic_flow();
 
         cic_flow.basic_flow.first_timestamp = chrono::Utc::now();
-        cic_flow.basic_flow.last_timestamp = chrono::Utc::now() + chrono::Duration::try_seconds(5).unwrap();
+        cic_flow.basic_flow.last_timestamp =
+            chrono::Utc::now() + chrono::Duration::try_seconds(5).unwrap();
 
         cic_flow.basic_flow.bwd_packet_count = 5;