Advance this code massively, no CFG analysis. Focus all your effort in improving the disassembler module.
main.rs:
use clap::Parser;
use goblin::Object;
use log::{error, info};
use std::fs;
use std::path::PathBuf;

mod disassembler;

use disassembler::Disassembler;

/// Rust Binary Analyzer with Custom Disassembler
#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
/// Path to the binary file to analyze
#[arg(short, long, value_name = "FILE")]
input: PathBuf,


/// Optionally specify the architecture (currently only x86_64 is supported)
#[arg(short, long, default_value = "x86_64")]
arch: String,
}

fn main() {
// Initialize the logger
env_logger::init();


// Parse command-line arguments
let args = Args::parse();

// Read the binary file
let buffer = match fs::read(&args.input) {
    Ok(buf) => buf,
    Err(e) => {
        error!("Failed to read file {}: {}", args.input.display(), e);
        return;
    }
};

// Parse the binary using goblin
match Object::parse(&buffer) {
    Ok(Object::Elf(elf)) => {
        info!("Parsed ELF binary");
        analyze_elf(&elf, &buffer, &args.arch);
    }
    Ok(Object::PE(pe)) => {
        info!("Parsed PE binary");
        analyze_pe(&pe, &buffer, &args.arch);
    }
    Ok(_) => {
        error!("Unsupported binary format");
    }
    Err(e) => {
        error!("Failed to parse binary: {}", e);
    }
}
}

/// Analyze ELF binaries
fn analyze_elf(elf: &goblin::elf::Elf, buffer: &[u8], arch: &str) {
// Initialize disassembler
let disassembler = Disassembler::new(arch);


// Iterate over executable sections
for section in elf.section_headers.iter() {
    if section.is_executable() {
        let section_name = match elf.shdr_strtab.get_at(section.sh_name) {
            Some(name) => name.to_string(),
            None => "<unknown>".to_string(),
        };
        let section_addr = section.sh_addr;
        let section_size = section.sh_size as usize;
        let section_data = &buffer[section.sh_offset as usize..section.sh_offset as usize + section_size];

        info!("Disassembling section: {} (0x{:x} - 0x{:x})", section_name, section_addr, section_addr + section.sh_size);
        let instructions = disassembler.disassemble(section_data, section_addr);

        for instr in instructions {
            println!("{:#010x}: {}\t{:?}", instr.address, instr.mnemonic, instr.operands);
        }
    }
}
}

/// Analyze PE binaries
fn analyze_pe(pe: &goblin::pe::PE, buffer: &[u8], arch: &str) {
// Initialize disassembler
let disassembler = Disassembler::new(arch);


// Iterate over executable sections
for section in pe.sections.iter() {
    if section.characteristics & goblin::pe::section_table::IMAGE_SCN_MEM_EXECUTE != 0 {
        let section_name = String::from_utf8_lossy(&section.name).trim_end_matches('\0').to_string();
        let section_addr = section.virtual_address as u64;
        let section_size = section.virtual_size as usize;
        let section_data = &buffer[section.pointer_to_raw_data as usize..section.pointer_to_raw_data as usize + section_size];

        info!("Disassembling section: {} (0x{:x} - 0x{:x})", section_name, section_addr, section_addr + section.virtual_size as u64);
        let instructions = disassembler.disassemble(section_data, section_addr);

        for instr in instructions {
            println!("{:#010x}: {}\t{:?}", instr.address, instr.mnemonic, instr.operands);
        }
    }
}
}

disassembler/mod.rs:
pub mod decoder;
pub mod instruction;

use decoder::Decoder;
use instruction::{Instruction, Operand};

/// Disassembler engine
pub struct Disassembler {
decoder: Decoder,
}

impl Disassembler {
/// Create a new Disassembler for the specified architecture
pub fn new(arch: &str) -> Self {
match arch {
"x86_64" => Disassembler {
decoder: Decoder::new_x86_64(),
},
_ => {
panic!("Unsupported architecture: {}", arch);
}
}
}


/// Disassemble a slice of bytes starting at a given address
pub fn disassemble(&self, code: &[u8], start_address: u64) -> Vec<Instruction> {
    let mut instructions = Vec::new();
    let mut offset = 0;
    let code_len = code.len();
    let mut current_address = start_address;

    while offset < code_len {
        match self.decoder.decode(&code[offset..], current_address) {
            Ok(instr) => {
                instructions.push(instr.clone());
                let instr_len = instr.bytes.len();
                if instr_len == 0 {
                    // Prevent infinite loop on zero-length instruction
                    break;
                }
                offset += instr_len;
                current_address += instr_len as u64;
            }
            Err(e) => {
                // Handle unknown opcode or decoding error as a single byte
                instructions.push(Instruction {
                    mnemonic: "db".to_string(),
                    operands: vec![Operand::Immediate(code[offset] as u64)],
                    address: current_address,
                    bytes: vec![code[offset]],
                });
                offset += 1;
                current_address += 1;
            }
        }
    }

    instructions
}
}
disassembler/instruction.rs:
#[derive(Debug, Clone)]
pub enum Operand {
Register(String),
Immediate(u64),
Memory {
base: String,
index: String,
scale: i32,
displacement: i64,
},
RelativeAddress(u64),
Label(String), // For symbolic displacement or jump targets
// Additional operand types can be added as needed
}

#[derive(Debug, Clone)]
pub struct Instruction {
pub mnemonic: String,
pub operands: Vec,
pub address: u64,
pub bytes: Vec,
}

disassembler/decoder.rs:
use crate::disassembler::instruction::{Instruction, Operand};
use std::collections::HashMap;

/// Decoder for x86_64 architecture
pub struct Decoder {
/// Opcode mapping: opcode bytes to [OpcodeInfo]
opcode_map: HashMap<Vec, Vec>,
}

#[derive(Debug, Clone)]
struct OpcodeInfo {
mnemonic: &'static str,
opcode_length: usize, // Number of bytes in the opcode
operands: Vec,
has_modrm: bool,
has_sib: bool,
rex_required: bool,
is_multibyte: bool, // Indicates if opcode uses prefixes like 0x0F
}

/// Specification for an operand in an instruction
#[derive(Debug, Clone)]
struct OperandSpec {
operand_type: OperandType,
size: OperandSize,
// Additional fields can specify operand position or encoding details
}

/// Sizes of operands
#[derive(Debug, Clone, PartialEq)]
enum OperandSize {
Byte,
Word,
Dword,
Qword,
// Other sizes can be added as needed
}

/// Types of operands for decoding purposes
#[derive(Debug, Clone)]
enum OperandType {
Register,
Immediate,
Memory,
RelativeAddress,
// Additional types like Segment, Pointer, etc., can be added
}

impl Decoder {
/// Initialize a new Decoder for x86_64
pub fn new_x86_64() -> Self {
let mut opcode_map: HashMap<Vec, Vec> = HashMap::new();


    // Example: Single-byte opcodes
    // NOP: 0x90
    opcode_map.insert(
        vec![0x90],
        vec![OpcodeInfo {
            mnemonic: "NOP",
            opcode_length: 1,
            operands: vec![],
            has_modrm: false,
            has_sib: false,
            rex_required: false,
            is_multibyte: false,
        }],
    );

    // RET: 0xC3
    opcode_map.insert(
        vec![0xC3],
        vec![OpcodeInfo {
            mnemonic: "RET",
            opcode_length: 1,
            operands: vec![],
            has_modrm: false,
            has_sib: false,
            rex_required: false,
            is_multibyte: false,
        }],
    );

    // MOV r64, imm32: 0xB8 + rd
    for rd in 0..8 {
        let opcode = vec![0xB8 + rd];
        opcode_map
            .entry(opcode.clone())
            .or_insert_with(Vec::new)
            .push(OpcodeInfo {
                mnemonic: "MOV",
                opcode_length: 1,
                operands: vec![
                    OperandSpec {
                        operand_type: OperandType::Register,
                        size: OperandSize::Qword,
                    },
                    OperandSpec {
                        operand_type: OperandType::Immediate,
                        size: OperandSize::Dword,
                    },
                ],
                has_modrm: false,
                has_sib: false,
                rex_required: rd >= 4, // REX prefix needed for extended registers
                is_multibyte: false,
            });
    }

    // CALL rel32: 0xE8
    opcode_map.insert(
        vec![0xE8],
        vec![OpcodeInfo {
            mnemonic: "CALL",
            opcode_length: 1,
            operands: vec![OperandSpec {
                operand_type: OperandType::RelativeAddress,
                size: OperandSize::Dword,
            }],
            has_modrm: false,
            has_sib: false,
            rex_required: false,
            is_multibyte: false,
        }],
    );

    // JMP rel8: 0xEB
    opcode_map.insert(
        vec![0xEB],
        vec![OpcodeInfo {
            mnemonic: "JMP",
            opcode_length: 1,
            operands: vec![OperandSpec {
                operand_type: OperandType::RelativeAddress,
                size: OperandSize::Byte,
            }],
            has_modrm: false,
            has_sib: false,
            rex_required: false,
            is_multibyte: false,
        }],
    );

    // ADD r/m64, r64: 0x01
    opcode_map.insert(
        vec![0x01],
        vec![OpcodeInfo {
            mnemonic: "ADD",
            opcode_length: 1,
            operands: vec![
                OperandSpec {
                    operand_type: OperandType::Memory, // r/m64
                    size: OperandSize::Qword,
                },
                OperandSpec {
                    operand_type: OperandType::Register, // r64
                    size: OperandSize::Qword,
                },
            ],
            has_modrm: true,
            has_sib: true, // Memory operand may have SIB
            rex_required: false,
            is_multibyte: false,
        }],
    );

    // More opcodes can be added here with similar structure...

    // Multi-byte opcodes (using 0x0F prefix)
    // Example: 0x0F 0xAF: IMUL r/m32, r32
    opcode_map.insert(
        vec![0x0F, 0xAF],
        vec![OpcodeInfo {
            mnemonic: "IMUL",
            opcode_length: 2,
            operands: vec![
                OperandSpec {
                    operand_type: OperandType::Memory, // r/m32
                    size: OperandSize::Dword,
                },
                OperandSpec {
                    operand_type: OperandType::Register, // r32
                    size: OperandSize::Dword,
                },
            ],
            has_modrm: true,
            has_sib: true,
            rex_required: false,
            is_multibyte: true,
        }],
    );

    // Initialize the Decoder with the opcode map
    Decoder { opcode_map }
}

/// Decode a single instruction from the given bytes starting at `address`
pub fn decode(&self, bytes: &[u8], address: u64) -> Result<Instruction, DecodeError> {
    if bytes.is_empty() {
        return Err(DecodeError::IncompleteInstruction);
    }

    let mut offset = 0;
    let mut rex_prefix = RexPrefix::default();
    let mut prefixes = Vec::new();

    // Handle possible prefixes (REX, operand-size, address-size, segment overrides, etc.)
    while offset < bytes.len() {
        match bytes[offset] {
            0x40..=0x4F => {
                // REX prefix
                rex_prefix = RexPrefix::from_byte(bytes[offset])?;
                prefixes.push(bytes[offset]);
                offset += 1;
            }
            0xF0 | 0xF2 | 0xF3 | 0x2E | 0x36 | 0x3E | 0x26 | 0x64 | 0x65 | 0x66 | 0x67 => {
                // Business other prefixes like LOCK, REP, segment overrides, etc.
                prefixes.push(bytes[offset]);
                offset += 1;
            }
            _ => break,
        }
    }

    // Attempt to match opcodes, starting with the longest possible
    let max_opcode_length = self
        .opcode_map
        .keys()
        .map(|k| k.len())
        .max()
        .unwrap_or(1);
    let mut matched_opcode_info: Option<(&OpcodeInfo, usize)> = None;

    // Search for the longest matching opcode
    for length in (1..=max_opcode_length).rev() {
        if offset + length > bytes.len() {
            continue;
        }
        let opcode_bytes = bytes[offset..offset + length].to_vec();
        if let Some(op_infos) = self.opcode_map.get(&opcode_bytes) {
            for op_info in op_infos {
                // Check REX prefix requirements
                if op_info.rex_required && !rex_prefix.is_present() {
                    continue;
                }
                matched_opcode_info = Some((op_info, length));
                break;
            }
            if matched_opcode_info.is_some() {
                break;
            }
        }
    }

    let (opcode_info, opcode_length) = match matched_opcode_info {
        Some(m) => m,
        None => return Err(DecodeError::UnknownOpcode),
    };

    offset += opcode_length;
    let mut operands = Vec::new();

    // Process operands
    for operand_spec in &opcode_info.operands {
        match operand_spec.operand_type {
            OperandType::Register => {
                let reg = if opcode_info.has_modrm {
                    // Register is typically encoded in the ModR/M reg field
                    let reg_field = self.get_modrm_reg_field(bytes, offset)?;
                    self.decode_register_field(reg_field, &rex_prefix)?
                } else {
                    // Otherwise, it might be encoded in the opcode
                    self.decode_register(&rex_prefix, opcode_info, &prefixes)?
                };
                operands.push(Operand::Register(reg));
            }
            OperandType::Immediate => {
                let imm = self.read_immediate(bytes, offset, &operand_spec.size)?;
                offset += operand_spec.size.bytes();
                operands.push(Operand::Immediate(imm));
            }
            OperandType::RelativeAddress => {
                let rel = self.read_immediate(bytes, offset, &operand_spec.size)? as i64;
                let instr_size = offset as usize; // Total bytes read so far
                let target = address
                    .wrapping_add(instr_size as u64)
                    .wrapping_add(rel as u64);
                operands.push(Operand::RelativeAddress(target));
                offset += operand_spec.size.bytes();
            }
            OperandType::Memory => {
                if !opcode_info.has_modrm {
                    return Err(DecodeError::MissingModRM);
                }
                let (mem_operand, new_offset) = self.decode_memory_operand(bytes, offset, &rex_prefix)?;
                operands.push(mem_operand);
                offset = new_offset;
            }
        }
    }

    // Extract the bytes of the instruction
    let instr_bytes = bytes[..offset].to_vec();

    Ok(Instruction {
        mnemonic: opcode_info.mnemonic.to_string(),
        operands,
        address,
        bytes: instr_bytes,
    })
}

/// Decode the register operand based on REX prefix and opcode context
fn decode_register(
    &self,
    rex: &RexPrefix,
    opcode_info: &OpcodeInfo,
    prefixes: &[u8],
) -> Result<String, DecodeError> {
    // Placeholder: In a complete implementation, decode based on opcode context
    // Here, we're returning 'rax' modified by REX.W for demonstration
    if rex.w {
        Ok("rax".to_string())
    } else {
        Ok("eax".to_string())
    }
}

/// Decode a register field from current instruction bytes
fn decode_register_field(
    &self,
    reg_field: u8,
    rex: &RexPrefix,
) -> Result<String, DecodeError> {
    let extended = if rex.has_r_extension() { 8 } else { 0 };
    let registers = [
        "rax", "rcx", "rdx", "rbx", "rsp", "rbp", "rsi", "rdi",
        "r8", "r9", "r10", "r11", "r12", "r13", "r14", "r15",
    ];
    if (reg_field as usize) + extended >= registers.len() {
        return Err(DecodeError::InvalidRegister);
    }
    Ok(registers[(reg_field as usize) + extended].to_string())
}

/// Decode the ModR/M byte's reg field
fn get_modrm_reg_field(&self, bytes: &[u8], offset: usize) -> Result<u8, DecodeError> {
    if offset >= bytes.len() {
        return Err(DecodeError::IncompleteInstruction);
    }
    let modrm = bytes[offset];
    Ok((modrm & 0b00111000) >> 3)
}

/// Decode the memory operand including ModR/M and SIB bytes
fn decode_memory_operand(
    &self,
    bytes: &[u8],
    mut offset: usize,
    rex: &RexPrefix,
) -> Result<(Operand, usize), DecodeError> {
    if offset >= bytes.len() {
        return Err(DecodeError::IncompleteInstruction);
    }
    let modrm = bytes[offset];
    offset += 1;

    let mod_field = (modrm & 0b11000000) >> 6;
    let reg_field = (modrm & 0b00111000) >> 3;
    let rm_field = modrm & 0b00000111;

    // Decode SIB if needed
    let (base, index, scale, displacement, sib_offset) = if rm_field == 4 && mod_field != 3 {
        // SIB byte is present
        if offset >= bytes.len() {
            return Err(DecodeError::IncompleteInstruction);
        }
        let sib = bytes[offset];
        offset += 1;
        self.decode_sib(sib, rex)?
    } else {
        ("".to_string(), "".to_string(), 1, 0, 0)
    };

    // Decode displacement
    let displacement = match mod_field {
        0 => {
            if rm_field == 5 {
                // disp32
                if offset + 3 >= bytes.len() {
                    return Err(DecodeError::IncompleteInstruction);
                }
                let disp = i32::from_le_bytes([
                    bytes[offset],
                    bytes[offset + 1],
                    bytes[offset + 2],
                    bytes[offset + 3],
                ]) as i64;
                offset += 4;
                disp
            } else {
                0
            }
        }
        1 => {
            // disp8
            if offset >= bytes.len() {
                return Err(DecodeError::IncompleteInstruction);
            }
            let disp = bytes[offset] as i8 as i64;
            offset += 1;
            disp
        }
        2 => {
            // disp32
            if offset + 3 >= bytes.len() {
                return Err(DecodeError::IncompleteInstruction);
            }
            let disp = i32::from_le_bytes([
                bytes[offset],
                bytes[offset + 1],
                bytes[offset + 2],
                bytes[offset + 3],
            ]) as i64;
            offset += 4;
            disp
        }
        _ => 0,
    };

    // Decode base register
    let base_reg = if mod_field == 0 && rm_field == 5 {
        // No base, [disp32]
        String::from("")
    } else {
        self.decode_register_field(rm_field, rex)?
    };

    let mem_operand = Operand::Memory {
        base,
        index,
        scale,
        displacement,
    };

    Ok((mem_operand, offset))
}

/// Decode the SIB byte
fn decode_sib(
    &self,
    sib: u8,
    rex: &RexPrefix,
) -> Result<(String, String, i32, i64, usize), DecodeError> {
    let scale = match (sib & 0b11000000) >> 6 {
        0 => 1,
        1 => 2,
        2 => 4,
        3 => 8,
        _ => 1, // Should not happen
    };
    let index = (sib & 0b00111000) >> 3;
    let base = sib & 0b00000111;

    let index_reg = if index == 4 {
        // No index
        "".to_string()
    } else {
        self.decode_sib_register(index, rex)?
    };

    let base_reg = if base == 5 && (self.is_no_base_required()) {
        // Special case: no base, disp32 follows
        "".to_string()
    } else {
        self.decode_register_field(base, rex)?
    };

    Ok((
        base_reg,
        index_reg,
        scale,
        0, // Placeholder for displacement if any
        0, // Placeholder for additional offset
    ))
}

/// Check if no base is required based on current decoding context
fn is_no_base_required(&self) -> bool {
    // Placeholder: Logic to determine if no base is required
    false
}

/// Decode SIB register with REX considerations
fn decode_sib_register(&self, reg: u8, rex: &RexPrefix) -> Result<String, DecodeError> {
    let extended = if rex.has_x_extension() { 8 } else { 0 };
    let registers = [
        "rax", "rcx", "rdx", "rbx", "rsp", "rbp", "rsi", "rdi",
        "r8", "r9", "r10", "r11", "r12", "r13", "r14", "r15",
    ];
    if (reg as usize) + extended >= registers.len() {
        return Err(DecodeError::InvalidRegister);
    }
    Ok(registers[(reg as usize) + extended].to_string())
}

/// Read immediate value from bytes based on operand size
fn read_immediate(
    &self,
    bytes: &[u8],
    offset: usize,
    size: &OperandSize,
) -> Result<u64, DecodeError> {
    match size {
        OperandSize::Byte => {
            if offset >= bytes.len() {
                return Err(DecodeError::IncompleteInstruction);
            }
            Ok(bytes[offset] as u64)
        }
        OperandSize::Word => {
            if offset + 1 >= bytes.len() {
                return Err(DecodeError::IncompleteInstruction);
            }
            Ok(u16::from_le_bytes([bytes[offset], bytes[offset + 1]]) as u64)
        }
        OperandSize::Dword => {
            if offset + 3 >= bytes.len() {
                return Err(DecodeError::IncompleteInstruction);
            }
            Ok(u32::from_le_bytes([
                bytes[offset],
                bytes[offset + 1],
                bytes[offset + 2],
                bytes[offset + 3],
            ]) as u64)
        }
        OperandSize::Qword => {
            if offset + 7 >= bytes.len() {
                return Err(DecodeError::IncompleteInstruction);
            }
            Ok(u64::from_le_bytes([
                bytes[offset],
                bytes[offset + 1],
                bytes[offset + 2],
                bytes[offset + 3],
                bytes[offset + 4],
                bytes[offset + 5],
                bytes[offset + 6],
                bytes[offset + 7],
            ]))
        }
    }
}
}

/// Structure representing the REX prefix
#[derive(Debug, Default)]
struct RexPrefix {
w: bool, // 64-bit operand
r: bool, // Extension of ModR/M reg field
x: bool, // Extension of SIB index field
b: bool, // Extension of ModR/M r/m field
}

impl OperandSize {
fn bytes(&self) -> usize {
match self {
OperandSize::Byte => 1,
OperandSize::Word => 2,
OperandSize::Dword => 4,
OperandSize::Qword => 8,
}
}
}

impl RexPrefix {
/// Create REX prefix from a byte
fn from_byte(byte: u8) -> Result<Self, DecodeError> {
if byte & 0b11001111 != 0x40 {
return Err(DecodeError::InvalidRexPrefix);
}
Ok(RexPrefix {
w: byte & 0b00001000 != 0,
r: byte & 0b00000100 != 0,
x: byte & 0b00000010 != 0,
b: byte & 0b00000001 != 0,
})
}


/// Check if any REX field is present
fn is_present(&self) -> bool {
    self.w || self.r || self.x || self.b
}

/// Extension for reg field
fn has_r_extension(&self) -> bool {
    self.r
}

/// Extension for r/m or base field
fn has_b_extension(&self) -> bool {
    self.b
}

/// Extension for index field in SIB
fn has_x_extension(&self) -> bool {
    self.x
}
}

/// Errors that can occur during decoding
#[derive(Debug)]
pub enum DecodeError {
UnknownOpcode,
IncompleteInstruction,
InvalidRegister,
MissingModRM,
InvalidRexPrefix,
InvalidSIB,
// More error variants can be added as needed
}

impl std::fmt::Display for DecodeError {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
match &self {
DecodeError::UnknownOpcode => write!(f, "Unknown opcode"),
DecodeError::IncompleteInstruction => write!(f, "Incomplete instruction bytes"),
DecodeError::InvalidRegister => write!(f, "Invalid register encoding"),
DecodeError::MissingModRM => write!(f, "Missing ModR/M byte for instruction"),
DecodeError::InvalidRexPrefix => write!(f, "Invalid REX prefix byte"),
DecodeError::InvalidSIB => write!(f, "Invalid SIB byte"),
}
}
}

impl std::error::Error for DecodeError {}