keep/src/meta_plugin/digest.rs

use crate::meta_plugin::{BaseMetaPlugin, MetaPlugin, MetaPluginType};
use md5;
use sha2::{Digest, Sha256, Sha512};
use std::io::Write;

#[derive(Clone)]
enum Hasher {
    Sha256(Sha256),
    Md5(md5::Context),
    Sha512(Sha512),
}

impl Default for Hasher {
    fn default() -> Self {
        Hasher::Sha256(Sha256::default())
    }
}

// Manual Debug implementation to avoid md5::Context not implementing Debug
impl std::fmt::Debug for Hasher {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Hasher::Sha256(_) => write!(f, "Hasher::Sha256"),
            Hasher::Md5(_) => write!(f, "Hasher::Md5"),
            Hasher::Sha512(_) => write!(f, "Hasher::Sha512"),
        }
    }
}

impl Hasher {
    fn update(&mut self, data: &[u8]) {
        match self {
            Hasher::Sha256(hasher) => hasher.update(data),
            Hasher::Md5(hasher) => {
                let _ = hasher.write(data);
            }
            Hasher::Sha512(hasher) => hasher.update(data),
        }
    }

    fn finalize(&mut self) -> String {
        match self {
            Hasher::Sha256(hasher) => {
                let result = std::mem::replace(hasher, Sha256::new()).finalize_reset();
                format!("{result:x}")
            }
            Hasher::Md5(hasher) => {
                let result = hasher.clone().compute();
                format!("{result:x}")
            }
            Hasher::Sha512(hasher) => {
                let result = std::mem::replace(hasher, Sha512::new()).finalize_reset();
                format!("{result:x}")
            }
        }
    }

    fn output_name(&self) -> &'static str {
        match self {
            Hasher::Sha256(_) => "digest_sha256",
            Hasher::Md5(_) => "digest_md5",
            Hasher::Sha512(_) => "digest_sha512",
        }
    }
}

#[derive(Debug, Default)]
pub struct DigestMetaPlugin {
    hasher: Option<Hasher>,
    is_finalized: bool,
    base: BaseMetaPlugin,
}

impl DigestMetaPlugin {
    pub fn new(
        options: Option<std::collections::HashMap<String, serde_yaml::Value>>,
        outputs: Option<std::collections::HashMap<String, serde_yaml::Value>>,
    ) -> DigestMetaPlugin {
        let mut base = BaseMetaPlugin::new();

        // Apply provided options
        if let Some(opts) = options {
            for (key, value) in opts {
                base.options.insert(key, value);
            }
        }

        // Determine the selected method
        let method = if let Some(method_value) = base.options.get("method") {
            if let Some(method_str) = method_value.as_str() {
                match method_str {
                    "md5" => "md5",
                    "sha256" => "sha256",
                    "sha512" => "sha512",
                    _ => "sha256",
                }
            } else {
                "sha256"
            }
        } else {
            "sha256"
        };

        // Initialize the hasher based on the method
        let hasher = match method {
            "md5" => Some(Hasher::Md5(md5::Context::new())),
            "sha256" => Some(Hasher::Sha256(Sha256::new())),
            "sha512" => Some(Hasher::Sha512(Sha512::new())),
            _ => Some(Hasher::Sha256(Sha256::new())),
        };

        // Add the method to options so it shows up in the status
        base.options.insert(
            "method".to_string(),
            serde_yaml::Value::String(method.to_string()),
        );

        // Set outputs based on the selected hash method
        // Only the selected method's output should be enabled, others should be None
        let all_outputs = vec!["digest_md5", "digest_sha256", "digest_sha512"];
        for output_name in &all_outputs {
            if output_name == &format!("digest_{method}") {
                base.outputs.insert(
                    output_name.to_string(),
                    serde_yaml::Value::String(output_name.to_string()),
                );
            } else {
                base.outputs
                    .insert(output_name.to_string(), serde_yaml::Value::Null);
            }
        }

        // Apply provided outputs, but only for enabled outputs
        if let Some(outs) = outputs {
            for (key, value) in outs {
                // Only update if the output is not disabled (not None)
                if let Some(current_value) = base.outputs.get_mut(&key)
                    && !current_value.is_null()
                {
                    *current_value = value;
                }
            }
        }

        DigestMetaPlugin {
            hasher,
            is_finalized: false,
            base,
        }
    }
}

impl MetaPlugin for DigestMetaPlugin {
    fn is_finalized(&self) -> bool {
        self.is_finalized
    }

    fn set_finalized(&mut self, finalized: bool) {
        self.is_finalized = finalized;
    }

    fn initialize(&mut self) -> crate::meta_plugin::MetaPluginResponse {
        crate::meta_plugin::MetaPluginResponse {
            metadata: Vec::new(),
            is_finalized: false,
        }
    }

    fn finalize(&mut self) -> crate::meta_plugin::MetaPluginResponse {
        if self.is_finalized {
            return crate::meta_plugin::MetaPluginResponse {
                metadata: Vec::new(),
                is_finalized: true,
            };
        }

        let mut metadata = Vec::new();

        // Update outputs based on the selected hash method
        if let Some(hasher) = &mut self.hasher {
            let hash_value = hasher.finalize();
            let output_name = hasher.output_name();

            // Use process_metadata_outputs to handle output mapping
            if let Some(meta_data) = crate::meta_plugin::process_metadata_outputs(
                output_name,
                serde_yaml::Value::String(hash_value),
                self.base.outputs(),
            ) {
                metadata.push(meta_data);
            }

            // Set all other digest outputs to None
            let all_outputs = vec!["digest_md5", "digest_sha256", "digest_sha512"];
            for output_name in all_outputs {
                if output_name != hasher.output_name() {
                    self.base
                        .outputs
                        .insert(output_name.to_string(), serde_yaml::Value::Null);
                }
            }
        }

        self.is_finalized = true;
        crate::meta_plugin::MetaPluginResponse {
            metadata,
            is_finalized: true,
        }
    }

    fn update(&mut self, data: &[u8]) -> crate::meta_plugin::MetaPluginResponse {
        if self.is_finalized {
            return crate::meta_plugin::MetaPluginResponse {
                metadata: Vec::new(),
                is_finalized: true,
            };
        }

        // Update the active hasher
        if let Some(hasher) = &mut self.hasher {
            hasher.update(data);
        }

        crate::meta_plugin::MetaPluginResponse {
            metadata: Vec::new(),
            is_finalized: false,
        }
    }

    fn meta_type(&self) -> MetaPluginType {
        MetaPluginType::Digest
    }

    fn outputs(&self) -> &std::collections::HashMap<String, serde_yaml::Value> {
        self.base.outputs()
    }

    fn outputs_mut(
        &mut self,
    ) -> anyhow::Result<&mut std::collections::HashMap<String, serde_yaml::Value>> {
        Ok(self.base.outputs_mut())
    }

    fn default_outputs(&self) -> Vec<String> {
        vec![
            "digest_md5".to_string(),
            "digest_sha256".to_string(),
            "digest_sha512".to_string(),
        ]
    }

    fn options(&self) -> &std::collections::HashMap<String, serde_yaml::Value> {
        self.base.options()
    }

    fn options_mut(
        &mut self,
    ) -> anyhow::Result<&mut std::collections::HashMap<String, serde_yaml::Value>> {
        Ok(self.base.options_mut())
    }

    fn parallel_safe(&self) -> bool {
        true
    }
}

use crate::meta_plugin::register_meta_plugin;

// Register the plugin at module initialization time
#[ctor::ctor]
fn register_digest_plugin() {
    register_meta_plugin(MetaPluginType::Digest, |options, outputs| {
        Box::new(DigestMetaPlugin::new(options, outputs))
    })
    .expect("Failed to register DigestMetaPlugin");
}