#!/usr/bin/env python3

import sys
import struct

if len(sys.argv) < 3:
    print("XTChain ExeTool for modifying PE/COFF image subsystem\nNot sufficient parametrs. '[PE/COFF Image File]' '[SubSystem]'")
    sys.exit(1)

ImageFile = sys.argv[1]
Subsystem = sys.argv[2].upper()

# Set proper subsystem
if Subsystem == "UNKNOWN":
    ImageSubsystem = 0x00
elif Subsystem == "NT_NATIVE":
    ImageSubsystem = 0x01
elif Subsystem == "WINDOWS_GUI":
    ImageSubsystem = 0x02
elif Subsystem == "WINDOWS_CLI":
    ImageSubsystem = 0x03
elif Subsystem == "WINDOWS_CE_OLD":
    ImageSubsystem = 0x04
elif Subsystem == "OS2_CUI":
    ImageSubsystem = 0x05
elif Subsystem == "POSIX_CUI":
    ImageSubsystem = 0x07
elif Subsystem == "NATIVE_WINDOWS":
    ImageSubsystem = 0x08
elif Subsystem == "WINDOWS_CE_GUI":
    ImageSubsystem = 0x09
elif Subsystem == "EFI_APPLICATION":
    ImageSubsystem = 0x0A
elif Subsystem == "EFI_BOOT_SERVICE_DRIVER":
    ImageSubsystem = 0x0B
elif Subsystem == "EFI_RUNTIME_DRIVER":
    ImageSubsystem = 0x0C
elif Subsystem == "EFI_ROM":
    ImageSubsystem = 0x0D
elif Subsystem == "XBOX":
    ImageSubsystem = 0x0E
elif Subsystem == "WINDOWS_BOOT_APPLICATION":
    ImageSubsystem = 0x10
elif Subsystem == "XT_NATIVE_KERNEL":
    ImageSubsystem = 0x14
elif Subsystem == "XT_NATIVE_APPLICATION":
    ImageSubsystem = 0x15
elif Subsystem == "XT_NATIVE_DRIVER":
    ImageSubsystem = 0x16
elif Subsystem == "XT_DYNAMIC_LIBRARY":
    ImageSubsystem = 0x17
elif Subsystem == "XT_APPLICATION_CLI":
    ImageSubsystem = 0x18
elif Subsystem == "XT_APPLICATION_GDI":
    ImageSubsystem = 0x19
else:
    print("Invalid subsystem privided")
    exit(2)

# Open PE/COFF image file
PeImage = open(sys.argv[1], "r+b")

# Get PE header
PeImage.seek(0x3C)
(PeHeader,)=struct.unpack("H", PeImage.read(2))

# Get PE signature
PeImage.seek(PeHeader)
(PeSignature,)=struct.unpack("I", PeImage.read(4))
if PeSignature != 0x4550:
    print("Invalid or corrupted PE header")

# Set new image subsystem
PeImage.seek(PeHeader + 0x5C)
print("Setting subsystem to " + str(ImageSubsystem))
PeImage.write(struct.pack("H", ImageSubsystem))

# Close PE/COFF image file
PeImage.close()