import sys 
import pefile
from capstone import *

def disassemble(file_path):
    # pefile info: https://github.com/erocarrera/pefile/blob/wiki/UsageExamples.md
    # load pe file or  pefile.PE(use data=str_object_with_pe_file_data) to load string buffer
    pe = pefile.PE(file_path)
    imports = {}
    # get all imports  
    for entry in pe.DIRECTORY_ENTRY_IMPORT:
        # uncomment to print dll name 
        # print entry.dll
        for imp in entry.imports:
            imports[hex(imp.address)] = imp.name
    # get entry point of PE
    eop = pe.OPTIONAL_HEADER.AddressOfEntryPoint
    code_section = pe.get_section_by_rva(eop)
    # get code data 
    code_dump = code_section.get_data()
    # allocate correct memory address 
    code_addr = pe.OPTIONAL_HEADER.ImageBase + code_section.VirtualAddress
    # Capstone engine info http://www.capstone-engine.org/lang_python.html
    # Initialize capstone class, arguments for class hardware architecture & the hardware mode
    # for 64bit change CS_MODE_32 to CS_MODE_64
    md = Cs(CS_ARCH_X86, CS_MODE_32)
    # detail must be enabled to get the operand values/details 
    md.detail = True
    # disassemble memory in code_dump at address code_addr
    for i in md.disasm(code_dump, code_addr):
        #  dir(i) for more details 
        op_str = i.op_str
        # compare mnemonic
        if i.mnemonic == "call" or i.mnemonic == "jmp" :
            # get the displacement within the operand
            temp_str = hex(i.operands[0].mem.disp)
            # hack to remove the L from Pythons Long
            if temp_str[-1] == "L":
                temp_str = temp_str[:-1]
            if imports.has_key(temp_str):
                # replace address with import name
                api_name = imports[temp_str]
                op_str = api_name
        # print address, mnemonic and operand string. See line 36 on accessing more 
        # details in the operand(s) 
        print("0x%x:\t%s\t%s" %(i.address, i.mnemonic, op_str))

disassemble(sys.argv[1])