import sys
sys.path.append('..\\..')

from PyStand import Struct
uint16 = Struct.uint16
uint32 = Struct.uint32

import XBE

CODE = 0x00000020
INITDATA = 0x00000040
UNINITDATA = 0x00000080
COMMENTS = 0x00000200
NOFINAL = 0x00000800
DISCARD = 0x02000000
SHAREABLE = 0x10000000
EXECUTABLE = 0x20000000
READABLE = 0x40000000
WRITABLE = 0x80000000

class ImageFileHeader(Struct):
	__endian__ = Struct.LE
	def __format__(self):
		self.machine, self.numberOfSections = uint16[2]
		self.timestamp, self.symtab, self.numsyms = uint32[3]
		self.optHdrSize, self.characteristics = uint16[2]

class ImageOptionalHeader(Struct):
	__endian__ = Struct.LE
	def __format__(self):
		self.Magic = uint16
		self.LinkerVersion = Struct.uint8[2]
		self.SizeOfCode, self.SizeOfInitializedData = uint32[2]
		self.SizeOfUninitializedData = uint32
		self.AddressOfEntryPoint = uint32
		self.BaseOfCode, self.BaseOfData = uint32[2]
		self.ImageBase, self.SectionAlignment, self.FileAlignment = uint32[3]
		self.OperatingSystemVersion = uint16[2]
		self.ImageVersion = uint16[2]
		self.SubsystemVersion = uint16[2]
		self.Reserved = uint32
		self.SizeOfImage, self.SizeOfHeaders, self.Checksum = uint32[3]
		self.Subsystem, self.DllCharacteristics = uint16[2]
		self.SizeOfStackReserve, self.SizeOfStackCommit = uint32[2]
		self.SizeOfHeapReserve, self.SizeOfHeapCommit = uint32[2]
		self.LoaderFlags, self.NumberOfRvaAndSizes = uint32[2]
		self.DataDirectory = uint32[32]
	
class ImageSectionHeader(Struct):
	__endian__ = Struct.LE
	def __format__(self):
		self.Name = Struct.string(8, stripNulls=True)
		self.VirtualSize, self.VirtualAddress = uint32[2]
		self.SizeOfRawData, self.RawData = uint32[2]
		self.Relocations, self.Linenumbers = uint32[2]
		self.NumRelocations, self.NumLinenumbers = uint16[2]
		self.Characteristics = uint32

class ImageImportDirectory(Struct):
	__endian__ = Struct.LE
	def __format__(self):
		self.RVAFunctionNameList = uint32
		self.useless = uint32[2]
		self.RVAModuleName = uint32
		self.RVAFunctionAddressList = uint32

# This is neither an XBE or PE structure, just a convenience struct
class XBEPEImportDirectory(Struct):
	__endian__ = Struct.LE
	def __format__(self):
		self.sclera, self.emptyimp = ImageImportDirectory()[2] # 0000
		self.dllname = Struct.string(16, value='KERNEL32.dll\0\0\0\0') # 0028
		self.importlist = self.uint32[4] # 0038
		self.iat = self.uint32[2] # 0048

class PE(object):
	def __init__(self, xbe):
		self.xbe = xbe
	
	def write(self, fp):
		fp.write('MZ')
		fp.write('\0' * 58)
		fp.write(uint32(64))
		
		imgheader = ImageFileHeader()
		imgheader.machine = 0x14C # x86
		imgheader.numberOfSections = len(self.xbe.sections) + 1 # +1 for import directory
		imgheader.timestamp = 0
		imgheader.symtab = 0
		imgheader.numsyms = 0
		imgheader.characteristics = 0x010F
		
		imgoptheader = ImageOptionalHeader()
		imgheader.optHdrSize = len(imgoptheader)
		
		imgoptheader.Magic = 0x010B
		imgoptheader.AddressOfEntryPoint = self.xbe.entry - self.xbe.base
		imgoptheader.ImageBase = self.xbe.base
		imgoptheader.SectionAlignment = 4
		imgoptheader.FileAlignment = 4
		imgoptheader.SubsystemVersion = [4, 0]
		imgoptheader.SizeOfImage = 0
		imgoptheader.Subsystem = 2 # Windows GUI
		imgoptheader.SizeOfStackCommit  = 0x1000
		imgoptheader.SizeOfHeapReserve  = 0x100000
		imgoptheader.NumberOfRvaAndSizes = 16
		
		imgsectheader = ImageSectionHeader()
		
		imgoptheader.SizeOfHeaders  = 64 + 4 + len(imgheader) + len(imgoptheader)
		imgoptheader.SizeOfHeaders += (len(imgsectheader) * imgheader.numberOfSections)
		
		dirBase = max(sect.addr + sect.size for sect in self.xbe.sections.values())
		imgoptheader.DataDirectory[2] = dirBase - self.xbe.base
		
		imgimpdirectory = XBEPEImportDirectory()
		imgimplen = len(imgimpdirectory)
		imgoptheader.DataDirectory[3] = imgimplen # len(ImageImportDirectory()) * 2
		
		imgoptheader.DataDirectory[24] = dirBase + 0x48 - self.xbe.base
		imgoptheader.DataDirectory[25] = 8
		
		imgoptheader.SizeOfCode = sum(len(sect.data) for sect in self.xbe.sections.values())
		imgoptheader.SizeOfCode += imgimplen
		imgoptheader.SizeOfImage = imgoptheader.SizeOfHeaders + imgoptheader.SizeOfCode
		
		fp.write('PE\0\0')
		fp.write(imgheader.pack())
		fp.write(imgoptheader.pack())
		
		sectpos = imgoptheader.SizeOfHeaders
		
		addrs = [sect.addr for sect in self.xbe.sections.values()]
		addrs.sort()
		
		sectByAddr = {}
		for sect in self.xbe.sections.values():
			sectByAddr[sect.addr] = sect
		
		for addr in addrs:
			sect = sectByAddr[addr]
			
			imgsectheader.Name = sect.name
			imgsectheader.VirtualSize = sect.size
			imgsectheader.VirtualAddress = sect.addr - self.xbe.base
			imgsectheader.SizeOfRawData = len(sect.data)
			imgsectheader.RawData = sectpos
			imgsectheader.Characteristics = READABLE
			imgsectheader.Characteristics |= (sect.flags & XBE.EXECUTABLE and EXECUTABLE) or 0
			imgsectheader.Characteristics |= (sect.flags & XBE.WRITABLE and WRITABLE) or 0
			imgsectheader.Characteristics |= (sect.flags & XBE.PRELOAD and INITDATA) or 0
			fp.write(imgsectheader.pack())
			sectpos += len(sect.data)
		
		imgimpdirectory.sclera.RVAModuleName = dirBase + 0x28 - self.xbe.base # Offset to 'sclera.dll\0\0'
		imgimpdirectory.sclera.RVAFunctionNameList = dirBase + 0x38 - self.xbe.base
		imgimpdirectory.sclera.RVAFunctionAddressList = dirBase + 0x48 - self.xbe.base
		imgimpdirectory.importlist[0] = 0x80000001
		imgimpdirectory.iat[0] = 0x80000001
		
		imgsectheader.Name = '.idata'
		imgsectheader.VirtualSize = imgimplen
		imgsectheader.VirtualAddress = dirBase - self.xbe.base
		imgsectheader.SizeOfRawData = imgimplen
		imgsectheader.RawData = sectpos
		imgsectheader.Characteristics = READABLE | INITDATA
		fp.write(imgsectheader.pack())
		sectpos += imgimplen
		
		for addr in addrs:
			sect = sectByAddr[addr]
			
			fp.write(sect.data)
		
		fp.write(imgimpdirectory.pack())