diff --git a/kunai/src/bin/main.rs b/kunai/src/bin/main.rs index 3724e3f..e356b84 100644 --- a/kunai/src/bin/main.rs +++ b/kunai/src/bin/main.rs @@ -11,6 +11,7 @@ use env_logger::Builder; use fs_walk::WalkOptions; use gene::rules::MAX_SEVERITY; use gene::{Compiler, Engine}; +use huby::ByteSize; use kunai::containers::Container; use kunai::events::{ BpfProgLoadData, BpfProgTypeInfo, BpfSocketFilterData, CloneData, ConnectData, DnsQueryData, @@ -1575,29 +1576,32 @@ impl<'s> EventConsumer<'s> { }; } - // we collect a vector of ioc matching - let matching_iocs = event - .iocs() - .iter() - .filter(|ioc| self.iocs.contains_key(&ioc.to_string())) - .map(|ioc| ioc.to_string()) - .collect::>(); - - let ioc_severity: usize = matching_iocs - .iter() - .map(|ioc| self.iocs.get(ioc).map(|e| *e as usize).unwrap_or_default()) - .sum(); - - if !matching_iocs.is_empty() { - // we create a new ScanResult if necessary - if scan_result.is_none() { - scan_result = Some(ScanResult::default()); - } + // no need to scan for IoC if not necessary + if !self.iocs.is_empty() { + // we collect a vector of ioc matching + let matching_iocs = event + .iocs() + .iter() + .filter(|ioc| self.iocs.contains_key(&ioc.to_string())) + .map(|ioc| ioc.to_string()) + .collect::>(); + + let ioc_severity: usize = matching_iocs + .iter() + .map(|ioc| self.iocs.get(ioc).map(|e| *e as usize).unwrap_or_default()) + .sum(); + + if !matching_iocs.is_empty() { + // we create a new ScanResult if necessary + if scan_result.is_none() { + scan_result = Some(ScanResult::default()); + } - // we add ioc matching to the list of matching rules - if let Some(sr) = scan_result.as_mut() { - sr.iocs = matching_iocs; - sr.severity = ioc_severity.clamp(0, MAX_SEVERITY as usize) as u8; + // we add ioc matching to the list of matching rules + if let Some(sr) = scan_result.as_mut() { + sr.iocs = matching_iocs; + sr.severity = ioc_severity.clamp(0, MAX_SEVERITY as usize) as u8; + } } } @@ -2441,6 +2445,10 @@ struct ReplayOpt { #[arg(short, long, value_name = "FILE")] config: Option, + /// Benchmark scanning engine + #[arg(long)] + bench: bool, + /// Detection/filtering rule file. Supersedes configuration file. #[arg(short, long, value_name = "FILE")] rule_file: Option>, @@ -2693,10 +2701,20 @@ enum Command { Logs(LogsOpt), } +fn time_it(mut f: F) -> Duration { + let start_time = std::time::Instant::now(); + f(); + let end_time = std::time::Instant::now(); + end_time - start_time +} + impl Command { fn replay(o: ReplayOpt) -> anyhow::Result<()> { + let bench = o.bench; let log_files = o.log_files.clone(); let conf: Config = o.try_into()?; + let mut kunai_scan_time = Duration::new(0, 0); + let mut data_size = ByteSize::from_bytes(0); let mut p = EventConsumer::with_config(conf.stdout_output())?; for f in log_files { @@ -2709,6 +2727,11 @@ impl Command { let mut de = serde_json::Deserializer::from_reader(reader); while let Ok(v) = serde_json::Value::deserialize(&mut de) { + // we need to know the size of the data we scan to compute throughput + if bench { + data_size += ByteSize::from_bytes(serde_json::to_string(&v)?.len() as u64); + } + // we attempt at getting event name from json if let Some(name) = v .get("info") @@ -2719,7 +2742,13 @@ impl Command { macro_rules! scan_event { ($scanner:expr, $into:ty) => {{ let mut e: UserEvent<$into> = serde_json::from_value(v)?; - $scanner.scan_and_print(&mut e); + if bench { + kunai_scan_time += time_it(|| { + let _ = $scanner.scan(&mut e); + }); + } else { + $scanner.scan_and_print(&mut e); + } }}; } @@ -2767,6 +2796,14 @@ impl Command { } } + if bench { + let throughput = ByteSize::from_bytes( + (data_size.in_bytes() as f64 / kunai_scan_time.as_secs_f64()) as u64, + ); + println!("scan duration: {kunai_scan_time:?}"); + println!("scan throughput: {throughput}/s"); + } + Ok(()) }