Call scan -- quickly scan for call instructions in text section of PE

scan.py

import pefile
import re
import sys

def extract_code_section(pe):
    code_section = None
    for section in pe.sections:
        if section.Name.startswith(b'.text'):
            code_section = section
            break
    if not code_section:
        raise ValueError("Code section (.text) not found in the PE file.")
    
    code_data = code_section.get_data()
    code_base = code_section.VirtualAddress
    return code_data, code_base

def scan_for_calls(code_data, code_base):
    call_pattern = re.compile(rb'\xE8[\x00-\xFF]{4}')
    calls = call_pattern.finditer(code_data)
    
    call_addresses = []
    for call in calls:
        call_offset = call.start()
        relative_address = int.from_bytes(code_data[call_offset+1:call_offset+5], byteorder='little', signed=True)
        call_address = code_base + call_offset + 5 + relative_address
        call_addresses.append((code_base + call_offset, call_address))
    
    return call_addresses

def main(pe_path):
    calls = []
    pe = pefile.PE(pe_path)
    image_base = pe.OPTIONAL_HEADER.ImageBase
    code_data, code_base = extract_code_section(pe)
    call_addresses = scan_for_calls(code_data, code_base)
    
    for offset, address in call_addresses:
        if code_base <= address < (code_base + len(code_data)):
            calls.append((offset + image_base, address + image_base))

    # Count number of calls to each address
    call_count = {}
    for offset, address in calls:
        if address not in call_count:
            call_count[address] = 0
        call_count[address] += 1
    
    # Print the to 10 most called addresses
    print("Top 10 most called addresses:")
    for address, count in sorted(call_count.items(), key=lambda x: x[1], reverse=True)[:10]:
        print(f"Address: 0x{address:08X}, Count: {count}")

    print("\n\nAll calls found:")
    print(f"Found {len(calls)} calls in the code section of the PE file.")
    for offset, address in calls:
        print(f"Call: 0x{offset:08X}, Address: 0x{address:08X}")


pe_path = sys.argv[1]
main(pe_path)