import sys
import os
import struct
import pefile

def disable_aslr(filename):
    # Read file as binary
    pe = pefile.PE(filename)
    with open(filename, 'rb') as f:
        data = bytearray(f.read())

    elfa_new = struct.unpack("<I", data[0x3C:0x40])[0]
    machine = struct.unpack("<H", data[elfa_new+4:elfa_new+6])[0]

    dllCharacteristicsOffset = pe.get_offset_from_elfa(elfa_new)
    dllCharacteristics = struct.unpack("<H", data[dllCharacteristicsOffset:dllCharacteristicsOffset+2])[0]
    dllCharacteristics = dllCharacteristics & ~0x40 # Clear the ASLR bit
    data[dllCharacteristicsOffset:dllCharacteristicsOffset+2] = struct.pack("<H", dllCharacteristics)

    # Create new filename with .noaslr before extension
    base, ext = os.path.splitext(filename)
    new_filename = f"{base}.noaslr{ext}"

    # Save modified file
    with open(new_filename, 'wb') as f:
        f.write(data)

    print(f"Saved: {new_filename}")

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python disable_aslr.py <filename>")
        sys.exit(1)
    filename = sys.argv[1]
    disable_aslr(filename)