Disable ASLR and make .nv_fatb section read-only to allow DLLs share address space and not overflow virtual memory on Windows

fix_pe.py

# Simple script to disable ASLR and make .nv_fatb sections read-only
# Requires: pefile  ( python -m pip install pefile )
# Usage:  fix_pe.py --input path/to/*.dll

import shutil
from argparse import ArgumentParser
from pathlib import Path

import pefile


def patch_headers(pe: pefile.PE) -> bool:
    aslr = pe.OPTIONAL_HEADER.IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE
    if aslr:
        print('- ASLR')

    sections = []
    for s in pe.sections:
        name = s.Name.decode()
        if not name.startswith('.nv_fatb'):
            continue
        size = s.Misc_VirtualSize
        writable = bool(s.Characteristics &
                        pefile.SECTION_CHARACTERISTICS['IMAGE_SCN_MEM_WRITE'])
        print(f'- {name}: {size/1_048_576:0.2f}MB,',
              'writable' if writable else 'read-only')
        if size > 0 and writable:
            sections += [s]

    if not aslr and not sections:  # PE is already good
        return False

    if aslr:  # Disable ASLR for DLL
        pe.OPTIONAL_HEADER.DllCharacteristics &= ~pefile.DLL_CHARACTERISTICS[
            'IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE'
        ]

    # Disable writing for all sections
    readonly_mask = ~pefile.SECTION_CHARACTERISTICS['IMAGE_SCN_MEM_WRITE']
    for s in sections:
        s.Characteristics &= readonly_mask

    return True


def main():
    parser = ArgumentParser(
        'Disable ASLR and make .nv_fatb sections read-only')
    parser.add_argument('root', type=Path)
    root = parser.parse_args().root

    for path in root.glob('*.dll'):
        bak = root / 'bak' / path.name
        new = path.with_suffix('.new')

        try:
            print(path.relative_to(root))
            with pefile.PE(path.as_posix(), fast_load=True) as pe:
                if not patch_headers(pe):
                    continue

                if bak.exists():  # Backup already exists, can't modify
                    print('- warning: backup file already exists, '
                          'skip modification')
                    continue

                pe.write(new)

            bak.parent.mkdir(exist_ok=True, parents=True)
            path.rename(bak)
            print(f'- {bak.relative_to(root)} backup created')

            new.rename(path)
            print(f'- {path.relative_to(root)} updated')

        except Exception as exc:  # noqa: PIE786
            print(f'- {path} failed with {exc!r}')


if __name__ == '__main__':
    main()