patchlaa 0.2.0

command line tool to force LAA for PE32/+ executables
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83

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

use clap::Parser;
use color_eyre::eyre::{self, Context};
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(long, help = "Unset the LAA flag instead of setting it")]
    pub no_laa: bool,

    #[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 && !args.no_laa {
        println!("executable already has the LAA flag set");
    } else if is_laa && args.no_laa {
        println!("executable has the LAA flag set, patching…");
    } else if !is_laa && args.no_laa {
        println!("executable already has the LAA flag unset");
    } else {
        println!("executable does not have the LAA flag set, patching…");
    }

    if !args.dry_run && is_laa == args.no_laa {
        // toggles the flag
        *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(())
}