new: implemented csv output format

evilsocket · Nov 14, 2023 · 26539c9 · 26539c9
1 parent c5f6c71
commit 26539c9
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 23 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -77,6 +77,7 @@ sibyl = { version = "0.6.16", optional = true, features = [
 rdp-rs = { version = "0.1.0", optional = true }
 scylla = { version = "0.10.1", optional = true }
 paho-mqtt = { version = "0.12.3", optional = true }
+csv = "1.3.0"
 
 [dev-dependencies]
 tempfile = "3.8.0"

diff --git a/src/session/loot.rs b/src/session/loot.rs
@@ -1,6 +1,6 @@
-use std::fmt;
 use std::fs::OpenOptions;
 use std::io::prelude::*;
+use std::{fmt, path::Path};
 
 use ansi_term::Colour;
 use chrono::{DateTime, Local};
@@ -14,6 +14,7 @@ use crate::session::Error;
 pub(crate) enum OutputFormat {
     #[default]
     Text,
+    CSV,
     JSONL,
 }
 
@@ -59,29 +60,57 @@ impl Loot {
         self.found_at.format("%Y-%m-%d %H:%M:%S").to_string()
     }
 
+    fn to_json(&self) -> Result<String, Error> {
+        serde_json::to_string(self).map_err(|e| e.to_string())
+    }
+
+    fn to_text(&self) -> Result<String, Error> {
+        let data = self
+            .data
+            .keys()
+            .map(|k| format!("{}={}", k, self.data.get(k).unwrap()))
+            .collect::<Vec<String>>()
+            .join("\t");
+
+        Ok(if self.target.is_empty() {
+            format!("[{}] ({}) {}", self.found_at_string(), &self.plugin, data)
+        } else {
+            format!(
+                "[{}] ({}) <{}> {}",
+                self.found_at_string(),
+                &self.plugin,
+                &self.target,
+                data
+            )
+        })
+    }
+
+    fn to_csv(&self, path: &str) -> Result<String, Error> {
+        let mut wtr = csv::Writer::from_writer(vec![]);
+
+        if !Path::new(path).exists() {
+            wtr.write_record(&["found_at", "plugin", "target", "data"])
+                .map_err(|e| e.to_string())?;
+        }
+
+        let data = self
+            .data
+            .keys()
+            .map(|k| format!("{}={}", k, self.data.get(k).unwrap()))
+            .collect::<Vec<String>>()
+            .join(";");
+
+        wtr.write_record(&[&self.found_at_string(), &self.plugin, &self.target, &data])
+            .map_err(|e| e.to_string())?;
+
+        String::from_utf8(wtr.into_inner().unwrap()).map_err(|e| e.to_string())
+    }
+
     pub fn append_to_file(&self, path: &str, format: &OutputFormat) -> Result<(), Error> {
         let data = match format {
-            OutputFormat::JSONL => serde_json::to_string(self).map_err(|e| e.to_string())?,
-            OutputFormat::Text => {
-                let data = self
-                    .data
-                    .keys()
-                    .map(|k| format!("{}={}", k, self.data.get(k).unwrap()))
-                    .collect::<Vec<String>>()
-                    .join("\t");
-
-                if self.target.is_empty() {
-                    format!("[{}] ({}) {}", self.found_at_string(), &self.plugin, data)
-                } else {
-                    format!(
-                        "[{}] ({}) <{}> {}",
-                        self.found_at_string(),
-                        &self.plugin,
-                        &self.target,
-                        data
-                    )
-                }
-            }
+            OutputFormat::JSONL => self.to_json()?,
+            OutputFormat::Text => self.to_text()?,
+            OutputFormat::CSV => self.to_csv(path)?,
         };
 
         let mut file = OpenOptions::new()
@@ -91,7 +120,7 @@ impl Loot {
             .open(path)
             .map_err(|e| e.to_string())?;
 
-        writeln!(file, "{}", data).map_err(|e| e.to_string())
+        writeln!(file, "{}", data.trim()).map_err(|e| e.to_string())
     }
 }