patchlaa 0.1.0

use std::fs::OpenOptions;
use std::path::PathBuf;

use clap::Parser;
use color_eyre::eyre::{self, Context, bail};
use goblin::pe::PE;
use goblin::pe::characteristic::IMAGE_FILE_LARGE_ADDRESS_AWARE;
use memmap2::{Mmap, MmapMut};
use scroll::Pwrite;

#[derive(Parser)]
struct Args {
    #[arg(
        short = 'n',
        help = "Only check if the file has the LAA flag set, don't write anything.",
        default_value_t = false
    )]
    pub dry_run: bool,

    #[arg(
        short,
        long,
        help = "Save new executable to this file instead of patching in place."
    )]
    pub output: Option<PathBuf>,

    #[arg(index = 1, required = true, help = "File to patch.")]
    pub file: PathBuf,
}

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

    let file = OpenOptions::new()
        .read(true)
        .write(args.output.is_none() && !args.dry_run)
        .append(false)
        .truncate(false)
        .create(false)
        .open(&args.file)
        .wrap_err("failed to open executable file")?;
    let source = unsafe { Mmap::map(&file) }.expect("failed to map file");

    let mut executable = PE::parse(&*source).expect("failed to parse executable");
    let char = &mut executable.header.coff_header.characteristics;
    let is_laa = (*char & IMAGE_FILE_LARGE_ADDRESS_AWARE) != 0;

    if is_laa {
        println!("executable already has the LAA flag set");
        if !args.dry_run {
            bail!("executable already has the LAA (Large Address Aware, 0x0020) flag set");
        }
    } else {
        println!("executable does not have the LAA flag set, patching");
        if !args.dry_run {
            *char |= IMAGE_FILE_LARGE_ADDRESS_AWARE;

            let file = match args.output.as_ref() {
                Some(path) => OpenOptions::new()
                    .write(true)
                    .create_new(true)
                    .open(path)
                    .wrap_err("failed to open output file")?,
                None => file,
            };
            let mut output = unsafe { MmapMut::map_mut(&file) }.expect("failed to map file");

            output
                .pwrite(executable, 0)
                .wrap_err("failed to save executable")?;
        }
    }

    println!("all done!");
    Ok(())
}