Fix the freezes

use eframe::{egui, Frame};
use memmap::Mmap;
use std::fs;
use iced_x86::{Decoder as IcedDecoder, DecoderOptions, Formatter, Instruction as IcedInstruction};
use std::sync::{Arc, Mutex};
use std::collections::{HashSet, HashMap};
use std::str;
use log::info;
use env_logger;
use rfd::FileDialog;
use std::thread;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize logging
    tracing_subscriber::fmt::init();

    let app = ReaverApp::default();
    let native_options = eframe::NativeOptions::default();
    Ok(eframe::run_native(
        "Reaver", // Replace with your app name
        native_options,
        Box::new(|cc| Ok(Box::new(app))),
    )?)
}

// Application State
struct ReaverApp {
    analysis: Option<Arc<Mutex<BinaryAnalysis>>>,
    selected_function: Option<usize>,
    selected_string: Option<usize>,
    selected_exception: Option<usize>,
}

impl Default for ReaverApp {
    fn default() -> Self {
        Self {
            analysis: None,
            selected_function: None,
            selected_string: None,
            selected_exception: None,
        }
    }
}

impl eframe::App for ReaverApp {


    fn update(&mut self, ctx: &egui::Context, frame: &mut Frame) {
        self.draw_menu_bar(ctx);
        self.draw_side_panel(ctx);
        self.draw_main_panel(ctx);
        self.draw_status_bar(ctx);
    }
}

impl ReaverApp {
    fn draw_menu_bar(&mut self, ctx: &egui::Context) {
        egui::TopBottomPanel::top("menu_bar").show(ctx, |ui| {
            ui.horizontal(|ui| {
                if ui.button("Open").clicked() {
                    self.open_file();
                }
                if ui.button("Exit").clicked() {
                    std::process::exit(0);
                }
            });
        });
    }

    fn draw_side_panel(&mut self, ctx: &egui::Context) {
        egui::SidePanel::left("navigation_panel").show(ctx, |ui| {
            if let Some(analysis) = &self.analysis {
                let analysis = analysis.lock().unwrap();
                ui.heading("Functions");
                egui::ScrollArea::vertical().show(ui, |ui| {
                    for (index, function) in analysis.functions.iter().enumerate() {
                        if ui
                            .selectable_label(
                                self.selected_function == Some(index),
                                format!("{:#X} {}", function.address, function.name),
                            )
                            .clicked()
                        {
                            self.selected_function = Some(index);
                            self.selected_string = None;
                            self.selected_exception = None;
                        }
                    }
                });
                ui.separator();
                ui.heading("Strings");
                egui::ScrollArea::vertical().show(ui, |ui| {
                    for (index, string_entry) in analysis.strings.iter().enumerate() {
                        if ui
                            .selectable_label(
                                self.selected_string == Some(index),
                                format!("{:#X} {}", string_entry.address, string_entry.value),
                            )
                            .clicked()
                        {
                            self.selected_string = Some(index);
                            self.selected_function = None;
                            self.selected_exception = None;
                        }
                    }
                });
                ui.separator();
                ui.heading("Exception Handlers");
                egui::ScrollArea::vertical().show(ui, |ui| {
                    for (index, exception) in analysis.exception_handlers.iter().enumerate() {
                        if ui
                            .selectable_label(
                                self.selected_exception == Some(index),
                                format!("{:#X} Handler at {:#X}", exception.address, exception.handler_address),
                            )
                            .clicked()
                        {
                            self.selected_exception = Some(index);
                            self.selected_function = None;
                            self.selected_string = None;
                        }
                    }
                });
            } else {
                ui.label("No binary loaded.");
            }
        });
    }

    fn draw_main_panel(&mut self, ctx: &egui::Context) {
        egui::CentralPanel::default().show(ctx, |ui| {
            if let Some(analysis) = &self.analysis {
                let analysis = analysis.lock().unwrap();
                if let Some(index) = self.selected_function {
                    let function = &analysis.functions[index];
                    ui.heading(format!("Function {:#X} {}", function.address, function.name));
                    egui::ScrollArea::vertical().show(ui, |ui| {
                        for instruction in &function.instructions {
                            ui.horizontal(|ui| {
                                ui.label(format!("{:#016X}:", instruction.address));
                                ui.monospace(&instruction.text);
                            });
                        }
                    });
                } else if let Some(index) = self.selected_string {
                    let string_entry = &analysis.strings[index];
                    ui.heading(format!("String at {:#X}", string_entry.address));
                    ui.monospace(&string_entry.value);
                } else if let Some(index) = self.selected_exception {
                    let exception = &analysis.exception_handlers[index];
                    ui.heading(format!("Exception Handler at {:#X}", exception.address));
                    ui.label(format!("End Address: {:#X}", exception.end_address));
                    ui.label(format!("Handler Address: {:#X}", exception.handler_address));
                } else {
                    ui.label("Select an item from the navigation panel.");
                }
            } else {
                ui.label("No binary loaded.");
            }
        });
    }

    fn draw_status_bar(&mut self, ctx: &egui::Context) {
        egui::TopBottomPanel::bottom("status_bar").show(ctx, |ui| {
            ui.horizontal(|ui| {
                ui.label("Status: Ready");
            });
        });
    }

    fn open_file(&mut self) {
        if let Some(file_path) = FileDialog::new()
            .add_filter("Executable", &["exe", "dll", "bin", "elf"])
            .pick_file()
        {
            let file_path_str = file_path.to_string_lossy().to_string();
            info!("Selected file: {}", file_path_str);

            let analysis = Arc::new(Mutex::new(BinaryAnalysis::new()));
            let analysis_clone = analysis.clone();
            thread::spawn(move || {
                let mut analysis = analysis_clone.lock().unwrap();
                if let Err(e) = analysis.load_binary(&file_path_str) {
                    eprintln!("Failed to load binary: {}", e);
                }
            });
            self.analysis = Some(analysis);
            self.selected_function = None;
            self.selected_string = None;
            self.selected_exception = None;
        }
    }
}

// Core Analysis Structures and Functions

struct BinaryAnalysis {
    functions: Vec<Function>,
    strings: Vec<StringEntry>,
    exception_handlers: Vec<ExceptionHandler>,
}

impl BinaryAnalysis {
    pub fn new() -> Self {
        Self {
            functions: Vec::new(),
            strings: Vec::new(),
            exception_handlers: Vec::new(),
        }
    }

    pub fn load_binary(&mut self, path: &str) -> Result<(), Box<dyn std::error::Error>> {
        info!("Loading binary from '{}'", path);
        let loader = BinaryLoader::load(path)?;
        self.process_binary(loader);
        Ok(())
    }

    fn process_binary(&mut self, loader: BinaryLoader) {
        for section in &loader.sections {
            if section.is_executable {
                info!("Disassembling section '{}'", section.name);
                let disassembler = Disassembler::new(&section.data, section.addr, loader.bitness);
                let instructions = disassembler.disassemble();
                info!(
                    "Disassembled {} instructions in section '{}'",
                    instructions.len(),
                    section.name
                );
                self.detect_functions(&instructions);
            }
        }
        self.extract_strings(&loader);
        self.analyze_exception_handlers(&loader);
    }

    fn detect_functions(&mut self, instructions: &[Instruction]) {
        let mut function_starts = HashSet::new();
        let mut addr_to_instr = HashMap::new();

        // Identify potential function entry points
        for instr in instructions {
            let mnemonic = instr.instr.mnemonic();
            // Simple heuristic: Looking for 'push rbp' or 'sub rsp, imm' as function prologue
            if mnemonic == iced_x86::Mnemonic::Push && instr.instr.op0_register() == iced_x86::Register::RBP {
                function_starts.insert(instr.address);
            } else if mnemonic == iced_x86::Mnemonic::Sub && instr.instr.op0_register() == iced_x86::Register::RSP {
                function_starts.insert(instr.address);
            }
            addr_to_instr.insert(instr.address, instr.clone());
        }

        // Build functions by following the control flow
        for &start_addr in &function_starts {
            if !self.functions.iter().any(|f| f.address == start_addr) {
                let function = self.build_function(start_addr, &addr_to_instr);
                self.functions.push(function);
            }
        }
    }

    fn build_function(&self, start_addr: u64, addr_to_instr: &HashMap<u64, Instruction>) -> Function {
        let mut visited = HashSet::new();
        let mut to_visit = vec![start_addr];
        let mut instructions = Vec::new();

        while let Some(addr) = to_visit.pop() {
            if visited.contains(&addr) {
                continue;
            }
            if let Some(instr) = addr_to_instr.get(&addr) {
                instructions.push(instr.clone());
                visited.insert(addr);

                let mnemonic = instr.instr.mnemonic();
                let instr_len = instr.instr.len() as u64;

                // Determine next addresses to visit based on the instruction type
                if mnemonic == iced_x86::Mnemonic::Jmp {
                    if let target_addr = instr.instr.near_branch_target() {
                        to_visit.push(target_addr);
                    }
                } else if mnemonic == iced_x86::Mnemonic::Call {
                    // Continue with next instruction
                    to_visit.push(addr + instr_len);
                } else if mnemonic == iced_x86::Mnemonic::Ret {
                    // End of function
                    continue;
                } else {
                    // Default case: proceed to next instruction
                    to_visit.push(addr + instr_len);
                }
            }
        }

        Function {
            address: start_addr,
            name: format!("sub{:X}", start_addr),
            instructions,
        }
    }

    fn extract_strings(&mut self, loader: &BinaryLoader) {
        info!("Extracting strings from binary");
        let min_length = 4;
        for section in &loader.sections {
            if !section.is_executable {
                let data = &section.data;
                let mut offset = 0;
                while offset < data.len() {
                    let byte = data[offset];
                    if byte.is_ascii() && !byte.is_ascii_control() {
                        let start = offset;
                        while offset < data.len() && data[offset].is_ascii_graphic() {
                            offset += 1;
                        }
                        if offset - start >= min_length {
                            if let Ok(string) = str::from_utf8(&data[start..offset]) {
                                self.strings.push(StringEntry {
                                    address: section.addr + start as u64,
                                    value: string.to_string(),
                                });
                            }
                        }
                    } else {
                        offset += 1;
                    }
                }
            }
        }
        info!("Extracted {} strings", self.strings.len());
    }

    fn analyze_exception_handlers(&mut self, loader: &BinaryLoader) {
        info!("Analyzing exception handlers");
            info!("Exception handler analysis is currently implemented only for PE files");
        }
}

// Data Structures

struct Function {
    address: u64,
    name: String,
    instructions: Vec<Instruction>,
}

#[derive(Clone)]
struct Instruction {
    address: u64,
    text: String,
    instr: IcedInstruction,
}

struct StringEntry {
    address: u64,
    value: String,
}

struct ExceptionHandler {
    address: u64,
    end_address: u64,
    handler_address: u64,
}

// Binary Loader

struct BinaryLoader {
    entry_point: u64,
    sections: Vec<BinarySection>,
    bitness: u32,
    mmap: memmap::Mmap,
}

struct BinarySection {
    name: String,
    addr: u64,
    size: u64,
    data: Vec<u8>,
    is_executable: bool,
}
impl BinaryLoader {
    pub fn load(path: &str) -> Result<Self, Box<dyn std::error::Error>> {
        use goblin::{Object, pe, elf};
        use goblin::pe::section_table::IMAGE_SCN_MEM_EXECUTE;

        let file = fs::File::open(path)?;
        let mmap = unsafe { Mmap::map(&file)? };
        let buffer = &mmap[..];

        let object = Object::parse(buffer)?;

        match object {
            Object::PE(pe) => {
                // Process PE file
                let image_base = pe.image_base as u64;
                let entry_point = image_base + pe.entry as u64;

                let mut sections = Vec::new();
                for section in &pe.sections {
                    let name = section.name().unwrap_or_default().to_string();
                    let addr = image_base + section.virtual_address as u64;
                    let size = section.virtual_size as u64;
                    let is_executable = section.characteristics & IMAGE_SCN_MEM_EXECUTE != 0;

                    let raw_data_ptr = section.pointer_to_raw_data as usize;
                    let raw_data_size = section.size_of_raw_data as usize;

                    let data = if raw_data_ptr + raw_data_size <= buffer.len() {
                        buffer[raw_data_ptr..raw_data_ptr + raw_data_size].to_vec()
                    } else {
                        Vec::new()
                    };

                    sections.push(BinarySection {
                        name,
                        addr,
                        size,
                        data,
                        is_executable,
                    });
                }

                let bitness = if pe.is_64 { 64 } else { 32 };

                Ok(BinaryLoader {
                    entry_point,
                    sections,
                    bitness,
                    mmap,
                })
            }
            Object::Elf(elf) => {
                // Process ELF file
                let entry_point = elf.entry as u64;
                let mut sections = Vec::new();

                for header in &elf.section_headers {
                    let name = if let Some(Ok(name)) = elf.shdr_strtab.get(header.sh_name) {
                        name.to_string()
                    } else {
                        "".to_string()
                    };

                    let addr = header.sh_addr as u64;
                    let size = header.sh_size as u64;
                    let is_executable = header.is_executable();

                    let offset = header.sh_offset as usize;
                    let size = header.sh_size as usize;

                    let data = if offset + size <= buffer.len() {
                        buffer[offset..offset + size].to_vec()
                    } else {
                        Vec::new()
                    };

                    sections.push(BinarySection {
                        name,
                        addr,
                        size: size.try_into().unwrap(),
                        data,
                        is_executable,
                    });
                }

                let bitness = match elf.header.e_ident[elf::header::EI_CLASS] {
                    elf::header::ELFCLASS64 => 64,
                    elf::header::ELFCLASS32 => 32,
                    _ => return Err("Unknown ELF class".into()),
                };

                Ok(BinaryLoader {
                    entry_point,
                    sections,
                    bitness,
                    mmap,
                })
            }
            _ => Err("Unsupported binary format".into()),
        }
    }
}


// Disassembler

struct Disassembler<'a> {
    bitness: u32,
    code: &'a [u8],
    address: u64,
}

impl<'a> Disassembler<'a> {
    pub fn new(code: &'a [u8], address: u64, bitness: u32) -> Self {
        Self { bitness, code, address }
    }

    pub fn disassemble(&self) -> Vec<Instruction> {
        let mut decoder = IcedDecoder::with_ip(self.bitness, self.code, self.address, DecoderOptions::NONE);
        let mut formatter = iced_x86::IntelFormatter::new();
        let mut output = String::new();
        let mut instructions = Vec::new();

        while decoder.can_decode() {
            let instr = decoder.decode();
            output.clear();
            formatter.format(&instr, &mut output);
            let instruction = Instruction {
                address: instr.ip(),
                text: output.clone(),
                instr,
            };
            instructions.push(instruction);
        }
        instructions
    }
}

PS C:\Users\Lukas\source\repos\GONE> cargo run --release
warning: unused import: `env_logger`
  --> src/main.rs:10:5
   |
10 | use env_logger;
   |     ^^^^^^^^^^
   |
   = note: `#[warn(unused_imports)]` on by default

warning: unused import: `pe`
   --> src/main.rs:385:30
    |
385 |         use goblin::{Object, pe, elf};
    |                              ^^

warning: use of deprecated method `goblin::strtab::Strtab::<'a>::get`: Bad performance, use get_at() instead
   --> src/main.rs:440:72
    |
440 |                     let name = if let Some(Ok(name)) = elf.shdr_strtab.get(header.sh_name) {
    |                                                                        ^^^
    |
    = note: `#[warn(deprecated)]` on by default

warning: unused variable: `cc`
  --> src/main.rs:23:19
   |
23 |         Box::new(|cc| Ok(Box::new(app))),
   |                   ^^ help: if this is intentional, prefix it with an underscore: `_cc`
   |
   = note: `#[warn(unused_variables)]` on by default

warning: unused variable: `frame`
  --> src/main.rs:49:47
   |
49 |     fn update(&mut self, ctx: &egui::Context, frame: &mut Frame) {
   |                                               ^^^^^ help: if this is intentional, prefix it with an underscore: `_frame`

warning: irrefutable `if let` pattern
   --> src/main.rs:281:24
    |
281 |                     if let target_addr = instr.instr.near_branch_target() {
    |                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |
    = note: this pattern will always match, so the `if let` is useless
    = help: consider replacing the `if let` with a `let`
    = note: `#[warn(irrefutable_let_patterns)]` on by default

warning: unused variable: `loader`
   --> src/main.rs:335:46
    |
335 |     fn analyze_exception_handlers(&mut self, loader: &BinaryLoader) {
    |                                              ^^^^^^ help: if this is intentional, prefix it with an underscore: `_loader`

warning: unused variable: `size`
   --> src/main.rs:447:25
    |
447 |                     let size = header.sh_size as u64;
    |                         ^^^^ help: if this is intentional, prefix it with an underscore: `_size`

warning: fields `entry_point` and `mmap` are never read
   --> src/main.rs:370:5
    |
369 | struct BinaryLoader {
    |        ------------ fields in this struct
370 |     entry_point: u64,
    |     ^^^^^^^^^^^
...
373 |     mmap: memmap::Mmap,
    |     ^^^^
    |
    = note: `#[warn(dead_code)]` on by default

warning: field `size` is never read
   --> src/main.rs:379:5
    |
376 | struct BinarySection {
    |        ------------- field in this struct
...
379 |     size: u64,
    |     ^^^^

warning: `GONE` (bin "GONE") generated 10 warnings (run `cargo fix --bin "GONE"` to apply 2 suggestions)
    Finished `release` profile [optimized] target(s) in 0.19s
     Running `target\release\GONE.exe`
2024-09-20T01:32:29.402315Z  INFO GONE: Selected file: C:\Users\Lukas\Downloads\roblox.dll    
2024-09-20T01:32:29.402559Z  INFO GONE: Loading binary from 'C:\Users\Lukas\Downloads\roblox.dll'
2024-09-20T01:32:29.402759Z  INFO GONE: Disassembling section '.text'
2024-09-20T01:32:29.404090Z  INFO GONE: Disassembled 1746 instructions in section '.text'
2024-09-20T01:32:29.404609Z  INFO GONE: Extracting strings from binary
( Freezes here )