zbasefind/src/main.rs

use std::fs::File;
use std::io::Write;
use std::os::fd::AsRawFd;
use std::path::PathBuf;

use clap::Parser;
use miette::{IntoDiagnostic, Result};

use zbasefind as lib;

fn get_cpu_cores() -> usize {
    nix::sched::sched_getaffinity(nix::unistd::Pid::this())
        .map(|x| {
            (0..nix::sched::CpuSet::count())
                .filter(|i| x.is_set(*i).unwrap())
                .count()
        })
        .unwrap_or(1)
}

fn page_size_parser(s: &str) -> std::result::Result<u32, String> {
    let page_size: u32 = s.parse().map_err(|_| format!("`{s}` isn't a number"))?;
    if page_size.count_ones() == 1 {
        Ok(page_size)
    } else {
        Err(format!(
            "page size must be a power of two: got {page_size:08x}"
        ))
    }
}

#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
    /// Assume target is big endian
    #[arg(short, long, default_value_t = false)]
    big_endian: bool,

    /// Target memory page size
    #[arg(short, long, default_value_t = 0x1000, value_parser = page_size_parser)]
    page_size: u32,

    /// Minimum string length
    #[arg(short = 'm', long, default_value_t = 10)]
    min_str_len: usize,

    /// Maximum matches to display
    #[arg(short = 'n', long, default_value_t = 10)]
    max_matches: usize,

    /// Number of threads to use
    #[arg(short, long, default_value_t = get_cpu_cores())]
    threads: usize,

    /// File to scan
    file: PathBuf,
}

fn main() -> Result<()> {
    let args = Args::parse();

    let term = if nix::unistd::isatty(std::io::stderr().as_raw_fd()).unwrap_or(false) {
        term::stderr()
    } else {
        None
    };

    let mut file = File::open(args.file).into_diagnostic()?;
    let strings = lib::get_strings(&mut file, args.min_str_len).into_diagnostic()?;
    eprintln!("Found {} strings", strings.len());
    let pointers = lib::get_pointers(&mut file, args.big_endian).into_diagnostic()?;
    eprintln!("Found {} pointers", pointers.len());

    let start = 0x0000_0000u64;
    let end = 0x1_0000_0000u64;
    let total_pages = ((end - start) / (args.page_size as u64)) as usize;
    let chunk_size = (end - start) / (args.threads as u64);

    let progress_counters: Vec<_> = (0..args.threads)
        .map(|_| lib::ComputeProgress::new())
        .collect();

    let mut thread_results = std::thread::scope(|s| {
        let ranges = (start..=end)
            .step_by(chunk_size as usize)
            .zip((start + chunk_size..=end).step_by(chunk_size as usize));

        let tasks: Vec<_> = ranges
            .zip(progress_counters.iter())
            .map(|((start, end), counter)| {
                let strings = &strings;
                let pointers = &pointers;
                s.spawn(move || {
                    lib::compute_matches(
                        strings,
                        pointers,
                        start,
                        end,
                        args.page_size,
                        Some(counter),
                    )
                })
            })
            .collect();

        if let Some(mut term) = term {
            loop {
                std::thread::sleep(std::time::Duration::from_millis(100));
                if tasks.iter().any(|thread| !thread.is_finished()) {
                    term.carriage_return().unwrap();
                    term.delete_line().unwrap();

                    let completed_pages: usize = progress_counters
                        .iter()
                        .map(|prg| prg.num_completed())
                        .sum();

                    eprint!(
                        "{} / {} ({}%)",
                        completed_pages,
                        total_pages,
                        completed_pages * 100 / total_pages
                    );
                    std::io::stderr().flush().unwrap();
                } else {
                    break;
                }
            }

            term.carriage_return().unwrap();
            term.delete_line().unwrap();
            eprintln!("Scan complete");
        } else {
            eprintln!("Scanning...");
        }

        let thread_results: Vec<_> = tasks
            .into_iter()
            .map(|thread| thread.join().map_err(std::panic::resume_unwind).unwrap())
            .collect::<Vec<_>>();
        thread_results
    });

    eprintln!("Results:");

    if thread_results.iter().map(|x| x.len()).sum::<usize>() == 0 {
        eprintln!("  No results :(");
    } else {
        let mut results_buf = vec![];
        for _ in 0..args.max_matches {
            results_buf.extend(thread_results.iter_mut().filter_map(|res| res.pop()));
            results_buf.sort();
            match results_buf.pop() {
                None => break,
                Some((ct, addr)) => println!("{:08x}: {}", addr, ct),
            }
        }
    }

    Ok(())
}