nabla-decompiler 0.1.2

//! Binary patching with safe state management and rollback

use anyhow::{anyhow, Result};
use nabla_scanner::binary::analysis::BinaryAnalysis;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::fs;
use std::path::{Path, PathBuf};
use std::sync::{Arc, RwLock};
use tempfile::TempDir;
use uuid::Uuid;
use chrono::{DateTime, Utc};

use crate::types::Address;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PatchSnapshot {
    pub original_hash: String,
    pub patches: Vec<AppliedPatch>,
    pub created_at: chrono::DateTime<chrono::Utc>,
    pub description: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AppliedPatch {
    pub address: Address,
    pub original_bytes: Vec<u8>,
    pub patched_bytes: Vec<u8>,
    pub description: String,
    pub pseudocode: String,
}

#[derive(Debug, Clone)]
pub struct PatchRequest {
    pub address: Address,
    pub hex_data: String,
    pub description: String,
    pub dry_run: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionPatch {
    pub id: Uuid,
    pub address: Address,
    pub original_bytes: Vec<u8>,
    pub patched_bytes: Vec<u8>,
    pub description: String,
    pub applied_at: DateTime<Utc>,
}

pub struct BinaryPatcher {
    binary_path: PathBuf,
    temp_dir: TempDir,
    snapshots: Vec<PatchSnapshot>,
    current_hash: String,
}

impl BinaryPatcher {
    pub fn new<P: AsRef<Path>>(binary_path: P) -> Result<Self> {
        let binary_path = binary_path.as_ref().to_path_buf();
        let temp_dir = TempDir::new()?;
        
        // Calculate initial hash
        let binary_data = fs::read(&binary_path)?;
        let current_hash = format!("{:x}", Sha256::digest(&binary_data));
        
        Ok(Self {
            binary_path,
            temp_dir,
            snapshots: Vec::new(),
            current_hash,
        })
    }
    
    /// Apply hex patch to binary
    pub fn apply_patch(&mut self, request: PatchRequest, analysis: &BinaryAnalysis) -> Result<String> {
        // Step 1: Validate the patch request
        self.validate_patch_request(&request, analysis)?;
        
        // Step 2: Parse hex data directly into bytes
        let machine_code = self.parse_hex_data(&request.hex_data)?;
        
        if request.dry_run {
            return Ok(format!(
                "DRY RUN - Would patch {} bytes at 0x{:x}:\nHex data: {}\nBytes: {:02x?}",
                machine_code.len(),
                request.address,
                request.hex_data,
                machine_code
            ));
        }
        
        // Step 3: Create backup snapshot
        let _snapshot = self.create_snapshot(&request.description)?;
        
        // Step 4: Apply the patch
        let original_bytes = self.patch_binary_at_address(request.address, &machine_code)?;
        
        // Step 5: Update snapshot with patch details
        let applied_patch = AppliedPatch {
            address: request.address,
            original_bytes,
            patched_bytes: machine_code.clone(),
            description: request.description.clone(),
            pseudocode: format!("hex: {}", request.hex_data),
        };
        
        if let Some(last_snapshot) = self.snapshots.last_mut() {
            last_snapshot.patches.push(applied_patch);
        }
        
        // Step 6: Update current hash
        let binary_data = fs::read(&self.binary_path)?;
        self.current_hash = format!("{:x}", Sha256::digest(&binary_data));
        
        Ok(format!(
            "✅ Successfully patched {} bytes at 0x{:x}\n📝 Description: {}\n🔧 Hex data: {}\n🔍 Bytes applied: {:02x?}",
            machine_code.len(),
            request.address,
            request.description,
            request.hex_data,
            machine_code
        ))
    }
    
    /// Rollback to a previous snapshot
    pub fn rollback(&mut self, snapshot_index: Option<usize>) -> Result<String> {
        let target_index = snapshot_index.unwrap_or(0);
        
        if target_index >= self.snapshots.len() {
            return Err(anyhow!("Invalid snapshot index: {}", target_index));
        }
        
        // Clone the data we need from the target snapshot before mutating
        let original_hash = self.snapshots[target_index].original_hash.clone();
        let created_at = self.snapshots[target_index].created_at;
        let description = self.snapshots[target_index].description.clone();
        
        // Restore from backup
        let backup_path = self.temp_dir.path().join(format!("backup_{}.bin", target_index));
        if backup_path.exists() {
            fs::copy(&backup_path, &self.binary_path)?;
            self.current_hash = original_hash;
            
            // Remove snapshots after the target
            self.snapshots.truncate(target_index + 1);
            
            Ok(format!(
                "✅ Rolled back to snapshot {}\n📅 Created: {}\n📝 Description: {}",
                target_index,
                created_at.format("%Y-%m-%d %H:%M:%S UTC"),
                description
            ))
        } else {
            Err(anyhow!("Backup file not found for snapshot {}", target_index))
        }
    }
    
    /// List all snapshots
    pub fn list_snapshots(&self) -> String {
        if self.snapshots.is_empty() {
            return "No snapshots available".to_string();
        }
        
        let mut result = String::from("📸 Binary Patch Snapshots:\n\n");
        
        for (i, snapshot) in self.snapshots.iter().enumerate() {
            result.push_str(&format!(
                "#{}: {} ({})\n  📅 {}\n  🔧 {} patches\n  🔍 Hash: {}...\n\n",
                i,
                snapshot.description,
                if i == self.snapshots.len() - 1 { "current" } else { "historical" },
                snapshot.created_at.format("%Y-%m-%d %H:%M:%S UTC"),
                snapshot.patches.len(),
                &snapshot.original_hash[..16]
            ));
        }
        
        result
    }
    
    /// Validate patch request
    fn validate_patch_request(&self, request: &PatchRequest, analysis: &BinaryAnalysis) -> Result<()> {
        // Check if address is in a code section first
        let target_section = analysis.code_sections.iter().find(|section| {
            request.address >= section.start_address && request.address < section.end_address
        });
        
        let Some(_section) = target_section else {
            return Err(anyhow!(
                "Address 0x{:x} is not in a known code section. This could corrupt data.",
                request.address
            ));
        };
        
        // Convert virtual address to file offset
        let binary_data = fs::read(&self.binary_path)?;
        let file_offset = self.virtual_address_to_file_offset(request.address, &binary_data)?;
        
        // Check if file offset is within binary bounds
        if file_offset >= binary_data.len() {
            return Err(anyhow!(
                "Address 0x{:x} maps to file offset 0x{:x}, which is beyond binary bounds (size: 0x{:x})",
                request.address,
                file_offset,
                binary_data.len()
            ));
        }
        
        // Validate hex data format
        if request.hex_data.trim().is_empty() {
            return Err(anyhow!("Hex data cannot be empty"));
        }
        
        Ok(())
    }
    
    /// Parse hex data directly into bytes with error handling
    fn parse_hex_data(&self, hex_data: &str) -> Result<Vec<u8>> {
        let cleaned = hex_data.trim().replace("0x", "").replace(" ", "");
        
        // Validate hex format
        if cleaned.is_empty() {
            return Err(anyhow!("Hex data cannot be empty"));
        }
        
        if !cleaned.chars().all(|c| c.is_ascii_hexdigit()) {
            return Err(anyhow!("Invalid hex data '{}': contains non-hex characters", hex_data));
        }
        
        if cleaned.len() % 2 != 0 {
            return Err(anyhow!("Invalid hex data '{}': must have even number of characters", hex_data));
        }
        
        // Convert hex string to bytes
        let mut bytes = Vec::new();
        for chunk in cleaned.as_bytes().chunks(2) {
            let hex_str = std::str::from_utf8(chunk)
                .map_err(|e| anyhow!("Invalid UTF-8 in hex data: {}", e))?;
            let byte = u8::from_str_radix(hex_str, 16)
                .map_err(|e| anyhow!("Invalid hex byte '{}': {}", hex_str, e))?;
            bytes.push(byte);
        }
        
        if bytes.is_empty() {
            return Err(anyhow!("Parsed hex data resulted in empty byte array"));
        }
        
        println!("Successfully parsed {} hex bytes: {:02x?}", bytes.len(), bytes);
        Ok(bytes)
    }
    
    
    
    /// Create a backup snapshot
    fn create_snapshot(&mut self, description: &str) -> Result<PatchSnapshot> {
        let binary_data = fs::read(&self.binary_path)?;
        let hash = format!("{:x}", Sha256::digest(&binary_data));
        
        // Save backup to temp directory
        let backup_path = self.temp_dir.path().join(format!("backup_{}.bin", self.snapshots.len()));
        fs::write(&backup_path, &binary_data)?;
        
        let snapshot = PatchSnapshot {
            original_hash: hash,
            patches: Vec::new(),
            created_at: chrono::Utc::now(),
            description: description.to_string(),
        };
        
        self.snapshots.push(snapshot.clone());
        Ok(snapshot)
    }
    
    /// Apply patch to binary at specific address
    fn patch_binary_at_address(&self, address: Address, new_bytes: &[u8]) -> Result<Vec<u8>> {
        let mut binary_data = fs::read(&self.binary_path)?;
        
        // Convert virtual address to file offset
        let file_offset = self.virtual_address_to_file_offset(address, &binary_data)?;
        let start_addr = file_offset;
        let end_addr = start_addr + new_bytes.len();
        
        if end_addr > binary_data.len() {
            return Err(anyhow!(
                "Patch would extend beyond binary bounds (file offset: 0x{:x}, patch size: {}, binary size: 0x{:x})",
                start_addr,
                new_bytes.len(),
                binary_data.len()
            ));
        }
        
        // Backup original bytes
        let original_bytes = binary_data[start_addr..end_addr].to_vec();
        
        // Apply patch
        binary_data.splice(start_addr..end_addr, new_bytes.iter().cloned());
        
        // Write patched binary
        fs::write(&self.binary_path, &binary_data)?;
        
        Ok(original_bytes)
    }
    
    /// Read bytes from binary at specified address
    pub fn read_bytes(&self, address: Address, length: usize) -> Result<Vec<u8>> {
        let binary_data = fs::read(&self.binary_path)?;
        let file_offset = self.virtual_address_to_file_offset(address, &binary_data)?;
        
        if file_offset + length > binary_data.len() {
            return Err(anyhow!(
                "Read would extend beyond binary bounds (file offset: 0x{:x}, read size: {}, binary size: 0x{:x})",
                file_offset,
                length,
                binary_data.len()
            ));
        }
        
        Ok(binary_data[file_offset..file_offset + length].to_vec())
    }
    
    /// Write bytes to binary at specified address
    pub fn write_bytes(&mut self, address: Address, bytes: &[u8]) -> Result<()> {
        let mut binary_data = fs::read(&self.binary_path)?;
        let file_offset = self.virtual_address_to_file_offset(address, &binary_data)?;
        let end_offset = file_offset + bytes.len();
        
        if end_offset > binary_data.len() {
            return Err(anyhow!(
                "Write would extend beyond binary bounds (file offset: 0x{:x}, write size: {}, binary size: 0x{:x})",
                file_offset,
                bytes.len(),
                binary_data.len()
            ));
        }
        
        // Replace bytes in binary data
        binary_data.splice(file_offset..end_offset, bytes.iter().cloned());
        
        // Write updated binary back to file
        fs::write(&self.binary_path, &binary_data)?;
        
        // Update current hash
        self.current_hash = format!("{:x}", Sha256::digest(&binary_data));
        
        Ok(())
    }
    
    /// Convert virtual address to file offset using heuristics
    fn virtual_address_to_file_offset(&self, virtual_address: u64, binary_data: &[u8]) -> Result<usize> {
        // This is a heuristic approach similar to the one used in disasm.rs
        // For more accurate mapping, we'd need to parse the binary format (ELF, PE, etc.)
        
        let potential_offset = if virtual_address > 0x400000 {
            // Typical Linux x86_64 binary base address
            (virtual_address - 0x400000) as usize
        } else if virtual_address > 0x8000000 {
            // Typical ARM binary base address  
            (virtual_address - 0x8000000) as usize
        } else if virtual_address > 0x10000000 {
            // Windows PE base address
            (virtual_address - 0x10000000) as usize
        } else if virtual_address > 0x1000 {
            // Small offset from base
            (virtual_address - 0x1000) as usize
        } else {
            // Assume it's already a file offset
            virtual_address as usize
        };
        
        // Ensure the calculated offset is within bounds
        if potential_offset < binary_data.len() {
            Ok(potential_offset)
        } else {
            // Fallback: try the address as-is (maybe it's already a file offset)
            let direct_offset = virtual_address as usize;
            if direct_offset < binary_data.len() {
                Ok(direct_offset)
            } else {
                Err(anyhow!(
                    "Cannot map virtual address 0x{:x} to valid file offset (tried 0x{:x}, binary size: 0x{:x})",
                    virtual_address,
                    potential_offset,
                    binary_data.len()
                ))
            }
        }
    }
}

pub struct MemoryPatcher {
    original_binary: Vec<u8>,
    working_binary: Arc<RwLock<Vec<u8>>>,
    applied_patches: Vec<SessionPatch>,
}

impl MemoryPatcher {
    pub fn new(binary_data: Vec<u8>) -> Self {
        Self {
            original_binary: binary_data.clone(),
            working_binary: Arc::new(RwLock::new(binary_data)),
            applied_patches: Vec::new(),
        }
    }

    pub fn apply_patch(&mut self, request: PatchRequest) -> Result<String> {
        let address = request.address;
        
        // Parse hex data into bytes
        let machine_code = self.parse_hex_data(&request.hex_data)?;
        
        if request.dry_run {
            return Ok(format!(
                "DRY RUN - Would patch {} bytes at 0x{:x}:\nHex data: {}\nBytes: {:02x?}",
                machine_code.len(),
                address,
                request.hex_data,
                machine_code
            ));
        }

        // Apply patch to in-memory binary
        let mut binary = self.working_binary.write().unwrap();
        
        // Convert virtual address to file offset
        let file_offset = self.virtual_address_to_file_offset(address, &binary)?;
        let start_idx = file_offset;
        let end_idx = start_idx + machine_code.len();
        
        if end_idx > binary.len() {
            return Err(anyhow!("Patch would extend beyond binary bounds"));
        }

        // Read original bytes before patching
        let original_bytes = binary[start_idx..end_idx].to_vec();

        // Apply the patch
        binary[start_idx..end_idx].copy_from_slice(&machine_code);

        // Record the patch
        let patch = SessionPatch {
            id: Uuid::new_v4(),
            address,
            original_bytes,
            patched_bytes: machine_code.clone(),
            description: request.description.clone(),
            applied_at: Utc::now(),
        };
        
        self.applied_patches.push(patch);

        Ok(format!(
            "✅ Successfully patched {} bytes at 0x{:x}\n📝 Description: {}\n🔧 Hex data: {}\n🔍 Bytes applied: {:02x?}",
            machine_code.len(),
            address,
            request.description,
            request.hex_data,
            machine_code
        ))
    }

    pub fn read_bytes(&self, address: Address, length: usize) -> Result<Vec<u8>> {
        let binary = self.working_binary.read().unwrap();
        let file_offset = self.virtual_address_to_file_offset(address, &binary)?;
        let start_idx = file_offset;
        let end_idx = start_idx + length;

        if end_idx > binary.len() {
            return Err(anyhow!("Read would extend beyond binary bounds"));
        }

        Ok(binary[start_idx..end_idx].to_vec())
    }

    pub fn get_working_binary(&self) -> Vec<u8> {
        self.working_binary.read().unwrap().clone()
    }

    pub fn get_original_binary(&self) -> &[u8] {
        &self.original_binary
    }

    pub fn get_applied_patches(&self) -> &[SessionPatch] {
        &self.applied_patches
    }

    pub fn get_applied_patches_mut(&mut self) -> &mut Vec<SessionPatch> {
        &mut self.applied_patches
    }

    pub fn get_working_binary_arc(&self) -> Arc<RwLock<Vec<u8>>> {
        Arc::clone(&self.working_binary)
    }

    pub fn get_original_binary_ref(&self) -> &[u8] {
        &self.original_binary
    }

    pub fn virtual_address_to_file_offset_public(&self, virtual_address: u64, binary_data: &[u8]) -> Result<usize> {
        self.virtual_address_to_file_offset(virtual_address, binary_data)
    }

    pub fn rollback_patch(&mut self, patch_id: Uuid) -> Result<String> {
        // Find the patch to rollback
        let patch_idx = self.applied_patches
            .iter()
            .position(|p| p.id == patch_id)
            .ok_or_else(|| anyhow!("Patch not found: {}", patch_id))?;

        // Rollback all patches from this point forward in reverse order
        let patches_to_rollback = self.applied_patches.split_off(patch_idx);
        
        for patch in patches_to_rollback.iter().rev() {
            let mut binary = self.working_binary.write().unwrap();
            let file_offset = self.virtual_address_to_file_offset(patch.address, &binary)?;
            let start_idx = file_offset;
            let end_idx = start_idx + patch.original_bytes.len();
            binary[start_idx..end_idx].copy_from_slice(&patch.original_bytes);
        }

        Ok(format!("✅ Rolled back patch and {} subsequent patches", patches_to_rollback.len() - 1))
    }

    pub fn rollback_all(&mut self) -> Result<String> {
        // Reset to original binary
        let mut binary = self.working_binary.write().unwrap();
        *binary = self.original_binary.clone();
        
        let patch_count = self.applied_patches.len();
        self.applied_patches.clear();

        Ok(format!("✅ Rolled back all {} patches", patch_count))
    }

    fn parse_hex_data(&self, hex_data: &str) -> Result<Vec<u8>> {
        let cleaned = hex_data.trim().replace("0x", "").replace(" ", "");
        
        if cleaned.is_empty() {
            return Err(anyhow!("Hex data cannot be empty"));
        }
        
        if !cleaned.chars().all(|c| c.is_ascii_hexdigit()) {
            return Err(anyhow!("Invalid hex data '{}': contains non-hex characters", hex_data));
        }
        
        if cleaned.len() % 2 != 0 {
            return Err(anyhow!("Invalid hex data '{}': must have even number of characters", hex_data));
        }
        
        let mut bytes = Vec::new();
        for chunk in cleaned.as_bytes().chunks(2) {
            let hex_str = std::str::from_utf8(chunk)?;
            let byte = u8::from_str_radix(hex_str, 16)
                .map_err(|e| anyhow!("Invalid hex byte '{}': {}", hex_str, e))?;
            bytes.push(byte);
        }
        
        Ok(bytes)
    }

    fn virtual_address_to_file_offset(&self, virtual_address: u64, binary_data: &[u8]) -> Result<usize> {
        let potential_offset = if virtual_address > 0x400000 {
            (virtual_address - 0x400000) as usize
        } else if virtual_address > 0x8000000 {
            (virtual_address - 0x8000000) as usize
        } else if virtual_address > 0x10000000 {
            (virtual_address - 0x10000000) as usize
        } else if virtual_address > 0x1000 {
            (virtual_address - 0x1000) as usize
        } else {
            virtual_address as usize
        };
        
        if potential_offset < binary_data.len() {
            Ok(potential_offset)
        } else {
            let direct_offset = virtual_address as usize;
            if direct_offset < binary_data.len() {
                Ok(direct_offset)
            } else {
                Err(anyhow!(
                    "Cannot map virtual address 0x{:x} to valid file offset",
                    virtual_address
                ))
            }
        }
    }
}

/// Utility functions for parsing addresses from strings
pub fn parse_address(addr_str: &str) -> Result<Address> {
    let cleaned = addr_str.trim().to_lowercase();
    
    if cleaned.starts_with("0x") {
        u64::from_str_radix(&cleaned[2..], 16)
            .map_err(|e| anyhow!("Invalid hex address '{}': {}", addr_str, e))
    } else {
        cleaned.parse::<u64>()
            .map_err(|e| anyhow!("Invalid address '{}': {}", addr_str, e))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::NamedTempFile;
    
    #[test]
    fn test_parse_address() {
        assert_eq!(parse_address("0x1000").unwrap(), 0x1000);
        assert_eq!(parse_address("4096").unwrap(), 4096);
        assert_eq!(parse_address("0X2000").unwrap(), 0x2000);
        assert!(parse_address("invalid").is_err());
    }
    
    #[test]
    fn test_parse_hex_data() {
        let patcher = create_test_patcher().unwrap();
        
        let machine_code = patcher.parse_hex_data("48c7c000000000").unwrap();
        assert_eq!(machine_code, vec![0x48, 0xc7, 0xc0, 0x00, 0x00, 0x00, 0x00]);
        
        let machine_code = patcher.parse_hex_data("0x90").unwrap();
        assert_eq!(machine_code, vec![0x90]);
        
        let machine_code = patcher.parse_hex_data("48 c7 c0 00").unwrap();
        assert_eq!(machine_code, vec![0x48, 0xc7, 0xc0, 0x00]);
        
        assert!(patcher.parse_hex_data("invalid").is_err());
        assert!(patcher.parse_hex_data("4").is_err()); // odd length
    }
    
    fn create_test_patcher() -> Result<BinaryPatcher> {
        let temp_file = NamedTempFile::new()?;
        let test_data = vec![0x48, 0xc7, 0xc0, 0x00, 0x00, 0x00, 0x00, 0xc3]; // mov rax, 0; ret
        fs::write(temp_file.path(), &test_data)?;
        BinaryPatcher::new(temp_file.path())
    }
}