maec-rs 0.1.0

//! MAEC Malware Family object implementation
//!
//! A Malware Family is a set of malware instances that are related by common
//! authorship and/or lineage.

use serde::{Deserialize, Serialize};

use crate::common::{ExternalReference, MaecObject};
use crate::error::{MaecError, Result};
use crate::objects::types::{FieldData, Name};
use crate::Capability;
use chrono::{DateTime, Utc};

/// MAEC Malware Family
///
/// A set of malware instances that are related by common authorship and/or lineage.
/// Malware Families are often named and may have components such as strings that
/// are common across all members of the family.
///
/// # Examples
///
/// ```
/// use maec::{MalwareFamily, Name};
///
/// let family = MalwareFamily::builder()
///     .name(Name::new("WannaCry"))
///     .description("Ransomware family first seen in May 2017")
///     .build()
///     .unwrap();
///
/// assert_eq!(family.name.value, "WannaCry");
/// ```
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub struct MalwareFamily {
    /// Common MAEC properties
    #[serde(flatten)]
    pub common: crate::common::CommonProperties,

    /// Name of the malware family
    pub name: Name,

    /// Alternative names/aliases
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub aliases: Vec<Name>,

    /// Labels describing the family (e.g., "worm", "ransomware")
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub labels: Vec<String>,

    /// Textual description
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,

    /// Field data (delivery vectors, timestamps)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub field_data: Option<FieldData>,

    /// Strings common to all family members
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub common_strings: Vec<String>,

    /// Capabilities common to all family members
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub common_capabilities: Vec<Capability>,

    /// References to common code artifacts (STIX artifact IDs)
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub common_code_refs: Vec<String>,

    /// References to common behavior IDs
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub common_behavior_refs: Vec<String>,

    /// External references (ATT&CK, research papers, etc.)
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub references: Vec<ExternalReference>,
}

impl MalwareFamily {
    /// Creates a new MalwareFamily builder
    pub fn builder() -> MalwareFamilyBuilder {
        MalwareFamilyBuilder::default()
    }

    /// Creates a minimal MalwareFamily with just a name
    pub fn new(name: impl Into<Name>) -> Self {
        Self {
            common: crate::common::CommonProperties::new("malware-family", None),
            name: name.into(),
            aliases: vec![],
            labels: vec![],
            description: None,
            field_data: None,
            common_strings: vec![],
            common_capabilities: vec![],
            common_code_refs: vec![],
            common_behavior_refs: vec![],
            references: vec![],
        }
    }

    /// Validates the MalwareFamily structure
    pub fn validate(&self) -> Result<()> {
        if self.common.r#type != "malware-family" {
            return Err(MaecError::ValidationError(format!(
                "type must be 'malware-family', got '{}'",
                self.common.r#type
            )));
        }

        if !crate::common::is_valid_maec_id(&self.common.id) {
            return Err(MaecError::InvalidId(self.common.id.clone()));
        }

        Ok(())
    }
}

impl MaecObject for MalwareFamily {
    fn id(&self) -> &str {
        &self.common.id
    }

    fn type_(&self) -> &str {
        &self.common.r#type
    }

    fn created(&self) -> DateTime<Utc> {
        self.common.created
    }
}

/// Builder for MalwareFamily objects
#[derive(Debug, Default)]
pub struct MalwareFamilyBuilder {
    id: Option<String>,
    name: Option<Name>,
    aliases: Vec<Name>,
    labels: Vec<String>,
    description: Option<String>,
    field_data: Option<FieldData>,
    common_strings: Vec<String>,
    common_capabilities: Vec<Capability>,
    common_code_refs: Vec<String>,
    common_behavior_refs: Vec<String>,
    references: Vec<ExternalReference>,
}

impl MalwareFamilyBuilder {
    /// Sets a custom ID
    pub fn id(mut self, id: impl Into<String>) -> Self {
        self.id = Some(id.into());
        self
    }

    /// Sets the family name (required)
    pub fn name(mut self, name: impl Into<Name>) -> Self {
        self.name = Some(name.into());
        self
    }

    /// Adds an alias
    pub fn add_alias(mut self, alias: impl Into<Name>) -> Self {
        self.aliases.push(alias.into());
        self
    }

    /// Sets all aliases at once
    pub fn aliases(mut self, aliases: Vec<Name>) -> Self {
        self.aliases = aliases;
        self
    }

    /// Adds a label
    pub fn add_label(mut self, label: impl Into<String>) -> Self {
        self.labels.push(label.into());
        self
    }

    /// Sets all labels at once
    pub fn labels(mut self, labels: Vec<String>) -> Self {
        self.labels = labels;
        self
    }

    /// Sets the description
    pub fn description(mut self, desc: impl Into<String>) -> Self {
        self.description = Some(desc.into());
        self
    }

    /// Sets field data
    pub fn field_data(mut self, field_data: FieldData) -> Self {
        self.field_data = Some(field_data);
        self
    }

    /// Adds a common string
    pub fn add_common_string(mut self, s: impl Into<String>) -> Self {
        self.common_strings.push(s.into());
        self
    }

    /// Sets all common strings
    pub fn common_strings(mut self, strings: Vec<String>) -> Self {
        self.common_strings = strings;
        self
    }

    /// Adds a common capability
    pub fn add_capability(mut self, capability: Capability) -> Self {
        self.common_capabilities.push(capability);
        self
    }

    /// Sets all common capabilities
    pub fn common_capabilities(mut self, capabilities: Vec<Capability>) -> Self {
        self.common_capabilities = capabilities;
        self
    }

    /// Adds a common code reference
    pub fn add_common_code_ref(mut self, ref_id: impl Into<String>) -> Self {
        self.common_code_refs.push(ref_id.into());
        self
    }

    /// Adds a common behavior reference
    pub fn add_common_behavior_ref(mut self, ref_id: impl Into<String>) -> Self {
        self.common_behavior_refs.push(ref_id.into());
        self
    }

    /// Adds an external reference
    pub fn add_reference(mut self, reference: ExternalReference) -> Self {
        self.references.push(reference);
        self
    }

    /// Sets all external references
    pub fn references(mut self, references: Vec<ExternalReference>) -> Self {
        self.references = references;
        self
    }

    /// Builds the MalwareFamily
    pub fn build(self) -> Result<MalwareFamily> {
        let name = self.name.ok_or(MaecError::MissingField("name"))?;

        let mut common = crate::common::CommonProperties::new("malware-family", None);
        if let Some(id) = self.id {
            common.id = id;
        }

        let family = MalwareFamily {
            common,
            name,
            aliases: self.aliases,
            labels: self.labels,
            description: self.description,
            field_data: self.field_data,
            common_strings: self.common_strings,
            common_capabilities: self.common_capabilities,
            common_code_refs: self.common_code_refs,
            common_behavior_refs: self.common_behavior_refs,
            references: self.references,
        };

        family.validate()?;
        Ok(family)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_malware_family_new() {
        let family = MalwareFamily::new("WannaCry");
        assert_eq!(family.name.value, "WannaCry");
        assert_eq!(family.common.r#type, "malware-family");
        assert!(family.common.id.starts_with("malware-family--"));
    }

    #[test]
    fn test_malware_family_builder() {
        let family = MalwareFamily::builder()
            .name(Name::new("Emotet"))
            .description("Banking trojan and malware loader")
            .add_label("trojan")
            .add_label("banking")
            .add_alias(Name::new("Geodo"))
            .build()
            .unwrap();

        assert_eq!(family.name.value, "Emotet");
        assert_eq!(family.labels.len(), 2);
        assert_eq!(family.aliases.len(), 1);
        assert!(family.description.is_some());
    }

    #[test]
    fn test_malware_family_validation() {
        let family = MalwareFamily::new("Test");
        assert!(family.validate().is_ok());
    }

    #[test]
    fn test_malware_family_serialize() {
        let family = MalwareFamily::builder()
            .name(Name::new("TestMalware"))
            .build()
            .unwrap();

        let json = serde_json::to_string(&family).unwrap();
        assert!(json.contains("\"type\":\"malware-family\""));
        assert!(json.contains("TestMalware"));

        let deserialized: MalwareFamily = serde_json::from_str(&json).unwrap();
        assert_eq!(family, deserialized);
    }
}