oletools_rs 0.1.0

//! OleID checker — runs all security checks on a document.
//!
//! Port of oletools/oleid.py.

use crate::ftguess::detector::FileTypeGuesser;
use crate::mraptor::analyzer::{MacroRaptor, MRaptorResult};
use crate::ole::container::OleFile;
use crate::oleid::indicator::{Indicator, RiskLevel};
use crate::oleobj::extractor::OleObjExtractor;
use crate::vba::parser::VbaParser;

/// Flash magic bytes.
const FLASH_MAGIC: &[&[u8]] = &[b"CWS", b"FWS", b"ZWS"];

/// OleID — comprehensive security indicator analysis.
pub struct OleID {
    data: Vec<u8>,
}

impl OleID {
    /// Create a new OleID analyzer from raw file data.
    pub fn new(data: &[u8]) -> Self {
        Self {
            data: data.to_vec(),
        }
    }

    /// Run all checks and return a list of indicators.
    pub fn analyze(&self) -> Vec<Indicator> {
        vec![
            self.check_format(),
            self.check_encrypted(),
            self.check_vba_macros(),
            self.check_macro_analysis(),
            self.check_external_relationships(),
            self.check_object_pool(),
            self.check_flash(),
        ]
    }

    /// Check document format.
    fn check_format(&self) -> Indicator {
        match FileTypeGuesser::from_bytes(&self.data) {
            Ok(result) => Indicator::new(
                "format",
                "File format",
                "Detected file format",
                format!("{:?}", result.file_type),
                RiskLevel::Info,
            ),
            Err(e) => Indicator::new(
                "format",
                "File format",
                "Detected file format",
                format!("Error: {}", e),
                RiskLevel::Error,
            ),
        }
    }

    /// Check if the document is encrypted.
    fn check_encrypted(&self) -> Indicator {
        // Check for EncryptedPackage stream (OLE) or encryption namespace (OOXML)
        let is_encrypted = if OleFile::is_ole(&self.data) {
            match OleFile::from_bytes(&self.data) {
                Ok(ole) => {
                    ole.exists("EncryptedPackage")
                        || ole.exists("/EncryptedPackage")
                        || ole.exists("EncryptionInfo")
                        || ole.exists("/EncryptionInfo")
                }
                Err(_) => false,
            }
        } else {
            false
        };

        let (value, risk) = if is_encrypted {
            ("True".to_string(), RiskLevel::Low)
        } else {
            ("False".to_string(), RiskLevel::None)
        };

        Indicator::new(
            "encrypted",
            "Encrypted",
            "Document is encrypted",
            value,
            risk,
        )
    }

    /// Check for VBA macros.
    fn check_vba_macros(&self) -> Indicator {
        match VbaParser::from_bytes(&self.data) {
            Ok(parser) => match parser.detect_vba_macros() {
                Ok(true) => Indicator::new(
                    "vba_macros",
                    "VBA Macros",
                    "Contains VBA macros",
                    "True",
                    RiskLevel::Medium,
                ),
                Ok(false) => Indicator::new(
                    "vba_macros",
                    "VBA Macros",
                    "Contains VBA macros",
                    "False",
                    RiskLevel::None,
                ),
                Err(_) => Indicator::new(
                    "vba_macros",
                    "VBA Macros",
                    "Contains VBA macros",
                    "Unknown",
                    RiskLevel::Unknown,
                ),
            },
            Err(_) => Indicator::new(
                "vba_macros",
                "VBA Macros",
                "Contains VBA macros",
                "Unknown",
                RiskLevel::Unknown,
            ),
        }
    }

    /// Run macro analysis (MacroRaptor A/W/X heuristic).
    fn check_macro_analysis(&self) -> Indicator {
        match MacroRaptor::scan_file(&self.data) {
            Ok((result, flags)) => {
                let (value, risk) = match result {
                    MRaptorResult::Suspicious => {
                        let mut parts = Vec::new();
                        if flags.autoexec {
                            parts.push("AutoExec");
                        }
                        if flags.write {
                            parts.push("Write");
                        }
                        if flags.execute {
                            parts.push("Execute");
                        }
                        (
                            format!("Suspicious ({})", parts.join(", ")),
                            RiskLevel::High,
                        )
                    }
                    MRaptorResult::Clean => ("Clean".to_string(), RiskLevel::None),
                    MRaptorResult::NoMacro => ("No macros".to_string(), RiskLevel::None),
                };

                Indicator::new(
                    "macro_analysis",
                    "Macro Analysis",
                    "MacroRaptor heuristic result (A+W/X)",
                    value,
                    risk,
                )
            }
            Err(_) => Indicator::new(
                "macro_analysis",
                "Macro Analysis",
                "MacroRaptor heuristic result (A+W/X)",
                "Error",
                RiskLevel::Error,
            ),
        }
    }

    /// Check for external relationships (OOXML).
    fn check_external_relationships(&self) -> Indicator {
        match OleObjExtractor::from_bytes(&self.data) {
            Ok(extractor) => match extractor.find_external_relationships() {
                Ok(rels) if rels.is_empty() => Indicator::new(
                    "ext_rels",
                    "External Relationships",
                    "External targets in OOXML relationships",
                    "False",
                    RiskLevel::None,
                ),
                Ok(rels) => {
                    let targets: Vec<_> = rels.iter().map(|r| r.target.as_str()).collect();
                    Indicator::new(
                        "ext_rels",
                        "External Relationships",
                        "External targets in OOXML relationships",
                        format!("{} found: {}", rels.len(), targets.join(", ")),
                        RiskLevel::High,
                    )
                }
                Err(_) => Indicator::new(
                    "ext_rels",
                    "External Relationships",
                    "External targets in OOXML relationships",
                    "N/A",
                    RiskLevel::None,
                ),
            },
            Err(_) => Indicator::new(
                "ext_rels",
                "External Relationships",
                "External targets in OOXML relationships",
                "N/A",
                RiskLevel::None,
            ),
        }
    }

    /// Check for ObjectPool storage (OLE embedded objects).
    fn check_object_pool(&self) -> Indicator {
        if !OleFile::is_ole(&self.data) {
            return Indicator::new(
                "object_pool",
                "ObjectPool",
                "Contains ObjectPool storage (embedded OLE objects)",
                "N/A",
                RiskLevel::None,
            );
        }

        match OleFile::from_bytes(&self.data) {
            Ok(ole) => {
                let has_pool = ole.exists("ObjectPool") || ole.exists("/ObjectPool");
                let (value, risk) = if has_pool {
                    ("True", RiskLevel::Low)
                } else {
                    ("False", RiskLevel::None)
                };
                Indicator::new(
                    "object_pool",
                    "ObjectPool",
                    "Contains ObjectPool storage (embedded OLE objects)",
                    value,
                    risk,
                )
            }
            Err(_) => Indicator::new(
                "object_pool",
                "ObjectPool",
                "Contains ObjectPool storage (embedded OLE objects)",
                "Error",
                RiskLevel::Error,
            ),
        }
    }

    /// Check for Flash content embedded in OLE streams.
    fn check_flash(&self) -> Indicator {
        if !OleFile::is_ole(&self.data) {
            return Indicator::new(
                "flash",
                "Flash Objects",
                "Contains Flash (SWF) objects",
                "N/A",
                RiskLevel::None,
            );
        }

        match OleFile::from_bytes(&self.data) {
            Ok(mut ole) => {
                let streams = ole.list_streams();
                let mut found_flash = false;

                for stream_path in &streams {
                    if let Ok(data) = ole.open_stream(stream_path) {
                        if data.len() >= 3 {
                            for magic in FLASH_MAGIC {
                                if data.starts_with(magic) {
                                    found_flash = true;
                                    break;
                                }
                            }
                        }
                        if found_flash {
                            break;
                        }
                    }
                }

                let (value, risk) = if found_flash {
                    ("True", RiskLevel::High)
                } else {
                    ("False", RiskLevel::None)
                };

                Indicator::new(
                    "flash",
                    "Flash Objects",
                    "Contains Flash (SWF) objects",
                    value,
                    risk,
                )
            }
            Err(_) => Indicator::new(
                "flash",
                "Flash Objects",
                "Contains Flash (SWF) objects",
                "Error",
                RiskLevel::Error,
            ),
        }
    }
}

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

    #[test]
    fn test_analyze_empty_data() {
        let oleid = OleID::new(&[]);
        let indicators = oleid.analyze();
        assert_eq!(indicators.len(), 7);
    }

    #[test]
    fn test_analyze_random_data() {
        let oleid = OleID::new(&[0x00, 0x01, 0x02, 0x03, 0x04]);
        let indicators = oleid.analyze();
        assert_eq!(indicators.len(), 7);

        // Format should be detected
        let format = indicators.iter().find(|i| i.id == "format").unwrap();
        assert_eq!(format.risk, RiskLevel::Info);
    }

    #[test]
    fn test_check_format_unknown() {
        let oleid = OleID::new(&[0xFF, 0xFE]);
        let indicator = oleid.check_format();
        assert_eq!(indicator.id, "format");
    }

    #[test]
    fn test_check_encrypted_non_ole() {
        let oleid = OleID::new(&[0x00, 0x01]);
        let indicator = oleid.check_encrypted();
        assert_eq!(indicator.value, "False");
        assert_eq!(indicator.risk, RiskLevel::None);
    }

    #[test]
    fn test_check_vba_unknown_format() {
        let oleid = OleID::new(&[0x00, 0x01, 0x02]);
        let indicator = oleid.check_vba_macros();
        assert_eq!(indicator.value, "Unknown");
    }

    #[test]
    fn test_check_object_pool_non_ole() {
        let oleid = OleID::new(&[0x50, 0x4B, 0x03, 0x04]);
        let indicator = oleid.check_object_pool();
        assert_eq!(indicator.value, "N/A");
    }

    #[test]
    fn test_check_flash_non_ole() {
        let oleid = OleID::new(&[0x00, 0x01]);
        let indicator = oleid.check_flash();
        assert_eq!(indicator.value, "N/A");
    }

    #[test]
    fn test_indicator_ids_unique() {
        let oleid = OleID::new(&[]);
        let indicators = oleid.analyze();
        let mut ids: Vec<_> = indicators.iter().map(|i| i.id.as_str()).collect();
        ids.sort();
        ids.dedup();
        assert_eq!(ids.len(), indicators.len(), "All indicator IDs should be unique");
    }
}