diff --git a/src/meta_plugin.rs b/src/meta_plugin.rs index 1797212..e95fd93 100644 --- a/src/meta_plugin.rs +++ b/src/meta_plugin.rs @@ -17,6 +17,8 @@ use strum::IntoEnumIterator; #[strum(ascii_case_insensitive)] pub enum MetaPluginType { FileMagic, + DigestSha256, + DigestMd5, } pub trait MetaPlugin { @@ -44,12 +46,19 @@ lazy_static! { let program = MetaPluginProgram::new("file", vec!["-bE", "-"], "file_magic".to_string(), true); if program.supported { Some(program) } else { None } } + MetaPluginType::DigestSha256 => None, + MetaPluginType::DigestMd5 => { + let program = MetaPluginProgram::new("md5sum", vec![], "digest_md5".to_string(), false); + if program.supported { Some(program) } else { None } + } }; } pub fn get_meta_plugin(meta_plugin_type: MetaPluginType) -> Box { match meta_plugin_type { MetaPluginType::FileMagic => Box::new(MetaPluginProgram::new("file", vec!["-bE", "-"], "file_magic".to_string(), true)), + MetaPluginType::DigestSha256 => Box::new(DigestSha256MetaPlugin::new()), + MetaPluginType::DigestMd5 => Box::new(MetaPluginProgram::new("md5sum", vec![], "digest_md5".to_string(), false)), } } diff --git a/src/meta_plugin/digest.rs b/src/meta_plugin/digest.rs index e69de29..ab6a88b 100644 --- a/src/meta_plugin/digest.rs +++ b/src/meta_plugin/digest.rs @@ -0,0 +1,60 @@ +use anyhow::Result; +use sha2::{Digest, Sha256}; +use std::io; +use std::io::Write; + +use crate::meta_plugin::MetaPlugin; + +#[derive(Debug, Clone, Default)] +pub struct DigestSha256MetaPlugin { + hasher: Sha256, + meta_name: String, +} + +impl DigestSha256MetaPlugin { + pub fn new() -> DigestSha256MetaPlugin { + DigestSha256MetaPlugin { + hasher: Sha256::new(), + meta_name: "digest_sha256".to_string(), + } + } +} + +impl MetaPlugin for DigestSha256MetaPlugin { + fn create(&self) -> Result> { + // For meta plugins, we don't actually create a writer since we're buffering data internally + // This method is required by the trait but not used in the same way as digest engines + Ok(Box::new(DummyWriter)) + } + + fn finalize(&mut self) -> io::Result { + let result = self.hasher.clone().finalize(); + Ok(format!("{:x}", result)) + } + + fn update(&mut self, data: &[u8]) { + self.hasher.update(data); + } + + fn meta_name(&mut self) -> String { + self.meta_name.clone() + } + + fn is_default(&self) -> bool { + true + } +} + +// Dummy writer that implements Write but doesn't do anything +// This is needed to satisfy the MetaPlugin trait requirements +struct DummyWriter; + +impl Write for DummyWriter { + fn write(&mut self, buf: &[u8]) -> io::Result { + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} diff --git a/src/modes/save.rs b/src/modes/save.rs index 9efae50..a5c20dd 100644 --- a/src/modes/save.rs +++ b/src/modes/save.rs @@ -12,9 +12,8 @@ use std::path::PathBuf; use crate::compression_engine::get_compression_engine; use crate::db::{self}; -use crate::digest_engine::get_digest_engine; use crate::meta_plugin::{MetaPlugin, MetaPluginType, get_meta_plugin}; -use crate::modes::common::{cmd_args_compression_type, cmd_args_digest_type, cmd_args_meta_plugin_types, store_item_digest_value, get_meta_from_env, store_item_meta_value}; +use crate::modes::common::{cmd_args_compression_type, cmd_args_digest_type, cmd_args_meta_plugin_types, get_meta_from_env, store_item_meta_value}; use chrono::Utc; pub fn mode_save( @@ -39,7 +38,18 @@ pub fn mode_save( let digest_type = cmd_args_digest_type(cmd, &args); debug!("MAIN: Digest type: {:?}", digest_type); - let mut digest_engine = get_digest_engine(digest_type.clone()); + + // Convert digest type to meta plugin type + let digest_meta_plugin_type = match digest_type { + crate::digest_engine::DigestType::Sha256 => Some(MetaPluginType::DigestSha256), + crate::digest_engine::DigestType::Md5 => Some(MetaPluginType::DigestMd5), + crate::digest_engine::DigestType::None => None, + }; + + // Add digest meta plugin to the list if needed + if let Some(digest_plugin_type) = digest_meta_plugin_type { + meta_plugin_types.push(digest_plugin_type); + } let compression_type = cmd_args_compression_type(cmd, &args); debug!("MAIN: Compression type: {:?}", compression_type); @@ -146,7 +156,6 @@ pub fn mode_save( stdout.write_all(&buffer[..n])?; item_out.write_all(&buffer[..n])?; - digest_engine.update(&buffer[..n]); for meta_plugin in meta_plugins.iter_mut() { meta_plugin.update(&buffer[..n]); @@ -157,11 +166,6 @@ pub fn mode_save( stdout.flush()?; item_out.flush()?; - let digest = digest_engine.finalize()?; - debug!("DIGEST: {}", digest); - - store_item_digest_value(conn, item.clone(), digest_type, digest)?; - for meta_plugin in meta_plugins.iter_mut() { let meta_name = meta_plugin.meta_name();