#!/usr/pkg/bin/python
"""pe.py - routines to create minimalistic Windows PE executables

   emphasis on symbol tables and hashtables for building postFORTH"""

Copyright = """
    pe.py - routines to create minimalistic ELF files
    Copyright (C) 2005  John Comeau <jc.jcomeau.com>

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
    """
errormessage = "Not all needed libraries found, upgrade or check path: "
try:
 True # not defined in older Python releases
except:
 True, False = 1, 0
try:
 import sys, os, types, re, pwd
 sys.path.append(os.path.join(pwd.getpwuid(os.geteuid())[5], 'lib', 'python'))
 errormessage = errormessage + repr(sys.path)
 from com.jcomeau import gpl, jclicense
except:
 try:
  sys.stderr.write("%s\n" % errormessage)
 except:
  print errormessage
 raise

# get name this program was called as
myself = os.path.split(sys.argv[0])[1]
command = os.path.splitext(myself)[0]  # chop any suffix (extension)
# now get name we gave it when we wrote it
originalself = re.compile('[0-9A-Za-z]+').search(Copyright).group()

# globals and routines that should be in every program
# (yes, you could import them, but there are problems in that approach too)

def DebugPrint(*whatever):
 return False  # defined instead by pytest module, use that for debugging

def join(*args):
 "for pythons without str.join"
 string, array = args
 if type(array) == types.StringType: array = eval(array)
 if hasattr(str, 'join'):
  return string.join(array)
 else:
  joined = ''
  for index in range(0, len(array)):
   joined = joined + array[index]
   if index != (len(array) - 1):
    joined = joined + string
  return joined

def split(*args):
 "for pythons without str.split"
 string, string_to_split = args
 if not len(string): string = None
 if hasattr('str', 'split'):
  return string_to_split.split(string)
 else:
  return re.compile(re.escape(string)).split(string_to_split)

# other globals, specific to this program
import struct
IMAGE_FILE_RELOCS_STRIPPED = 1
IMAGE_FILE_EXECUTABLE_IMAGE = 2
IMAGE_FILE_32BIT_MACHINE = 0x100
IMAGE_FILE_DEBUG_STRIPPED = 0x200

def null_dos_stub(*args):
 """create a non-runnable DOS stub

    required by spec, doesn't really need to run on DOS"""
 length = 64  # minimum? not sure
 e_lfanew = struct.pack('<L', length)  # pointer to PE header
 header = 'MZ'  # the magic 'number' for a DOS executable
 skiplength = len(header) + len(e_lfanew)
 header += '\0' * (length - skiplength)  # pad with zeroes
 header += e_lfanew  # pointer to PE header
 return header

def pe_header(*args):
 """24-byte PE header"""
 i386 = 0x14c  # don't ask me, found it online
 header = 'PE\0\0'
 header += struct.pack('<H', i386)  # Machine it should run on
 header += struct.pack('<H', 2)  # NumberOfSections, we use code and data only
 header += struct.pack('<L', 0)  # TimeDateStamp not necessary
 header += struct.pack('<L', 0)  # PointerToSymbolTable unused
 header += struct.pack('<L', 0)  # NumberOfSymbols unused
 header += struct.pack('<H', 0xe0)  # SizeOfOptionalHeader (includes datadirs)
 header += struct.pack('<H', IMAGE_FILE_RELOCS_STRIPPED | \
  IMAGE_FILE_EXECUTABLE_IMAGE | \
  IMAGE_FILE_32BIT_MACHINE | \
  IMAGE_FILE_DEBUG_STRIPPED)
 return header

def optional_header(*args):
 """not optional for PE files, the name may be from COFF

    this header specifies some vital runtime parameters of this image"""
 magic = 0x10b  # PE32 = 0x10b; PE32+ = 0x20b
 header = struct.pack('<H', magic)  # Magic
 header += struct.pack('BB', 0, 0)  # MajorLinkerVersion, MinorLinkerVersion
 header += struct.pack('<L', 0x20)  # SizeOfCode
 header += struct.pack('<L', 0xa0)  # SizeOfInitializedData
 header += struct.pack('<L', 0)  # SizeOfUninitializedData, AKA .bss
 header += struct.pack('<L', 0x1a0)  # AddressOfEntryPoint
 header += struct.pack('<L', 0x1a0)  # BaseOfCode
 header += struct.pack('<L', 0x1c0)  # BaseOfData
 header += struct.pack('<L', 0x100000)  # ImageBase (arbitrarily chosen)
 header += struct.pack('<LL', 32, 32)  # SectionAlignment, FileAlignment
 # prefix the next 3 with Major and Minor, respectively
 header += struct.pack('<HH', 4, 0)  # OperatingSystemVersion (NT4)
 header += struct.pack('<HH', 0, 0)  # ImageVersion
 header += struct.pack('<HH', 4, 0)  # SubsystemVersion
 header += struct.pack('<L', 0)  # Win32VersionValue
 header += struct.pack('<L', 0xc0)  # SizeOfImage, sum of all section sizes
 header += struct.pack('<L', 0x1a0)  # SizeOfHeaders, offset to first section
 header += struct.pack('<L', 0)  # CheckSum, only needed for drivers
 header += struct.pack('<H', 3)  # SubSystem, Win32 console
 header += struct.pack('<H', 0)  # DllCharacteristics, unused for executables
 header += struct.pack('<LLLL', 0x100000, 0x1000,  # stack
  0x100000, 0x1000)  # heap
 header += struct.pack('<L', 0)  # LoaderFlags
 header += struct.pack('<L', 0x10)  # NumberOfRvaAndSizes
 return header

def data_directories(*args):
 """generate default values for the data directories

    all have names in the format IMAGE_DIRECTORY_ENTRY_x,
    where x has the names, in order:
    EXPORT, IMPORT, RESOURCE, EXCEPTION, SECURITY, BASERELOC, DEBUG, COPYRIGHT,
    GLOBALPTR, TLS, LOAD_CONFIG, BOUND_IMPORT, IAT. offsets 13-15 are unused.
 """
 dirs = struct.pack('<LL', 0, 0)  # EXPORT
 dirs += struct.pack('<LL', 0x1e0, 0x6f)  # IMPORT: offset, size
 dirs += struct.pack('<LL', 0, 0) * 14  # the rest unused
 return dirs

def nullpad(scalar, alignment):
 if alignment:  # no-op with alignment=0
  modulus = len(scalar) % alignment
  if modulus:
   scalar += '\0' * (alignment - modulus)
 return scalar

def code_header(*args):
 """create a sectionheader for the code section"""
 header = nullpad(".text", 8)
 header += struct.pack('<L', 0)  # VirtualSize, unused
 header += struct.pack('<L', 0x1a0)  # VirtualAddress, RVA to code
 header += struct.pack('<L', 0x20)  # SizeOfRawData
 header += struct.pack('<L', 0x1a0)  # PointerToRawData, offset to code
 header += struct.pack('<LLHH', 0, 0, 0, 0)  # unused
 header += struct.pack('<L', 0x60000020)  # Characteristics: code, exec, read
 return header

def data_header(*args):
 """create a sectionheader for the data section"""
 header = nullpad(".data", 8)
 header += struct.pack('<L', 0)  # VirtualSize, unused
 header += struct.pack('<L', 0x1c0)  # VirtualAddress, RVA to code
 header += struct.pack('<L', 0xa0)  # SizeOfRawData
 header += struct.pack('<L', 0x1c0)  # PointerToRawData, offset to code
 header += struct.pack('<LLHH', 0, 0, 0, 0)  # unused
 header += struct.pack('<L', 0xc0000040)  # Characteristics: init, read, write
 return header

def code_section(*args):
 code = '\x6a\x00'  # push 0
 code += '\x68\xd0\x01\x10\x00'  # push offset written
 code += '\x6a\x0d'  # push 0xd
 code += '\x68\xc0\x01\x10\x00'  # push offset hello
 code += '\x6a\xf5'  # push 0xfffffff5
 code += '\x2e\xff\x15\x28\x02\x10\x00'  # call dword ptr GetStdHandle
 code += '\x50'  # push eax
 code += '\x2e\xff\x15\x24\x02\x10\x00'  # call dword ptr WriteConsoleA
 code += '\xc3'  # ret
 return code

def data_section(*args):
 data = nullpad('hello, world\n', 16)
 data += struct.pack('<L', 0)  # variable 'written'
 return data

def import_descriptors(*args):
 descriptors = struct.pack('<L', 0x218)  # OriginalFirstThunk
 descriptors += struct.pack('<L', 0)  # TimeDateStamp
 descriptors += struct.pack('<L', -1)  # ForwarderChain, no forwarders
 descriptors += struct.pack('<L', 0x208)  # Name, RVA to DLL name
 descriptors += struct.pack('<L', 0x224)  # FirstThunk, RVA to 1st thunk
 descriptors += struct.pack('<LLLLL', 0, 0, 0, 0, 0)  # terminates list
 descriptors += nullpad('kernel32.dll', 16)  # name of DLL
 descriptors += struct.pack('<L', 0x230)  # AddressOfData, RVA to function name
 descriptors += struct.pack('<L', 0x240)  # RVA to function name GetStdHandle
 descriptors += struct.pack('<L', 0)  # terminator of OriginalFirstThunk
 descriptors += struct.pack('<L', 0x230)  # AddressOfData, RVA to function name
 descriptors += struct.pack('<L', 0x240)  # RVA to function name GetStdHandle
 descriptors += struct.pack('<L', 0)  # terminator of FirstThunk
 descriptors += struct.pack('<H', 1) + 'WriteConsoleA\0'
 descriptors += struct.pack('<H', 2) + 'GetStdHandle\0'
 return descriptors

def pe(*args):
 """create a default PE file, output to stdout or file

    the PE file consists of:
     DOS stub
     PE header
     "optional" header (actually vital)
     data directories
     section headers
     sections
 """
 pe_file = null_dos_stub()
 pe_file += pe_header()
 pe_file += optional_header()
 pe_file += data_directories()
 pe_file += code_header()
 pe_file += data_header()
 pe_file = nullpad(pe_file, 0x1a0)
 pe_file += code_section()
 pe_file += data_section()
 pe_file = nullpad(pe_file, 0x10)
 pe_file += import_descriptors()
 pe_file = nullpad(pe_file, 0x20)
 DebugPrint('pe_file', pe_file)
 output = sys.stdout
 if len(args):
  output = open(args[0], 'wb')
 output.write(pe_file)

if __name__ == '__main__':
 # if this program was imported by another, the above test will fail,
 # and this following code won't be used...
 function = command; args = sys.argv[1:]  # default case
 if command == originalself:
  try:
   if len(args) and eval('type(%s) == types.FunctionType' % args[0]):
    function = sys.argv[1]; args = sys.argv[2:]
  except: pass
 print eval('%s%s' % (function, repr(tuple(args)))) or ''
else:
 # if you want something to be done on import, do it here; otherwise pass
 pass