diff --git a/Cargo.toml b/Cargo.toml index 8dfc001..bf33aa8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -44,6 +44,7 @@ serde = { version = "1.0.219", features = ["derive"] } serde_json = "1.0.142" serde_yaml = "0.9.34" sha2 = "0.10.0" +md5 = "0.7.0" stderrlog = "0.6.0" strum = { version = "0.27.2", features = ["derive"] } strum_macros = "0.27.2" diff --git a/src/meta_plugin/digest.rs b/src/meta_plugin/digest.rs index 92aa354..931746d 100644 --- a/src/meta_plugin/digest.rs +++ b/src/meta_plugin/digest.rs @@ -1,56 +1,116 @@ -use sha2::{Digest, Sha256}; - +use sha2::{Digest, Sha256, Sha512}; +use md5::Md5; use crate::meta_plugin::MetaPlugin; -#[derive(Debug, Clone, Default)] -pub struct DigestSha256MetaPlugin { - hasher: Sha256, +#[derive(Debug, Clone)] +enum HashMethod { + Md5, + Sha256, + Sha512, +} + +impl Default for HashMethod { + fn default() -> Self { + HashMethod::Sha256 + } +} + +impl HashMethod { + fn from_str(s: &str) -> Option { + match s { + "md5" => Some(HashMethod::Md5), + "sha256" => Some(HashMethod::Sha256), + "sha512" => Some(HashMethod::Sha512), + _ => None, + } + } +} + +#[derive(Debug, Clone)] +pub struct DigestMetaPlugin { + md5_hasher: Option, + sha256_hasher: Option, + sha512_hasher: Option, is_finalized: bool, meta_name: String, outputs: std::collections::HashMap, options: std::collections::HashMap, + methods: Vec, } -impl DigestSha256MetaPlugin { - pub fn new( - _options: Option>, - outputs: Option>, - ) -> DigestSha256MetaPlugin { - // Start with default options - let mut final_options = std::collections::HashMap::new(); - if let Some(opts) = _options { - for (key, value) in opts { - final_options.insert(key, value); - } - } - - // Start with default outputs - let mut final_outputs = std::collections::HashMap::new(); - let default_outputs = vec!["digest_sha256".to_string()]; - for output_name in default_outputs { - final_outputs.insert(output_name.clone(), serde_yaml::Value::String(output_name)); - } - if let Some(outs) = outputs { - for (key, value) in outs { - final_outputs.insert(key, value); - } - } - - DigestSha256MetaPlugin { - hasher: Sha256::new(), +impl Default for DigestMetaPlugin { + fn default() -> Self { + Self { + md5_hasher: None, + sha256_hasher: None, + sha512_hasher: None, is_finalized: false, - meta_name: "digest_sha256".to_string(), - outputs: final_outputs, - options: final_options, + meta_name: "digest".to_string(), + outputs: std::collections::HashMap::new(), + options: std::collections::HashMap::new(), + methods: vec![HashMethod::Sha256], } } +} + +impl DigestMetaPlugin { + pub fn new( + options: Option>, + outputs: Option>, + ) -> DigestMetaPlugin { + let mut plugin = DigestMetaPlugin::default(); + + // Set default outputs + let default_outputs = vec![ + ("digest_md5".to_string(), serde_yaml::Value::String("digest_md5".to_string())), + ("digest_sha256".to_string(), serde_yaml::Value::String("digest_sha256".to_string())), + ("digest_sha512".to_string(), serde_yaml::Value::String("digest_sha512".to_string())), + ]; + for (key, value) in default_outputs { + plugin.outputs.insert(key, value); + } + + // Apply provided options + if let Some(opts) = options { + for (key, value) in opts { + plugin.options.insert(key, value); + } + } + + // Apply provided outputs + if let Some(outs) = outputs { + for (key, value) in outs { + plugin.outputs.insert(key, value); + } + } + + // Configure methods based on options + if let Some(method_value) = plugin.options.get("method") { + if let Some(method_str) = method_value.as_str() { + if let Some(method) = HashMethod::from_str(method_str) { + plugin.methods = vec![method]; + } + } + } + + // Initialize hashers based on selected methods + for method in &plugin.methods { + match method { + HashMethod::Md5 => plugin.md5_hasher = Some(Md5::new()), + HashMethod::Sha256 => plugin.sha256_hasher = Some(Sha256::new()), + HashMethod::Sha512 => plugin.sha512_hasher = Some(Sha512::new()), + } + } + + plugin + } - pub fn new_simple() -> DigestSha256MetaPlugin { + pub fn new_simple() -> DigestMetaPlugin { Self::new(None, None) } } -impl MetaPlugin for DigestSha256MetaPlugin { +impl MetaPlugin for DigestMetaPlugin { fn is_finalized(&self) -> bool { self.is_finalized } @@ -67,7 +127,6 @@ impl MetaPlugin for DigestSha256MetaPlugin { } fn finalize(&mut self) -> crate::meta_plugin::MetaPluginResponse { - // If already finalized, don't process again if self.is_finalized { return crate::meta_plugin::MetaPluginResponse { metadata: Vec::new(), @@ -77,22 +136,52 @@ impl MetaPlugin for DigestSha256MetaPlugin { let mut metadata = Vec::new(); - // Finalize the hash - let hash_result = self.hasher.finalize_reset(); - let hex_string = format!("{:x}", hash_result); - - // Use process_metadata_outputs to handle output mapping - if let Some(meta_data) = crate::meta_plugin::process_metadata_outputs( - "digest_sha256", - hex_string, - &self.outputs - ) { - metadata.push(meta_data); + // Process each method + for method in &self.methods { + match method { + HashMethod::Md5 => { + if let Some(hasher) = self.md5_hasher.take() { + let hash_result = hasher.finalize(); + let hex_string = format!("{:x}", hash_result); + if let Some(meta_data) = crate::meta_plugin::process_metadata_outputs( + "digest_md5", + hex_string, + &self.outputs + ) { + metadata.push(meta_data); + } + } + } + HashMethod::Sha256 => { + if let Some(hasher) = self.sha256_hasher.take() { + let hash_result = hasher.finalize_reset(); + let hex_string = format!("{:x}", hash_result); + if let Some(meta_data) = crate::meta_plugin::process_metadata_outputs( + "digest_sha256", + hex_string, + &self.outputs + ) { + metadata.push(meta_data); + } + } + } + HashMethod::Sha512 => { + if let Some(hasher) = self.sha512_hasher.take() { + let hash_result = hasher.finalize_reset(); + let hex_string = format!("{:x}", hash_result); + if let Some(meta_data) = crate::meta_plugin::process_metadata_outputs( + "digest_sha512", + hex_string, + &self.outputs + ) { + metadata.push(meta_data); + } + } + } + } } - // Mark as finalized self.is_finalized = true; - crate::meta_plugin::MetaPluginResponse { metadata, is_finalized: true, @@ -100,7 +189,6 @@ impl MetaPlugin for DigestSha256MetaPlugin { } fn update(&mut self, data: &[u8]) -> crate::meta_plugin::MetaPluginResponse { - // If already finalized, don't process more data if self.is_finalized { return crate::meta_plugin::MetaPluginResponse { metadata: Vec::new(), @@ -108,7 +196,26 @@ impl MetaPlugin for DigestSha256MetaPlugin { }; } - self.hasher.update(data); + for method in &self.methods { + match method { + HashMethod::Md5 => { + if let Some(hasher) = &mut self.md5_hasher { + hasher.update(data); + } + } + HashMethod::Sha256 => { + if let Some(hasher) = &mut self.sha256_hasher { + hasher.update(data); + } + } + HashMethod::Sha512 => { + if let Some(hasher) = &mut self.sha512_hasher { + hasher.update(data); + } + } + } + } + crate::meta_plugin::MetaPluginResponse { metadata: Vec::new(), is_finalized: false, @@ -128,11 +235,17 @@ impl MetaPlugin for DigestSha256MetaPlugin { } fn default_outputs(&self) -> Vec { - vec!["digest_sha256".to_string()] + vec![ + "digest_md5".to_string(), + "digest_sha256".to_string(), + "digest_sha512".to_string(), + ] } fn default_options(&self) -> std::collections::HashMap { - std::collections::HashMap::new() + let mut options = std::collections::HashMap::new(); + options.insert("method".to_string(), serde_yaml::Value::String("sha256".to_string())); + options } fn options(&self) -> &std::collections::HashMap { @@ -144,4 +257,3 @@ impl MetaPlugin for DigestSha256MetaPlugin { } } - diff --git a/src/meta_plugin/mod.rs b/src/meta_plugin/mod.rs index 2c0f6c1..0b54e38 100644 --- a/src/meta_plugin/mod.rs +++ b/src/meta_plugin/mod.rs @@ -296,7 +296,7 @@ pub fn get_meta_plugin(meta_plugin_type: MetaPluginType) -> Box MetaPluginType::Shell => Box::new(ShellMetaPlugin::new_simple()), MetaPluginType::ShellPid => Box::new(ShellPidMetaPlugin::new_simple()), MetaPluginType::KeepPid => Box::new(KeepPidMetaPlugin::new_simple()), - MetaPluginType::DigestSha256 => Box::new(DigestSha256MetaPlugin::new_simple()), + MetaPluginType::DigestSha256 => Box::new(DigestMetaPlugin::new_simple()), MetaPluginType::DigestMd5 => Box::new(MetaPluginProgram::new_simple("md5sum", vec![], "digest_md5".to_string(), true)), MetaPluginType::ReadTime => Box::new(ReadTimeMetaPlugin::new_simple()), MetaPluginType::ReadRate => Box::new(ReadRateMetaPlugin::new_simple()),