diff --git a/Cargo.lock b/Cargo.lock index e74e195..7cd19a0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -829,6 +829,27 @@ dependencies = [ "subtle 1.0.0", ] +[[package]] +name = "csv" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac574ff4d437a7b5ad237ef331c17ccca63c46479e5b5453eb8e10bb99a759fe" +dependencies = [ + "csv-core", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "csv-core" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5efa2b3d7902f4b634a20cae3c9c4e6209dc4779feb6863329607560143efa70" +dependencies = [ + "memchr", +] + [[package]] name = "ctor" version = "0.2.5" @@ -1982,6 +2003,7 @@ dependencies = [ "chrono", "cidr-utils", "clap", + "csv", "ctor", "ctrlc", "env_logger", diff --git a/Cargo.toml b/Cargo.toml index fb9b1d3..761e43e 100644 --- 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 index 146f027..6ae8089 100644 --- 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 { + serde_json::to_string(self).map_err(|e| e.to_string()) + } + + fn to_text(&self) -> Result { + let data = self + .data + .keys() + .map(|k| format!("{}={}", k, self.data.get(k).unwrap())) + .collect::>() + .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 { + 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::>() + .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::>() - .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()) } }