import struct

__all__ = ['DataType', 'Array', 'Struct', 'Byte', 'Word', 'Dword', 'ArrayOf',
           'RT_RCDATA', 'RT_VERSION', 'UpdateResource', 'SetSubsystem',
           ]

try:
    # Python 2.3+
    issubclass(type, ())
except TypeError:
    # Python 2.2
    import __builtin__
    def issubclass(subclass, class_or_tuple):
        if isinstance(class_or_tuple, tuple):
            for class_ in class_or_tuple:
                if issubclass(subclass, class_):
                    return True
            return False
        return __builtin__.issubclass(subclass, class_or_tuple)

class DataType(type):
    def __init__(cls, name, bases, namespace):
        if Struct in bases:
            assert '__fields__' in namespace
            slots = []
            format = '<'
            for datatype, field in namespace['__fields__']:
                slots.append(field)
                if issubclass(datatype, (Struct, Array)):
                    format += '%ds' % datatype.__size__
                else:
                    format += datatype.__format__[1:]
            cls.__slots__ = tuple(slots)
            cls.__format__ = format
        else:
            assert '__format__' in namespace
            cls.__slots__ = ()
        # Update (cache) the record size
        cls.__size__ = struct.calcsize(cls.__format__)
        return super(DataType, cls).__init__(name, bases, namespace)

# allow the DataType metaclass to do its thing on the 'Struct' class
Struct = None
class Struct(object):
    __metaclass__ = DataType
    __format__ = ''

    def __init__(self, bytes=None):
        if bytes is None:
            bytes = '\0' * self.__size__
        values = struct.unpack(self.__format__, bytes)
        for (datatype, name), value in zip(self.__fields__, values):
            setattr(self, name, datatype(value))
        return

    def __repr__(self):
        values = []
        for format, name in self.__fields__:
            value = getattr(self, name)
            if issubclass(format, Struct):
                members = repr(value).split('\n')
                members.insert(0, '<%s>' % format.__name__)
                value = '\n    '.join(members)
            values.append('%s: %s' % (name, value))
        return '\n'.join(values)

    #@classmethod
    def load(cls, stream):
        return cls(stream.read(cls.__size__))
    load = classmethod(load)

    def pack(self):
        values = []
        for datatype, field in self.__fields__:
            value = getattr(self, field)
            if issubclass(datatype, (Struct, Array)):
                value = value.pack()
            values.append(value)
        return struct.pack(self.__format__, *values)

class Array(list):
    def __init__(self, bytes=None):
        if bytes is None:
            bytes = '\0' * self.__size__
        values = struct.unpack(self.__format__, bytes)
        if self.__itemtype__:
            values = map(self.__itemtype__, values)
        list.__init__(self, values)

    def pack(self):
        values = self
        if self.__itemtype__ is not None:
            values = [ item.pack() for item in values ]
        return struct.pack(self.__format__, *values)

class Byte(int):
    __metaclass__ = DataType
    __format__ = '<B'

class Word(int):
    __metaclass__ = DataType
    __format__ = '<H'

class Dword(long):
    __metaclass__ = DataType
    __format__ = '<L'

def ArrayOf(datatype, size):
    name = '%s[%d]' % (datatype.__name__, size)
    if issubclass(datatype, (Struct, Array)):
        format = '%ds' % datatype.__size__
        itemtype = datatype
    else:
        format = datatype.__format__[1:]
        itemtype = None
    namespace = {
        '__metaclass__' : DataType,
        '__format__' : '<' + (format * size),
        '__itemtype__' : itemtype,
        }
    return DataType(name, (Array,), namespace)

IMAGE_NT_OPTIONAL_HDR32_MAGIC = 0x010B
IMAGE_NUMBEROF_DIRECTORY_ENTRIES = 16
IMAGE_SIZEOF_SHORT_NAME = 8
IMAGE_DIRECTORY_ENTRY_RESOURCE = 2

IMAGE_SUBSYSTEM_WINDOWS_GUI = 2
IMAGE_SUBSYSTEM_WINDOWS_CUI = 3

# Predefined Resource Types
RT_RCDATA = 10
RT_VERSION = 16

class ImageDosHeader(Struct):
    __fields__ = [
        (Word, 'e_magic'),
        (Word, 'e_cblp'),
        (Word, 'e_cp'),
        (Word, 'e_crlc'),
        (Word, 'e_cparhdr'),
        (Word, 'e_minalloc'),
        (Word, 'e_maxalloc'),
        (Word, 'e_ss'),
        (Word, 'e_sp'),
        (Word, 'e_csnum'),
        (Word, 'e_ip'),
        (Word, 'e_cs'),
        (Word, 'e_lfarlc'),
        (Word, 'e_ovno'),
        (ArrayOf(Word, 4), 'e_res'),
        (Word, 'e_oemid'),
        (Word, 'e_oeminfo'),
        (ArrayOf(Word, 10), 'e_res2'),
        (Dword, 'e_lfanew'),
        ]

class ImageFileHeader(Struct):
    __fields__ = [
        (Word, 'Machine'),
        (Word, 'NumberOfSections'),
        (Dword, 'TimeDateStamp'),
        (Dword, 'PointerToSymbolTable'),
        (Dword, 'NumberOfSymbols'),
        (Word, 'SizeOfOptionalHeader'),
        (Word, 'Characteristics'),
        ]

class ImageDataDirectory(Struct):
    __fields__ = [
        (Dword, 'VirtualAddress'),
        (Dword, 'Size'),
        ]

class ImageOptionalHeader(Struct):
    __fields__ = [
        (Word, 'Magic'),
        (Byte, 'MajorLinkerVersion'),
        (Byte, 'MinorLinkerVersion'),
        (Dword, 'SizeOfCode'),
        (Dword, 'SizeOfInitializedData'),
        (Dword, 'SizeOfUninitializedData'),
        (Dword, 'AddressOfEntryPoint'),
        (Dword, 'BaseOfCode'),
        (Dword, 'BaseOfData'),
        (Dword, 'ImageBase'),
        (Dword, 'SectionAlignment'),
        (Dword, 'FileAlignment'),
        (Word, 'MajorOperatingSystemVersion'),
        (Word, 'MinorOperatingSystemVersion'),
        (Word, 'MajorImageVersion'),
        (Word, 'MinorImageVersion'),
        (Word, 'MajorSubsystemVersion'),
        (Word, 'MinorSubsystemVersion'),
        (Dword, 'Win32VersionValue'),
        (Dword, 'SizeOfImage'),
        (Dword, 'SizeOfHeaders'),
        (Dword, 'CheckSum'),
        (Word, 'Subsystem'),
        (Word, 'DllCharacteristics'),
        (Dword, 'SizeOfStackReserve'),
        (Dword, 'SizeOfStackCommit'),
        (Dword, 'SizeOfHeapReserve'),
        (Dword, 'SizeOfHepCommit'),
        (Dword, 'LoaderFlags'),
        (Dword, 'NumberOfRvaAndSizes'),
        (ArrayOf(ImageDataDirectory, IMAGE_NUMBEROF_DIRECTORY_ENTRIES),
         'DataDirectory'),
        ]

class ImageNTHeaders(Struct):
    __fields__ = [
        (Dword, 'Signature'),
        (ImageFileHeader, 'FileHeader'),
        (ImageOptionalHeader, 'OptionalHeader'),
        ]

class ImageSectionHeader(Struct):
    __fields__ = [
        (ArrayOf(Byte, IMAGE_SIZEOF_SHORT_NAME), 'Name'),
        (Dword, 'Misc'),
        (Dword, 'VirtualAddress'),
        (Dword, 'SizeOfRawData'),
        (Dword, 'PointerToRawData'),
        (Dword, 'PointerToRelocations'),
        (Dword, 'PointerToLinenumbers'),
        (Word, 'NumberOfRelocations'),
        (Word, 'NumberOfLinenumbers'),
        (Dword, 'Characteristics'),
        ]

class ImageResourceDirectory(Struct):
    __fields__ = [
        (Dword, 'Characteristics'),
        (Dword, 'TimeDateStamp'),
        (Word, 'MajorVersion'),
        (Word, 'MinorVersion'),
        (Word, 'NumberOfNamedEntries'),
        (Word, 'NumberOfIdEntries'),
        ]

class ImageResourceDirectoryEntry(Struct):
    __fields__ = [
        (Dword, 'Id'),
        (Dword, 'Offset'),
        ]

class ImageResourceDataEntry(Struct):
    __fields__ = [
        (Dword, 'OffsetToData'),
        (Dword, 'Size'),
        (Dword, 'CodePage'),
        (Dword, 'Reserved'),
        ]

def FindImageHeaders(image):
    # Read DOS header
    image.seek(0)
    dos_header = ImageDosHeader.load(image)
    if dos_header.e_magic != 0x5A4D: # "MZ"
        raise ValueError("not an EXE file")
    if dos_header.e_lfanew == 0:
        raise ValueError("not a PE file")
    # Read PE header and optional header
    image.seek(dos_header.e_lfanew)
    nt_headers = ImageNTHeaders.load(image)
    if nt_headers.Signature != 0x00004550: # "PE\0\0"
        raise ValueError("invalid PE signature")
    file_header = nt_headers.FileHeader
    if file_header.SizeOfOptionalHeader != ImageOptionalHeader.__size__:
        raise ValueError("invalid header size")
    opt_header = nt_headers.OptionalHeader
    if opt_header.Magic != IMAGE_NT_OPTIONAL_HDR32_MAGIC:
        raise ValueError("invaid header magic");
    return dos_header, nt_headers

def FindResourceSection(image):
    # Scan section headers for resource section
    dos_header, nt_headers = FindImageHeaders(image)
    data_directory = nt_headers.OptionalHeader.DataDirectory
    directory_entry = data_directory[IMAGE_DIRECTORY_ENTRY_RESOURCE]
    if directory_entry.VirtualAddress == 0 or directory_entry.Size == 0:
        raise ValueError("no resources")
    address = directory_entry.VirtualAddress
    for i in xrange(nt_headers.FileHeader.NumberOfSections):
        section = ImageSectionHeader.load(image)
        if address == section.VirtualAddress and section.SizeOfRawData != 0:
            return section
    raise ValueError("no resources")

def FindResourceOffset(image, anyId, resourceId, findSubdir):
    directory = ImageResourceDirectory.load(image)
    # Skip over named entries
    for i in xrange(directory.NumberOfNamedEntries):
        entry = ImageResourceDirectoryEntry.load(image)
    # Now process ID entries
    for i in xrange(directory.NumberOfIdEntries):
        entry = ImageResourceDirectoryEntry.load(image)
        if ((anyId or entry.Id == resourceId) and
            ((entry.Offset & 0x80000000L != 0) == findSubdir)):
            return (entry.Offset & 0x7FFFFFFFL)
    raise ValueError("resource not found")

def SeekToResource(image, resourceType, resourceId):
    """Sets stream `f` to the location of the resource source"""
    section = FindResourceSection(image)

    # Scan the resource directory
    image.seek(section.PointerToRawData)
    offset = FindResourceOffset(image, False, resourceType, True)
    image.seek(section.PointerToRawData + offset)
    offset = FindResourceOffset(image, False, resourceId, True)
    image.seek(section.PointerToRawData + offset)
    offset = FindResourceOffset(image, True, 0, False)
    image.seek(section.PointerToRawData + offset)
    entry = ImageResourceDataEntry.load(image)

    # Sanity check: data_entry.OffsetToData is an RVA.  It's
    # technically possible for the RVA to point to a different section,
    # but we don't support that.
    offset = entry.OffsetToData - section.VirtualAddress
    if offset < 0 or section.SizeOfRawData < (offset + entry.Size):
        raise ValueError("invalid resource")

    # Seek to the resource
    image.seek(section.PointerToRawData + offset)
    return

def UpdateResource(image, resourceType, resourceId, resourceData):
    """
    Replaces the resource in the executable file.
    """
    SeekToResource(image, resourceType, resourceId)
    image.write(resourceData.pack())
    return

def SetSubsystem(image, subsystem):
    """
    Sets the Windows subsystem required to run the executable file.
    """
    if subsystem not in (IMAGE_SUBSYSTEM_WINDOWS_GUI,
                         IMAGE_SUBSYSTEM_WINDOWS_CUI):
        raise ValueError("invalid executable subsystem")

    dos_header, nt_headers = FindImageHeaders(image)
    nt_headers.OptionalHeader.Subsystem = subsystem
    image.seek(dos_header.e_lfanew)
    image.write(nt_headers.pack())
    return