nabla-decompiler 0.1.2

//! Binary disassembly using Capstone

use anyhow::Result;
use capstone::prelude::*;
use nabla_scanner::binary::analysis::{BinaryAnalysis, CodeSection, demangle_function_name};
use std::collections::HashMap;
use goblin;

use crate::types::{Address, Disassembly, Instruction, InstructionGroup, Section};

/// Disassemble entire binary
pub fn disassemble_binary(binary: &BinaryAnalysis) -> Result<Disassembly> {
    let cs = create_capstone_engine(&binary.architecture)?;

    let mut instructions = Vec::new();
    let mut symbols = HashMap::new();

    // Add known symbols from binary analysis with actual symbol addresses
    if let Some(binary_data) = &binary.binary_data {
        match goblin::Object::parse(binary_data) {
            Ok(goblin::Object::Elf(elf)) => {
                // ELF symbol extraction
                tracing::debug!("Extracting ELF symbols: {} static, {} dynamic", 
                    elf.syms.len(), elf.dynsyms.len());
                
                // Add symbols from ELF symbol table with their actual addresses
                for sym in elf.syms.iter() {
                    if let Some(name) = elf.strtab.get_at(sym.st_name) {
                        if !name.is_empty() && sym.st_value != 0 {
                            let demangled_name = demangle_function_name(name);
                            symbols.insert(sym.st_value, demangled_name);
                        }
                    }
                }
                
                // Add dynamic symbols
                for sym in elf.dynsyms.iter() {
                    if let Some(name) = elf.dynstrtab.get_at(sym.st_name) {
                        if !name.is_empty() && sym.st_value != 0 {
                            let demangled_name = demangle_function_name(name);
                            symbols.insert(sym.st_value, demangled_name);
                        }
                    }
                }
            }
            Ok(goblin::Object::PE(pe)) => {
                // PE (Windows executable) symbol extraction
                tracing::debug!("Extracting PE symbols from exports and imports");
                
                // Add exported symbols
                for export in pe.exports {
                    if let Some(name) = export.name {
                        let rva = export.rva;
                        // Convert RVA to address (PE images typically load at 0x400000)
                        let address = 0x400000u64 + rva as u64;
                        let demangled_name = demangle_function_name(&name);
                        symbols.insert(address, demangled_name);
                    }
                }
                
                // Add imported function names (though these won't have meaningful addresses in this context)
                for import in pe.imports {
                    let name = import.name;
                    // Use a placeholder address space for imports
                    let import_addr = 0x10000000u64 + (symbols.len() as u64 * 8);
                    let demangled_name = demangle_function_name(&name);
                    symbols.insert(import_addr, format!("import_{}", demangled_name));
                }
            }
            Ok(goblin::Object::Mach(mach)) => {
                // Mach-O (macOS/iOS executable) symbol extraction
                match mach {
                    goblin::mach::Mach::Binary(macho) => {
                        tracing::debug!("Extracting Mach-O symbols");
                        
                        // Extract symbols from symbol table
                        let symbols_iter = macho.symbols();
                        for symbol_result in symbols_iter {
                            if let Ok((name, nlist)) = symbol_result {
                                if !name.is_empty() && nlist.n_value != 0 {
                                    let demangled_name = demangle_function_name(name);
                                    symbols.insert(nlist.n_value, demangled_name);
                                }
                            }
                        }
                        
                        // Add exported symbols from dylib
                        for segment in &macho.segments {
                            if let Ok(sections) = segment.sections() {
                                for (section, _) in sections {
                                    if let Ok(section_name) = section.name() {
                                        if section_name.contains("__text") {
                                            // Add section as a symbol
                                            let addr = section.addr;
                                            symbols.insert(addr, format!("section_{}", section_name));
                                        }
                                    }
                                }
                            }
                        }
                    }
                    goblin::mach::Mach::Fat(_) => {
                        tracing::debug!("Fat Mach-O binary detected, skipping symbol extraction for now");
                    }
                }
            }
            Ok(goblin::Object::Archive(_)) => {
                tracing::debug!("Archive format detected, skipping symbol extraction");
            }
            Ok(goblin::Object::Unknown(_)) => {
                tracing::warn!("Unknown binary format, skipping symbol extraction");
            }
            Ok(_) => {
                tracing::debug!("Unsupported binary format, skipping symbol extraction");
            }
            Err(e) => {
                tracing::warn!("Failed to parse binary for symbol extraction: {}", e);
            }
        }
        
        tracing::debug!("Extracted {} symbols from binary", symbols.len());
    } else {
        tracing::warn!("No binary data available for symbol extraction");
    }

    // Disassemble code sections
    for section in &binary.code_sections {
        let section_data = get_section_data(binary, section)?;
        let insns = cs.disasm_all(&section_data, section.start_address)?;
        
        for insn in insns.iter() {
            instructions.push(Instruction {
                address: insn.address(),
                bytes: insn.bytes().to_vec(),
                mnemonic: insn.mnemonic().unwrap_or("").to_string(),
                operands: insn.op_str().unwrap_or("").to_string(),
                size: insn.len(),
                group: classify_instruction(&insn, &cs)?,
            });
        }
    }

    // Create sections from binary analysis
    let sections = binary.code_sections.iter().map(|s| Section {
        name: s.name.clone(),
        address: s.start_address,
        size: s.size,
        permissions: s.permissions.clone(),
    }).collect();

    instructions.sort_by_key(|i| i.address);

    Ok(Disassembly {
        instructions,
        sections,
        symbols,
    })
}

/// Disassemble from a specific address
pub fn disassemble_from_address(binary: &BinaryAnalysis, address: Address) -> Result<Disassembly> {
    let cs = create_capstone_engine(&binary.architecture)?;

    let mut instructions = Vec::new();
    let symbols = HashMap::new();

    // Find the section containing this address
    if let Some(section) = binary.code_sections.iter().find(|s| {
        address >= s.start_address && address < s.end_address
    }) {
        let section_data = get_section_data(binary, section)?;
        let offset = (address - section.start_address) as usize;
        
        if offset < section_data.len() {
            // Disassemble from the offset, limiting to reasonable size (1KB)
            let data = &section_data[offset..std::cmp::min(offset + 1024, section_data.len())];
            let insns = cs.disasm_all(data, address)?;
            
            for insn in insns.iter() {
                instructions.push(Instruction {
                    address: insn.address(),
                    bytes: insn.bytes().to_vec(),
                    mnemonic: insn.mnemonic().unwrap_or("").to_string(),
                    operands: insn.op_str().unwrap_or("").to_string(),
                    size: insn.len(),
                    group: classify_instruction(&insn, &cs)?,
                });
            }
        }
    }

    let sections = vec![Section {
        name: format!("function_at_0x{:x}", address),
        address,
        size: instructions.len() as u64 * 4, // Approximate
        permissions: "rx".to_string(),
    }];

    Ok(Disassembly {
        instructions,
        sections,
        symbols,
    })
}

/// Create Capstone engine based on architecture string with comprehensive support
fn create_capstone_engine(arch_str: &str) -> Result<Capstone> {
    let arch_lower = arch_str.to_lowercase();
    
    let cs = match arch_lower.as_str() {
        // ARM variants (most common in IoT firmware)
        "arm_thumb" | "armthumb" => {
            Capstone::new()
                .arm()
                .mode(arch::arm::ArchMode::Thumb)
                .detail(true)
                .build()?
        }
        "arm32" | "arm" => {
            Capstone::new()
                .arm()
                .mode(arch::arm::ArchMode::Arm)
                .detail(true)
                .build()?
        }
        "arm64" | "aarch64" => {
            Capstone::new()
                .arm64()
                .mode(arch::arm64::ArchMode::Arm)
                .detail(true)
                .build()?
        }
        "arm_cortex_m" | "cortex-m" => {
            // Cortex-M is ARM Thumb
            Capstone::new()
                .arm()
                .mode(arch::arm::ArchMode::Thumb)
                .detail(true)
                .build()?
        }
        
        // x86 variants
        "x86_64" | "x64" | "amd64" => {
            Capstone::new()
                .x86()
                .mode(arch::x86::ArchMode::Mode64)
                .detail(true)
                .build()?
        }
        "x86" | "i386" | "i686" => {
            Capstone::new()
                .x86()
                .mode(arch::x86::ArchMode::Mode32)
                .detail(true)
                .build()?
        }
        
        // MIPS variants
        "mips" | "mips32" => {
            Capstone::new()
                .mips()
                .mode(arch::mips::ArchMode::Mips32)
                .detail(true)
                .build()?
        }
        "mips64" => {
            Capstone::new()
                .mips()
                .mode(arch::mips::ArchMode::Mips64)
                .detail(true)
                .build()?
        }
        
        // PowerPC variants  
        "powerpc" | "ppc" | "ppc32" => {
            Capstone::new()
                .ppc()
                .mode(arch::ppc::ArchMode::Mode32)
                .detail(true)
                .build()?
        }
        "ppc64" | "powerpc64" => {
            Capstone::new()
                .ppc()
                .mode(arch::ppc::ArchMode::Mode64)
                .detail(true)
                .build()?
        }
        
        // Other embedded architectures
        "riscv" | "riscv32" => {
            // RISC-V support in Capstone (if available)
            // For now, default to ARM Thumb as fallback
            Capstone::new()
                .arm()
                .mode(arch::arm::ArchMode::Thumb)  
                .detail(true)
                .build()?
        } 
        
        "embedded" | "unknown" | _ => {
            tracing::warn!("Unknown architecture '{}', defaulting to ARM Thumb (most common in IoT)", arch_str);
            Capstone::new()
                .arm()
                .mode(arch::arm::ArchMode::Thumb)
                .detail(true)
                .build()?
        }
    };
    
    tracing::debug!("Created Capstone engine for architecture: {}", arch_str);
    Ok(cs)
}

/// Classify instruction for control flow analysis  
fn classify_instruction(insn: &capstone::Insn, cs: &Capstone) -> Result<InstructionGroup> {
    let mnemonic = insn.mnemonic().unwrap_or("").to_lowercase();
    
    // Check instruction groups using Capstone's detail info
    if let Ok(detail) = cs.insn_detail(insn) {
        for &group in detail.groups() {
            if let Some(group_name) = cs.group_name(group) {
                match group_name.to_lowercase().as_str() {
                    "jump" => return Ok(InstructionGroup::Jump),
                    "call" => return Ok(InstructionGroup::Call),
                    "ret" => return Ok(InstructionGroup::Return),
                    _ => {}
                }
            }
        }
    }
    
    // Fallback to mnemonic-based classification
    if mnemonic.starts_with("b") && (mnemonic.contains("l") || mnemonic == "bl") {
        Ok(InstructionGroup::Call)
    } else if mnemonic.starts_with("b") && !mnemonic.contains("l") {
        Ok(InstructionGroup::Jump)
    } else if mnemonic.contains("ret") || mnemonic == "bx" {
        Ok(InstructionGroup::Return)
    } else if mnemonic.starts_with("mov") || mnemonic.starts_with("ldr") || mnemonic.starts_with("str") {
        Ok(InstructionGroup::Move)
    } else if mnemonic.starts_with("add") || mnemonic.starts_with("sub") || mnemonic.starts_with("mul") {
        Ok(InstructionGroup::Arithmetic)
    } else if mnemonic.starts_with("and") || mnemonic.starts_with("orr") || mnemonic.starts_with("eor") {
        Ok(InstructionGroup::Logical)
    } else if mnemonic.starts_with("cmp") || mnemonic.starts_with("tst") {
        Ok(InstructionGroup::Compare)
    } else if mnemonic == "nop" {
        Ok(InstructionGroup::Nop)
    } else {
        Ok(InstructionGroup::Other)
    }
}

/// Get section data from binary
fn get_section_data(binary: &BinaryAnalysis, section: &CodeSection) -> Result<Vec<u8>> {
    // Get the binary data from the BinaryAnalysis
    let binary_data = binary.binary_data.as_ref()
        .ok_or_else(|| anyhow::anyhow!("Binary data not available in BinaryAnalysis"))?;
    
    let section_size = std::cmp::min(section.size, 64 * 1024) as usize; // Limit to 64KB for safety
    
    // Strategy 0: Use file offset if available (most reliable)
    if let Some(file_offset) = section.file_offset {
        let offset = file_offset as usize;
        if offset < binary_data.len() {
            let end_offset = std::cmp::min(offset + section_size, binary_data.len());
            if end_offset > offset {
                let section_data = &binary_data[offset..end_offset];
                tracing::debug!("Extracted section '{}' data using file offset 0x{:x} (size: {} bytes)", 
                    section.name, offset, section_data.len());
                return Ok(section_data.to_vec());
            }
        }
    }
    
    // Try to find section data using different strategies
    
    // Strategy 1: For small binaries, try to find patterns near the section address
    if binary_data.len() < 1024 * 1024 { // Less than 1MB
        // For small files, sections might be laid out sequentially
        // Try to find executable code patterns in the binary
        for window_start in (0..binary_data.len().saturating_sub(section_size)).step_by(16) {
            let window_end = std::cmp::min(window_start + section_size, binary_data.len());
            let candidate_data = &binary_data[window_start..window_end];
            
            // Check if this looks like executable code
            if looks_like_executable_code(candidate_data, &binary.architecture) {
                tracing::debug!("Found potential code section at file offset 0x{:x} for section '{}' (VA: 0x{:x})", 
                    window_start, section.name, section.start_address);
                return Ok(candidate_data.to_vec());
            }
        }
    }
    
    // Strategy 2: Try file offset calculation based on section address
    // This is a heuristic - for ELF files, we'd ideally parse the program headers
    // to get the proper virtual address to file offset mapping
    let potential_offset = if section.start_address > 0x400000 {
        // Typical Linux x86_64 binary base address
        (section.start_address - 0x400000) as usize
    } else if section.start_address > 0x8000000 {
        // Typical ARM binary base address
        (section.start_address - 0x8000000) as usize
    } else {
        section.start_address as usize
    };
    
    if potential_offset < binary_data.len() {
        let end_offset = std::cmp::min(potential_offset + section_size, binary_data.len());
        if end_offset > potential_offset {
            let section_data = &binary_data[potential_offset..end_offset];
            if !section_data.iter().all(|&b| b == 0) { // Not all zeros
                tracing::debug!("Extracted section '{}' data from file offset 0x{:x} (size: {} bytes)", 
                    section.name, potential_offset, section_data.len());
                return Ok(section_data.to_vec());
            }
        }
    }
    
    // Strategy 3: Fallback - search for non-zero data that might be code
    for chunk_start in (0..binary_data.len().saturating_sub(256)).step_by(64) {
        let chunk_end = std::cmp::min(chunk_start + std::cmp::min(section_size, 256), binary_data.len());
        let chunk = &binary_data[chunk_start..chunk_end];
        
        // Look for patterns that suggest executable code
        let non_zero_bytes = chunk.iter().filter(|&&b| b != 0).count();
        let entropy = calculate_simple_entropy(chunk);
        
        if non_zero_bytes > chunk.len() / 4 && entropy > 3.0 { // Reasonable entropy and non-zero content
            tracing::debug!("Found potential code at offset 0x{:x} for section '{}' (entropy: {:.2})", 
                chunk_start, section.name, entropy);
            return Ok(chunk.to_vec());
        }
    }
    
    // Final fallback: return empty data
    tracing::warn!("Could not locate section data for '{}' at 0x{:x}, returning empty data", 
        section.name, section.start_address);
    Ok(Vec::new())
}

/// Simple heuristic to check if data looks like executable code
fn looks_like_executable_code(data: &[u8], architecture: &str) -> bool {
    if data.len() < 4 {
        return false;
    }
    
    match architecture.to_lowercase().as_str() {
        "x86_64" | "i386" | "x86" => {
            // Look for common x86 instruction patterns
            data.windows(2).any(|w| {
                matches!(w, 
                    [0x48, _] | // REX prefix
                    [0x55, _] | // push rbp/ebp
                    [0x89, _] | // mov instructions
                    [0x83, _] | // immediate arithmetic
                    [0xff, _] | // various instructions
                    [0x8b, _] | // mov from memory
                    [0xc3, _] | // ret
                    [0xe8, _] | // call
                    [0x74, _] | [0x75, _] // conditional jumps
                )
            })
        },
        "arm" | "arm_thumb" | "armthumb" => {
            // Look for common ARM/Thumb instruction patterns
            data.windows(2).any(|w| {
                matches!(w,
                    [0x70, 0x47] | // bx lr (Thumb)
                    [0x00, 0x20] | // movs r0, #0 (Thumb)
                    [0x08, 0x44] | // add r0, r1 (Thumb)
                    [0x00, 0xbf] | // nop (Thumb)
                    [_, 0xe0..=0xef] // ARM branch instructions
                )
            })
        },
        _ => {
            // Generic heuristic: look for varied bytes (not all zero/same)
            let unique_bytes: std::collections::HashSet<u8> = data.iter().copied().collect();
            unique_bytes.len() > 4 && !data.iter().all(|&b| b == 0 || b == 0xff)
        }
    }
}

/// Calculate simple entropy for data
fn calculate_simple_entropy(data: &[u8]) -> f64 {
    if data.is_empty() {
        return 0.0;
    }
    
    let mut counts = [0u32; 256];
    for &byte in data {
        counts[byte as usize] += 1;
    }
    
    let len = data.len() as f64;
    let mut entropy = 0.0;
    
    for &count in &counts {
        if count > 0 {
            let probability = count as f64 / len;
            entropy -= probability * probability.log2();
        }
    }
    
    entropy
}