import pefile
import sys
import os

# Convenience functions
def rol32(val, amt):
	return ( (val << amt) & 0xffffffff ) | ( ( val >> (32 - amt) ) & 0xffffffff )

def ror32(val, amt):
	return ( (val >> amt) & 0xffffffff ) | ( ( val << (32 - amt) ) & 0xffffffff )

def add32(val, amt):
	return (val + amt) & 0xffffffff

# Typical ror-13 hash
def StandardZ0mbieHash(name):
	edi = 0
	idx = 0
	while(idx < len(name)):
		edi  = add32(ror32(edi, 13), ord(name[idx]) & 0xff)
		idx += 1
			
	return edi
	
# Workhorse:  create IDC script from DLL's export name hashes
def HashExportNames(pe_path, dll_name, idc_path, hashfunc):
	# Open the PE file and create the IDC file
	pe = pefile.PE(pe_path, fast_load=False) 
	f = open(idc_path, 'w')
	
	# Write stock beginning for IDC file
	f.write("#include <idc.idc>\n")
	f.write("static main() {\n")
	f.write("\tauto id;\n")
	f.write("\tid = AddEnum(-1, \"%s_apihashes_t\", 0x1100000);\n" % dll_name)
	
	# Create an enum element for each exported name
	for entry in pe.DIRECTORY_ENTRY_EXPORT.symbols: 
		if entry.name != None:
			f.write("\tAddConstEx(id, \"%s_apihashes_%s\", 0x%lx, -1);\n" % (dll_name, entry.name, hashfunc(entry.name)))
		
	# Close the file
	f.write("}")
	f.close()
	return

# Extract the DLL's name, use it to name the IDC structures uniquely
# and decide the IDC file name, then call the function above
def HashExportNamesWrapper(pe_path, hashfunc = StandardZ0mbieHash):
	base_name = os.path.basename(pe_path)
	dll_name = base_name.split('.')[0]
	idc_path = dll_name + '.idc'
	HashExportNames(pe_path, dll_name, idc_path, hashfunc)

# main():  hash the export names from a DLL; produce [dllname].idc
if __name__ == "__main__":
	HashExportNamesWrapper(sys.argv[1])