diff --git a/src/meta_plugin.rs b/src/meta_plugin.rs index c41bd56..0e8b9f8 100644 --- a/src/meta_plugin.rs +++ b/src/meta_plugin.rs @@ -6,11 +6,13 @@ pub mod program; pub mod digest; pub mod system; pub mod magic; +pub mod binary; use crate::meta_plugin::program::MetaPluginProgram; use crate::meta_plugin::digest::{DigestSha256MetaPlugin, ReadTimeMetaPlugin, ReadRateMetaPlugin}; -use crate::meta_plugin::system::{CwdMetaPlugin, BinaryMetaPlugin, UidMetaPlugin, UserMetaPlugin, GidMetaPlugin, GroupMetaPlugin, ShellMetaPlugin, ShellPidMetaPlugin, KeepPidMetaPlugin, HostnameMetaPlugin, FullHostnameMetaPlugin}; +use crate::meta_plugin::system::{CwdMetaPlugin, UidMetaPlugin, UserMetaPlugin, GidMetaPlugin, GroupMetaPlugin, ShellMetaPlugin, ShellPidMetaPlugin, KeepPidMetaPlugin, HostnameMetaPlugin, FullHostnameMetaPlugin}; use crate::meta_plugin::magic::MagicFileMetaPlugin; +use crate::meta_plugin::binary::BinaryMetaPlugin; #[derive(Debug, Eq, PartialEq, Clone, strum::EnumIter, strum::Display, strum::EnumString)] #[strum(ascii_case_insensitive)] diff --git a/src/meta_plugin/binary.rs b/src/meta_plugin/binary.rs index e69de29..0e600df 100644 --- a/src/meta_plugin/binary.rs +++ b/src/meta_plugin/binary.rs @@ -0,0 +1,92 @@ +use anyhow::Result; +use rusqlite::Connection; + +use crate::common::is_binary::is_binary; +use crate::meta_plugin::{MetaPlugin, output_metadata}; + +#[derive(Debug, Clone, Default)] +pub struct BinaryMetaPlugin { + meta_name: String, + buffer: Vec, + max_buffer_size: usize, + is_saved: bool, + item_id: Option, + output_names: std::collections::HashMap, +} + +impl BinaryMetaPlugin { + pub fn new() -> BinaryMetaPlugin { + BinaryMetaPlugin { + meta_name: "binary".to_string(), + buffer: Vec::new(), + max_buffer_size: 4096, // 4KB + is_saved: false, + item_id: None, + output_names: std::collections::HashMap::new(), + } + } + +} + +impl MetaPlugin for BinaryMetaPlugin { + fn is_internal(&self) -> bool { + true + } + + fn finalize(&mut self, conn: &Connection) -> Result<()> { + // Save the binary detection result when finalizing, after all data has been collected + if !self.is_saved { + if let Some(item_id) = self.item_id { + let is_binary = is_binary(&self.buffer); + let value = if is_binary { "true".to_string() } else { "false".to_string() }; + + // Save to database immediately using central output handler + let _ = output_metadata(conn, item_id, &self.meta_name, value, &self.output_names); + + self.is_saved = true; + } + } + Ok(()) + } + + fn update(&mut self, data: &[u8], _conn: &Connection) { + // Only collect up to max_buffer_size + let remaining_capacity = self.max_buffer_size.saturating_sub(self.buffer.len()); + if remaining_capacity > 0 { + let bytes_to_copy = std::cmp::min(data.len(), remaining_capacity); + self.buffer.extend_from_slice(&data[..bytes_to_copy]); + } + } + + fn meta_name(&mut self) -> String { + self.meta_name.clone() + } + + fn initialize(&mut self, _conn: &Connection, item_id: i64) -> Result<()> { + self.item_id = Some(item_id); + Ok(()) + } + + fn configure(&mut self, options: &std::collections::HashMap) -> Result<()> { + if let Some(max_buffer_size) = options.get("max_buffer_size") { + if let Some(size) = max_buffer_size.as_u64() { + self.max_buffer_size = size as usize; + } + } + + if let Some(outputs) = options.get("outputs") { + if let Some(outputs_map) = outputs.as_mapping() { + for (key, value) in outputs_map { + if let (Some(key_str), Some(value_str)) = (key.as_str(), value.as_str()) { + self.output_names.insert(key_str.to_string(), value_str.to_string()); + } + } + } + } + Ok(()) + } + + fn get_output_name(&self, default_name: &str) -> String { + self.output_names.get(default_name).cloned().unwrap_or_else(|| default_name.to_string()) + } +} diff --git a/src/meta_plugin/system.rs b/src/meta_plugin/system.rs index 89f456a..f38967f 100644 --- a/src/meta_plugin/system.rs +++ b/src/meta_plugin/system.rs @@ -7,7 +7,6 @@ use std::process; use uzers::{get_current_uid, get_current_gid, get_current_username, get_current_groupname}; use rusqlite::Connection; -use crate::common::is_binary::is_binary; use crate::meta_plugin::{MetaPlugin, output_metadata}; #[derive(Debug, Clone, Default)] @@ -17,93 +16,6 @@ pub struct CwdMetaPlugin { output_names: std::collections::HashMap, } -#[derive(Debug, Clone, Default)] -pub struct BinaryMetaPlugin { - meta_name: String, - buffer: Vec, - max_buffer_size: usize, - is_saved: bool, - item_id: Option, - output_names: std::collections::HashMap, -} - -impl BinaryMetaPlugin { - pub fn new() -> BinaryMetaPlugin { - BinaryMetaPlugin { - meta_name: "binary".to_string(), - buffer: Vec::new(), - max_buffer_size: 4096, // 4KB - is_saved: false, - item_id: None, - output_names: std::collections::HashMap::new(), - } - } - -} - -impl MetaPlugin for BinaryMetaPlugin { - fn is_internal(&self) -> bool { - true - } - - fn finalize(&mut self, conn: &Connection) -> Result<()> { - // Save the binary detection result when finalizing, after all data has been collected - if !self.is_saved { - if let Some(item_id) = self.item_id { - let is_binary = is_binary(&self.buffer); - let value = if is_binary { "true".to_string() } else { "false".to_string() }; - - // Save to database immediately using central output handler - let _ = output_metadata(conn, item_id, &self.meta_name, value, &self.output_names); - - self.is_saved = true; - } - } - Ok(()) - } - - fn update(&mut self, data: &[u8], _conn: &Connection) { - // Only collect up to max_buffer_size - let remaining_capacity = self.max_buffer_size.saturating_sub(self.buffer.len()); - if remaining_capacity > 0 { - let bytes_to_copy = std::cmp::min(data.len(), remaining_capacity); - self.buffer.extend_from_slice(&data[..bytes_to_copy]); - } - } - - fn meta_name(&mut self) -> String { - self.meta_name.clone() - } - - fn initialize(&mut self, _conn: &Connection, item_id: i64) -> Result<()> { - self.item_id = Some(item_id); - Ok(()) - } - - fn configure(&mut self, options: &std::collections::HashMap) -> Result<()> { - if let Some(max_buffer_size) = options.get("max_buffer_size") { - if let Some(size) = max_buffer_size.as_u64() { - self.max_buffer_size = size as usize; - } - } - - if let Some(outputs) = options.get("outputs") { - if let Some(outputs_map) = outputs.as_mapping() { - for (key, value) in outputs_map { - if let (Some(key_str), Some(value_str)) = (key.as_str(), value.as_str()) { - self.output_names.insert(key_str.to_string(), value_str.to_string()); - } - } - } - } - Ok(()) - } - - fn get_output_name(&self, default_name: &str) -> String { - self.output_names.get(default_name).cloned().unwrap_or_else(|| default_name.to_string()) - } -} - impl CwdMetaPlugin { pub fn new() -> CwdMetaPlugin { CwdMetaPlugin {