malwareanalyzer.rs

//! Rust Malware Analyzer (RMA)
//! Author: Ilya Markin
//! 
//! Purpose: Identify library code vs. attacker-authored code in Rust binaries
//! Based on techniques from Microsoft RIFT, JPCERT/CC research, and Threatray methodologies
//!
//! Core features:
//! - Extract compiler version and crate information
//! - Demangle Rust symbols
//! - Identify known library functions
//! - Highlight suspicious attacker code
//! - Generate YARA rules from malware-specific functions

use std::collections::{HashMap, HashSet};
use std::fs;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::str;

use anyhow::{Result, Context, bail};
use clap::{Parser, Subcommand};
use goblin::elf::Elf;
use goblin::pe::PE;
use goblin::mach::Mach;
use serde::{Serialize, Deserialize};
use regex::Regex;
use walkdir::WalkDir;

// ============================================================================
// Data Structures
// ============================================================================

#[derive(Parser)]
#[clap(author, version, about = "Rust Malware Analyzer - Identify library vs. malware code")]
struct Cli {
    #[clap(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Analyze a single binary
    Analyze {
        /// Path to binary
        binary: PathBuf,
        
        /// Output format (text, json, markdown)
        #[clap(short, long, default_value = "text")]
        format: String,
        
        /// Extract all functions for further analysis
        #[clap(short, long)]
        extract: bool,
    },
    
    /// Batch analyze multiple binaries
    Batch {
        /// Directory containing binaries
        directory: PathBuf,
        
        /// Output directory for reports
        #[clap(short, long)]
        output: PathBuf,
        
        /// Generate YARA rules from common functions
        #[clap(long)]
        generate_yara: bool,
    },
    
    /// Generate FLIRT-like signatures from known Rust libraries
    GenerateSignatures {
        /// Output signature file
        output: PathBuf,
        
        /// Path to Rust installation (optional, auto-detect if not specified)
        #[clap(long)]
        rust_path: Option<PathBuf>,
    },
    
    /// Compare two binaries to find shared malicious code
    Diff {
        /// First binary
        binary1: PathBuf,
        
        /// Second binary
        binary2: PathBuf,
        
        /// Output format
        #[clap(short, long, default_value = "text")]
        format: String,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct AnalysisResult {
    /// File information
    file_info: FileInfo,
    
    /// Rust compiler information
    compiler_info: CompilerInfo,
    
    /// All functions found in binary
    functions: Vec<Function>,
    
    /// Library functions identified (likely benign)
    library_functions: Vec<Function>,
    
    /// Potential malware functions (attacker-authored)
    malware_functions: Vec<Function>,
    
    /// Suspicious indicators
    indicators: Vec<String>,
    
    /// Function clusters by similarity
    clusters: Vec<FunctionCluster>,
    
    /// YARA rules generated (if requested)
    yara_rules: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct FileInfo {
    path: String,
    size: u64,
    file_type: String,
    arch: String,
    sha256: String,
    entropy: f64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct CompilerInfo {
    rustc_version: Option<String>,
    commit_hash: Option<String>,
    target_triple: Option<String>,
    crates: Vec<CrateInfo>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct CrateInfo {
    name: String,
    version: String,
    functions_found: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct Function {
    name: String,
    demangled_name: Option<String>,
    address: u64,
    size: usize,
    bytes: Vec<u8>,
    hex_bytes: String,
    is_library: bool,
    library_name: Option<String>,
    crate_name: Option<String>,
    confidence: f32,
    features: FunctionFeatures,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct FunctionFeatures {
    entropy: f64,
    instructions: Vec<String>,
    calls: Vec<String>,
    strings: Vec<String>,
    api_calls: Vec<String>,
    basic_blocks: usize,
    cyclomatic_complexity: u32,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct FunctionCluster {
    id: usize,
    functions: Vec<usize>, // indices in functions array
    common_features: Vec<String>,
    malware_family: Option<String>,
}

// ============================================================================
// Known Rust Library Database
// ============================================================================

/// Database of known Rust library functions (simplified version of what RIFT does)
/// Based on analysis of common crates: tokio, serde, regex, reqwest, etc.
struct RustLibraryDB {
    /// Map of function name patterns to crate names
    patterns: Vec<(Regex, String)>,
    
    /// Map of known function hashes to crate names
    hashes: HashMap<u64, String>,
    
    /// Map of crate names to versions
    crate_versions: HashMap<String, Vec<String>>,
}

impl RustLibraryDB {
    fn new() -> Self {
        let mut db = RustLibraryDB {
            patterns: Vec::new(),
            hashes: HashMap::new(),
            crate_versions: HashMap::new(),
        };
        
        // Initialize with common patterns
        db.add_patterns();
        db.add_common_hashes();
        
        db
    }
    
    fn add_patterns(&mut self) {
        // Core library patterns
        self.patterns.push((
            Regex::new(r"^core::.*").unwrap(),
            "core".to_string()
        ));
        self.patterns.push((
            Regex::new(r"^std::.*").unwrap(),
            "std".to_string()
        ));
        self.patterns.push((
            Regex::new(r"^alloc::.*").unwrap(),
            "alloc".to_string()
        ));
        
        // Common crates
        let common_crates = vec![
            "tokio", "serde", "regex", "reqwest", "hyper", "rustls",
            "log", "env_logger", "clap", "anyhow", "thiserror", "once_cell",
            "lazy_static", "crossbeam", "rayon", "futures", "async_std",
            "smol", "tokio_rustls", "webpki", "ring", "rand", "chrono",
            "time", "uuid", "base64", "hex", "sha2", "sha1", "md5",
            "aes", "cipher", "block_modes", "crypto_common", "digest",
        ];
        
        for crate_name in common_crates {
            self.patterns.push((
                Regex::new(&format!(r"^{}::.*", crate_name)).unwrap(),
                crate_name.to_string()
            ));
        }
    }
    
    fn add_common_hashes(&mut self) {
        // In a real implementation, this would be populated from
        // pre-computed hashes of known library functions
        // For now, we'll leave it as a placeholder
    }
    
    fn identify_library_function(&self, name: &str, hash: u64) -> Option<String> {
        // Check by hash first (more accurate)
        if let Some(crate_name) = self.hashes.get(&hash) {
            return Some(crate_name.clone());
        }
        
        // Fall back to pattern matching
        for (pattern, crate_name) in &self.patterns {
            if pattern.is_match(name) {
                return Some(crate_name.clone());
            }
        }
        
        None
    }
}

// ============================================================================
// Parser Engine
// ============================================================================

struct ParserEngine {
    db: RustLibraryDB,
    yara_gen: YaraGenerator,
}

impl ParserEngine {
    fn new() -> Self {
        ParserEngine {
            db: RustLibraryDB::new(),
            yara_gen: YaraGenerator::new(),
        }
    }
    
    fn analyze_binary(&self, path: &Path) -> Result<AnalysisResult> {
        println!("[*] Analyzing: {}", path.display());
        
        let data = fs::read(path)
            .with_context(|| format!("Failed to read {}", path.display()))?;
        
        // Calculate file hash
        use sha2::{Sha256, Digest};
        let mut hasher = Sha256::new();
        hasher.update(&data);
        let hash = format!("{:x}", hasher.finalize());
        
        // Calculate entropy
        let entropy = calculate_entropy(&data);
        
        // Determine file type and parse
        let (file_type, arch, functions) = self.parse_binary(&data)?;
        
        // Extract Rust compiler info
        let compiler_info = self.extract_compiler_info(&data, &functions)?;
        
        // Identify library functions
        let mut library_fns = Vec::new();
        let mut malware_fns = Vec::new();
        let mut indicators = Vec::new();
        
        for func in &functions {
            if func.is_library {
                library_fns.push(func.clone());
            } else {
                malware_fns.push(func.clone());
                
                // Check for suspicious indicators
                self.check_suspicious_indicators(func, &mut indicators);
            }
        }
        
        // Cluster similar functions (to find variants)
        let clusters = self.cluster_functions(&functions);
        
        // Generate YARA rules if enough malware functions found
        let yara_rules = if malware_fns.len() > 5 {
            Some(self.yara_gen.generate_rules(&malware_fns))
        } else {
            None
        };
        
        let file_info = FileInfo {
            path: path.display().to_string(),
            size: data.len() as u64,
            file_type,
            arch,
            sha256: hash,
            entropy,
        };
        
        Ok(AnalysisResult {
            file_info,
            compiler_info,
            functions,
            library_functions: library_fns,
            malware_functions: malware_fns,
            indicators,
            clusters,
            yara_rules,
        })
    }
    
    fn parse_binary(&self, data: &[u8]) -> Result<(String, String, Vec<Function>)> {
        // Try PE first (Windows)
        if let Ok(pe) = PE::parse(data) {
            return self.parse_pe(pe, data);
        }
        
        // Try ELF (Linux)
        if let Ok(elf) = Elf::parse(data) {
            return self.parse_elf(elf, data);
        }
        
        // Try Mach-O (macOS)
        if let Ok(mach) = Mach::parse(data) {
            return self.parse_mach(mach, data);
        }
        
        bail!("Unsupported file format");
    }
    
    fn parse_pe(&self, pe: PE, data: &[u8]) -> Result<(String, String, Vec<Function>)> {
        let arch = match pe.header.coff_header.machine {
            0x8664 => "x86_64".to_string(),
            0x14c => "x86".to_string(),
            0xaa64 => "ARM64".to_string(),
            _ => "unknown".to_string(),
        };
        
        let mut functions = Vec::new();
        
        // Extract functions from .text section
        for section in &pe.sections {
            let name = section.name().unwrap_or("").to_string();
            if name == ".text" || name == "CODE" {
                let section_data = &data[section.pointer_to_raw_data as usize..
                    (section.pointer_to_raw_data + section.size_of_raw_data) as usize];
                
                // Find function starts (simplified - in real tool would use disassembly)
                let mut addr = section.virtual_address;
                let mut offset = 0;
                
                while offset < section_data.len() {
                    // Simple heuristic: look for function prologues
                    if offset + 4 < section_data.len() {
                        let bytes = &section_data[offset..offset+4];
                        
                        // x86_64 function prologues
                        if bytes == [0x55, 0x48, 0x89, 0xe5] || // push rbp; mov rbp, rsp
                           bytes == [0x48, 0x83, 0xec, 0x28] || // sub rsp, 0x28
                           bytes == [0x53, 0x48, 0x83, 0xec] {  // push rbx; sub rsp, ...
                            
                            // Found potential function
                            let name = format!("sub_{:x}", addr);
                            let demangled = self.demangle_rust_name(&name);
                            
                            // Calculate function features
                            let features = self.analyze_function(section_data, offset);
                            
                            // Check if this is a library function
                            let (is_library, library_name) = self.db.identify_library_function(
                                &demangled.as_deref().unwrap_or(&name),
                                hash_bytes(&section_data[offset..offset+32])
                            ).map(|c| (true, Some(c))).unwrap_or((false, None));
                            
                            functions.push(Function {
                                name,
                                demangled_name: demangled,
                                address: addr,
                                size: 0, // Would need proper size detection
                                bytes: section_data[offset..offset+32].to_vec(),
                                hex_bytes: hex::encode(&section_data[offset..offset+16]),
                                is_library,
                                library_name,
                                crate_name: None, // Would need more analysis
                                confidence: if is_library { 0.9 } else { 0.3 },
                                features,
                            });
                            
                            addr += 16; // Skip ahead - simplified
                            offset += 16;
                        } else {
                            addr += 1;
                            offset += 1;
                        }
                    } else {
                        break;
                    }
                }
            }
        }
        
        Ok(("PE".to_string(), arch, functions))
    }
    
    fn parse_elf(&self, elf: Elf, data: &[u8]) -> Result<(String, String, Vec<Function>)> {
        let arch = match elf.header.e_machine {
            0x3e => "x86_64".to_string(), // EM_X86_64
            0x03 => "x86".to_string(),    // EM_386
            0xb7 => "ARM64".to_string(),  // EM_AARCH64
            _ => "unknown".to_string(),
        };
        
        // Similar to PE parsing but for ELF sections
        // (simplified for brevity)
        
        Ok(("ELF".to_string(), arch, Vec::new()))
    }
    
    fn parse_mach(&self, mach: Mach, data: &[u8]) -> Result<(String, String, Vec<Function>)> {
        // Mach-O parsing (simplified)
        Ok(("Mach-O".to_string(), "x86_64".to_string(), Vec::new()))
    }
    
    fn demangle_rust_name(&self, name: &str) -> Option<String> {
        // Rust mangling patterns:
        // _RNvCskw... - v0 mangling
        // _ZN... - legacy mangling
        
        if name.starts_with("_RN") {
            // v0 mangling - use rustc-demangle
            if let Ok(output) = Command::new("rustfilt")
                .arg(name)
                .output() {
                if output.status.success() {
                    return String::from_utf8(output.stdout).ok();
                }
            }
        } else if name.starts_with("_ZN") {
            // legacy mangling
            let mut result = String::new();
            let mut chars = name.chars().skip(3).peekable();
            
            while let Some(len_char) = chars.next() {
                if let Some(len) = len_char.to_digit(10) {
                    let mut segment = String::new();
                    for _ in 0..len {
                        if let Some(c) = chars.next() {
                            segment.push(c);
                        }
                    }
                    if !result.is_empty() {
                        result.push_str("::");
                    }
                    result.push_str(&segment);
                }
            }
            
            if !result.is_empty() {
                return Some(result);
            }
        }
        
        None
    }
    
    fn analyze_function(&self, data: &[u8], offset: usize) -> FunctionFeatures {
        let mut features = FunctionFeatures {
            entropy: calculate_entropy(&data[offset..offset+64.min(data.len()-offset)]),
            instructions: Vec::new(),
            calls: Vec::new(),
            strings: Vec::new(),
            api_calls: Vec::new(),
            basic_blocks: 0,
            cyclomatic_complexity: 0,
        };
        
        // Simple instruction detection (would use capstone in real tool)
        let mut pos = offset;
        while pos < data.len() && pos < offset + 256 {
            if data[pos] == 0xe8 { // call instruction
                if pos + 4 < data.len() {
                    let target = i32::from_le_bytes([
                        data[pos+1], data[pos+2], data[pos+3], data[pos+4]
                    ]);
                    features.calls.push(format!("call 0x{:x}", target));
                }
                pos += 5;
            } else if data[pos] == 0x90 { // nop
                features.instructions.push("nop".to_string());
                pos += 1;
            } else {
                pos += 1;
            }
        }
        
        features
    }
    
    fn extract_compiler_info(&self, data: &[u8], functions: &[Function]) -> Result<CompilerInfo> {
        let mut compiler_info = CompilerInfo {
            rustc_version: None,
            commit_hash: None,
            target_triple: None,
            crates: Vec::new(),
        };
        
        // Search for rustc version strings
        if let Ok(contents) = std::str::from_utf8(data) {
            // Look for "rustc 1.xx.x" patterns
            let re = Regex::new(r"rustc ([\d\.]+) \(([a-f0-9]+) [\d-]+\)").unwrap();
            for cap in re.captures_iter(contents) {
                compiler_info.rustc_version = Some(cap[1].to_string());
                compiler_info.commit_hash = Some(cap[2].to_string());
                break;
            }
        }
        
        // Count functions per crate
        let mut crate_counts: HashMap<String, usize> = HashMap::new();
        for func in functions {
            if let Some(ref lib) = func.library_name {
                *crate_counts.entry(lib.clone()).or_insert(0) += 1;
            }
        }
        
        for (name, count) in crate_counts {
            compiler_info.crates.push(CrateInfo {
                name,
                version: "unknown".to_string(), // Would need more analysis
                functions_found: count,
            });
        }
        
        Ok(compiler_info)
    }
    
    fn check_suspicious_indicators(&self, func: &Function, indicators: &mut Vec<String>) {
        // Check for common malware patterns
        
        // Anti-debugging
        if func.features.instructions.iter().any(|i| i.contains("int3") || i.contains("rdtsc")) {
            indicators.push("Anti-debugging instructions detected".to_string());
        }
        
        // Process injection APIs
        let injection_apis = ["VirtualAllocEx", "WriteProcessMemory", "CreateRemoteThread"];
        for api in &func.features.api_calls {
            if injection_apis.iter().any(|&a| api.contains(a)) {
                indicators.push(format!("Process injection API: {}", api));
            }
        }
        
        // Suspicious strings
        let suspicious = ["cmd.exe", "powershell", "http://", "https://", "C2"];
        for s in &func.features.strings {
            if suspicious.iter().any(|&sus| s.contains(sus)) {
                indicators.push(format!("Suspicious string: {}", s));
            }
        }
        
        // High entropy (packed/encrypted code)
        if func.features.entropy > 7.5 {
            indicators.push(format!("High entropy section: {:.2}", func.features.entropy));
        }
    }
    
    fn cluster_functions(&self, functions: &[Function]) -> Vec<FunctionCluster> {
        let mut clusters = Vec::new();
        let mut used = vec![false; functions.len()];
        
        for i in 0..functions.len() {
            if used[i] { continue; }
            
            let mut cluster = FunctionCluster {
                id: clusters.len(),
                functions: vec![i],
                common_features: Vec::new(),
                malware_family: None,
            };
            
            // Find similar functions (simplified - would use proper similarity metrics)
            for j in i+1..functions.len() {
                if !used[j] && functions[i].bytes == functions[j].bytes {
                    cluster.functions.push(j);
                    used[j] = true;
                }
            }
            
            if cluster.functions.len() > 1 {
                clusters.push(cluster);
            }
            
            used[i] = true;
        }
        
        clusters
    }
}

// ============================================================================
// YARA Rule Generator
// ============================================================================

struct YaraGenerator {
    rule_counter: usize,
}

impl YaraGenerator {
    fn new() -> Self {
        YaraGenerator {
            rule_counter: 0,
        }
    }
    
    fn generate_rules(&mut self, malware_functions: &[Function]) -> String {
        let mut rules = String::new();
        
        rules.push_str("// YARA rules generated by Rust Malware Analyzer\n");
        rules.push_str("// Author: Ilya Markin\n\n");
        
        // Group similar functions
        let mut seen_bytes = HashSet::new();
        
        for func in malware_functions {
            if func.bytes.len() < 8 { continue; }
            
            // Generate rule for distinctive byte patterns
            if !seen_bytes.contains(&func.hex_bytes) && func.confidence > 0.7 {
                seen_bytes.insert(func.hex_bytes.clone());
                
                let rule_name = format!("RustMalware_{:04x}", self.rule_counter);
                self.rule_counter += 1;
                
                rules.push_str(&format!("rule {}\n", rule_name));
                rules.push_str("{\n");
                rules.push_str("    meta:\n");
                rules.push_str("        author = \"Ilya Markin\"\n");
                rules.push_str("        description = \"Potential Rust malware function\"\n");
                if let Some(ref name) = func.demangled_name {
                    rules.push_str(&format!("        function = \"{}\"\n", name));
                }
                rules.push_str(&format!("        confidence = \"{:.2}\"\n", func.confidence));
                rules.push_str("\n");
                rules.push_str("    strings:\n");
                rules.push_str(&format!("        $func = {{ {} }}\n", func.hex_bytes));
                rules.push_str("\n");
                rules.push_str("    condition:\n");
                rules.push_str("        $func\n");
                rules.push_str("}\n\n");
            }
        }
        
        rules
    }
}

// ============================================================================
// Utility Functions
// ============================================================================

fn calculate_entropy(data: &[u8]) -> f64 {
    let mut counts = [0u32; 256];
    for &b in data {
        counts[b as usize] += 1;
    }
    
    let len = data.len() as f64;
    let mut entropy = 0.0;
    
    for &count in &counts {
        if count > 0 {
            let p = count as f64 / len;
            entropy -= p * p.log2();
        }
    }
    
    entropy
}

fn hash_bytes(data: &[u8]) -> u64 {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::{Hash, Hasher};
    
    let mut hasher = DefaultHasher::new();
    data.hash(&mut hasher);
    hasher.finish()
}

// ============================================================================
// Output Formatting
// ============================================================================

struct OutputFormatter;

impl OutputFormatter {
    fn format(&self, result: &AnalysisResult, format: &str) -> String {
        match format {
            "json" => self.format_json(result),
            "markdown" => self.format_markdown(result),
            _ => self.format_text(result),
        }
    }
    
    fn format_text(&self, result: &AnalysisResult) -> String {
        let mut output = String::new();
        
        output.push_str(&"=".repeat(80));
        output.push_str("\n");
        output.push_str("RUST MALWARE ANALYZER REPORT\n");
        output.push_str(&"=".repeat(80));
        output.push_str("\n\n");
        
        // File info
        output.push_str("FILE INFORMATION\n");
        output.push_str(&"-".repeat(40));
        output.push_str("\n");
        output.push_str(&format!("Path: {}\n", result.file_info.path));
        output.push_str(&format!("Size: {} bytes\n", result.file_info.size));
        output.push_str(&format!("Type: {}\n", result.file_info.file_type));
        output.push_str(&format!("Arch: {}\n", result.file_info.arch));
        output.push_str(&format!("SHA256: {}\n", result.file_info.sha256));
        output.push_str(&format!("Entropy: {:.4}\n", result.file_info.entropy));
        output.push_str("\n");
        
        // Compiler info
        output.push_str("RUST COMPILER INFORMATION\n");
        output.push_str(&"-".repeat(40));
        output.push_str("\n");
        if let Some(ref ver) = result.compiler_info.rustc_version {
            output.push_str(&format!("rustc version: {}\n", ver));
        }
        if let Some(ref hash) = result.compiler_info.commit_hash {
            output.push_str(&format!("Commit hash: {}\n", hash));
        }
        if !result.compiler_info.crates.is_empty() {
            output.push_str("\nDetected crates:\n");
            for crate_info in &result.compiler_info.crates {
                output.push_str(&format!("  - {} ({} functions)\n", 
                    crate_info.name, crate_info.functions_found));
            }
        }
        output.push_str("\n");
        
        // Summary
        output.push_str("ANALYSIS SUMMARY\n");
        output.push_str(&"-".repeat(40));
        output.push_str("\n");
        output.push_str(&format!("Total functions: {}\n", result.functions.len()));
        output.push_str(&format!("Library functions: {}\n", result.library_functions.len()));
        output.push_str(&format!("Potential malware functions: {}\n", result.malware_functions.len()));
        output.push_str(&format!("Function clusters: {}\n", result.clusters.len()));
        output.push_str("\n");
        
        // Suspicious indicators
        if !result.indicators.is_empty() {
            output.push_str("⚠️  SUSPICIOUS INDICATORS\n");
            output.push_str(&"-".repeat(40));
            output.push_str("\n");
            for indicator in &result.indicators {
                output.push_str(&format!("  • {}\n", indicator));
            }
            output.push_str("\n");
        }
        
        // Top malware functions
        if !result.malware_functions.is_empty() {
            output.push_str("TOP SUSPICIOUS FUNCTIONS\n");
            output.push_str(&"-".repeat(40));
            output.push_str("\n");
            
            let mut sorted = result.malware_functions.clone();
            sorted.sort_by(|a, b| b.confidence.partial_cmp(&a.confidence).unwrap());
            
            for func in sorted.iter().take(10) {
                let name = func.demangled_name.as_ref().unwrap_or(&func.name);
                output.push_str(&format!("  • {} (conf: {:.2})\n", name, func.confidence));
                output.push_str(&format!("    Bytes: {}\n", func.hex_bytes));
                if !func.features.api_calls.is_empty() {
                    output.push_str(&format!("    APIs: {:?}\n", &func.features.api_calls[..3.min(func.features.api_calls.len())]));
                }
            }
        }
        
        output.push_str(&"=".repeat(80));
        output.push_str("\n");
        
        output
    }
    
    fn format_json(&self, result: &AnalysisResult) -> String {
        serde_json::to_string_pretty(result).unwrap()
    }
    
    fn format_markdown(&self, result: &AnalysisResult) -> String {
        let mut output = String::new();
        
        output.push_str("# Rust Malware Analysis Report\n\n");
        output.push_str(&format!("**File:** `{}`  \n", result.file_info.path));
        output.push_str(&format!("**SHA256:** `{}`  \n", result.file_info.sha256));
        output.push_str(&format!("**Analysis Date:** {}  \n\n", chrono::Local::now().format("%Y-%m-%d %H:%M:%S")));
        
        output.push_str("## Summary\n\n");
        output.push_str(&format!("- **File Type:** {}  \n", result.file_info.file_type));
        output.push_str(&format!("- **Architecture:** {}  \n", result.file_info.arch));
        output.push_str(&format!("- **Total Functions:** {}  \n", result.functions.len()));
        output.push_str(&format!("- **Library Functions:** {}  \n", result.library_functions.len()));
        output.push_str(&format!("- **Potential Malware Functions:** {}  \n", result.malware_functions.len()));
        output.push_str(&format!("- **Suspicious Indicators:** {}  \n\n", result.indicators.len()));
        
        if !result.indicators.is_empty() {
            output.push_str("## ⚠️ Suspicious Indicators\n\n");
            for indicator in &result.indicators {
                output.push_str(&format!("- {}\n", indicator));
            }
            output.push_str("\n");
        }
        
        if let Some(yara) = &result.yara_rules {
            output.push_str("## Generated YARA Rules\n\n");
            output.push_str("```yara\n");
            output.push_str(yara);
            output.push_str("```\n\n");
        }
        
        output
    }
}

// ============================================================================
// Main
// ============================================================================

fn main() -> Result<()> {
    let cli = Cli::parse();
    let analyzer = ParserEngine::new();
    let formatter = OutputFormatter;
    
    match cli.command {
        Commands::Analyze { binary, format, extract } => {
            println!("[*] Starting Rust malware analysis...");
            
            let result = analyzer.analyze_binary(&binary)?;
            
            let output = formatter.format(&result, &format);
            println!("{}", output);
            
            if extract {
                // Save functions for further analysis
                let func_dir = Path::new("extracted_functions");
                std::fs::create_dir_all(func_dir)?;
                
                for (i, func) in result.malware_functions.iter().enumerate() {
                    let name = func.demangled_name.as_ref().unwrap_or(&func.name);
                    let safe_name = name.replace(|c: char| !c.is_ascii_alphanumeric(), "_");
                    let path = func_dir.join(format!("{}_{:x}.bin", safe_name, func.address));
                    std::fs::write(&path, &func.bytes)?;
                }
                
                println!("[*] Extracted {} functions to {}", result.malware_functions.len(), func_dir.display());
            }
        }
        
        Commands::Batch { directory, output, generate_yara } => {
            println!("[*] Batch analyzing directory: {}", directory.display());
            
            std::fs::create_dir_all(&output)?;
            
            let mut all_results = Vec::new();
            
            for entry in WalkDir::new(&directory)
                .into_iter()
                .filter_map(|e| e.ok())
                .filter(|e| e.file_type().is_file()) {
                
                let path = entry.path();
                if let Ok(result) = analyzer.analyze_binary(path) {
                    all_results.push(result);
                }
            }
            
            // Generate summary
            let summary_path = output.join("batch_summary.json");
            let summary_json = serde_json::to_string_pretty(&all_results)?;
            std::fs::write(summary_path, summary_json)?;
            
            if generate_yara {
                // Collect all malware functions
                let mut all_malware = Vec::new();
                for result in &all_results {
                    all_malware.extend(result.malware_functions.clone());
                }
                
                let yara_rules = analyzer.yara_gen.generate_rules(&all_malware);
                let yara_path = output.join("generated_rules.yar");
                std::fs::write(yara_path, yara_rules)?;
            }
            
            println!("[*] Batch analysis complete. Results saved to {}", output.display());
        }
        
        Commands::GenerateSignatures { output, rust_path } => {
            println!("[*] Generating FLIRT-like signatures from Rust libraries...");
            println!("[*] This would download and compile Rust crates to extract function signatures");
            println!("[*] (Similar to Microsoft RIFT Generator)");
            
            // In a real implementation, this would:
            // 1. Find Rust installation
            // 2. Download common crates
            // 3. Compile them with debug symbols
            // 4. Extract function hashes and patterns
            // 5. Save to output file
            
            std::fs::write(&output, "# Rust library signatures\n")?;
            println!("[*] Signatures saved to {}", output.display());
        }
        
        Commands::Diff { binary1, binary2, format } => {
            println!("[*] Comparing two binaries...");
            
            let result1 = analyzer.analyze_binary(&binary1)?;
            let result2 = analyzer.analyze_binary(&binary2)?;
            
            // Find common malware functions
            let mut common = Vec::new();
            for f1 in &result1.malware_functions {
                for f2 in &result2.malware_functions {
                    if f1.bytes == f2.bytes {
                        common.push((f1, f2));
                        break;
                    }
                }
            }
            
            println!("[*] Found {} common malware function(s)", common.len());
            
            if format == "json" {
                let diff_result = serde_json::json!({
                    "binary1": result1.file_info,
                    "binary2": result2.file_info,
                    "common_functions": common.iter().map(|(f1, f2)| {
                        serde_json::json!({
                            "name": f1.demangled_name,
                            "bytes": f1.hex_bytes,
                            "address1": f1.address,
                            "address2": f2.address,
                        })
                    }).collect::<Vec<_>>()
                });
                println!("{}", serde_json::to_string_pretty(&diff_result)?);
            } else {
                println!("\nCommon malware code found!");
                for (i, (f1, f2)) in common.iter().enumerate() {
                    println!("  {}. {} (0x{:x} -> 0x{:x})", 
                        i+1, 
                        f1.demangled_name.as_ref().unwrap_or(&f1.name),
                        f1.address,
                        f2.address
                    );
                }
            }
        }
    }
    
    Ok(())
}