diagnose.py

                      

import sys
import pefile
import struct
from capstone import Cs, CS_ARCH_X86, CS_MODE_64

def diagnose_pe(filepath):
    print(f"\n{'='*60}")
    print(f"Diagnosing file: {filepath}")
    print(f"{'='*60}\n")
    
    try:
        pe = pefile.PE(filepath)
    except Exception as e:
        print(f"❌ Failed to load PE file: {e}")
        return False
    
    issues = []
    warnings = []
    
               
    print("[1] Check basic PE structure...")
    try:
        if not pe.OPTIONAL_HEADER:
            issues.append("Missing OPTIONAL_HEADER")
        else:
            print(f"  ✓ ImageBase: 0x{pe.OPTIONAL_HEADER.ImageBase:X}")
            print(f"  ✓ EntryPoint: 0x{pe.OPTIONAL_HEADER.AddressOfEntryPoint:X}")
            print(f"  ✓ SizeOfImage: 0x{pe.OPTIONAL_HEADER.SizeOfImage:X}")
    except Exception as e:
        issues.append(f"OPTIONAL_HEADER check failed: {e}")
    
                 
    print("\n[2] Check .text section...")
    text_section = None
    for section in pe.sections:
        name = section.Name.decode('utf-8', errors='ignore').rstrip('\x00')
        if name == '.text':
            text_section = section
            print("  ✓ .text found")
            print(f"    - VirtualAddress: 0x{section.VirtualAddress:X}")
            print(f"    - VirtualSize: 0x{section.Misc_VirtualSize:X}")
            print(f"    - SizeOfRawData: 0x{section.SizeOfRawData:X}")
            print(f"    - PointerToRawData: 0x{section.PointerToRawData:X}")
            
                          
            if section.SizeOfRawData < section.Misc_VirtualSize:
                warnings.append(
                    f".text has {section.Misc_VirtualSize - section.SizeOfRawData} bytes of uninitialized slack"
                )
            break
    
    if not text_section:
        issues.append("Missing .text section")
        return False
    
              
    print("\n[3] Check entry point...")
    entry_rva = pe.OPTIONAL_HEADER.AddressOfEntryPoint
    entry_offset = pe.get_offset_from_rva(entry_rva)
    
    if entry_offset:
                     
        pe_data = pe.__data__
        entry_bytes = pe_data[entry_offset:entry_offset+16]
        print(f"  ✓ EntryPoint RVA: 0x{entry_rva:X}")
        print(f"    File offset: 0x{entry_offset:X}")
        print(f"    First 16 bytes: {' '.join(f'{b:02X}' for b in entry_bytes)}")
        
                   
        try:
            md = Cs(CS_ARCH_X86, CS_MODE_64)
            md.detail = True
            disasm = list(md.disasm(bytes(entry_bytes), entry_rva + pe.OPTIONAL_HEADER.ImageBase))
            if disasm:
                print("    Disassembly:")
                for insn in disasm[:5]:
                    print(f"      {insn.mnemonic} {insn.op_str}")
            else:
                issues.append("Failed to disassemble entry point")
        except Exception as e:
            warnings.append(f"Entry point disassembly failed: {e}")
    else:
        issues.append(f"Failed to resolve entry point file offset (RVA=0x{entry_rva:X})")
    
              
    print("\n[4] Check import table...")
    if hasattr(pe, 'DIRECTORY_ENTRY_IMPORT'):
        import_count = len(pe.DIRECTORY_ENTRY_IMPORT)
        print(f"  ✓ Found {import_count} imported DLL(s)")
        for entry in pe.DIRECTORY_ENTRY_IMPORT[:5]:
            dll_name = entry.dll.decode('utf-8', errors='ignore')
            func_count = len(entry.imports)
            print(f"    - {dll_name}: {func_count} function(s)")
    else:
        warnings.append("No import table found")
    
               
    print("\n[5] Check base relocations...")
    if hasattr(pe, 'DIRECTORY_ENTRY_BASERELOC'):
        reloc_count = sum(len(entry.entries) for entry in pe.DIRECTORY_ENTRY_BASERELOC)
        print(f"  ✓ Found {reloc_count} relocation item(s)")
    else:
        warnings.append("No base relocation table found")
    
              
    print("\n[6] Check exception table...")
    if hasattr(pe, 'DIRECTORY_ENTRY_EXCEPTION'):
        exc_count = len(pe.DIRECTORY_ENTRY_EXCEPTION)
        print(f"  ✓ Found {exc_count} exception handler(s)")
    else:
        warnings.append("No exception table found")
    
                          
    print("\n[7] Scan for potential jump tables...")
    text_data = text_section.get_data()
    text_rva = text_section.VirtualAddress
    image_base = pe.OPTIONAL_HEADER.ImageBase
    
    md = Cs(CS_ARCH_X86, CS_MODE_64)
    md.detail = True
    
    jump_table_patterns = []
    try:
        insns = list(md.disasm(text_data, text_rva + image_base))
        for i in range(len(insns) - 2):
                                      
            if (insns[i].mnemonic == 'movsxd' and 
                insns[i+1].mnemonic == 'add' and 
                insns[i+2].mnemonic == 'jmp'):
                
                       
                if len(insns[i+2].operands) > 0 and insns[i+2].operands[0].type == 1:       
                    rva = insns[i].address - image_base
                    jump_table_patterns.append(rva)
                    print(f"  → Found jump-table pattern at RVA 0x{rva:X}")
    except Exception as e:
        warnings.append(f"Jump-table scan failed: {e}")
    
    if jump_table_patterns:
        print(f"  Found {len(jump_table_patterns)} potential jump table(s)")
    else:
        print("  No obvious jump-table pattern found")
    
                         
    print("\n[8] Check .rdata section...")
    rdata_section = None
    for section in pe.sections:
        name = section.Name.decode('utf-8', errors='ignore').rstrip('\x00')
        if name == '.rdata':
            rdata_section = section
            print("  ✓ .rdata found")
            print(f"    - VirtualAddress: 0x{section.VirtualAddress:X}")
            print(f"    - VirtualSize: 0x{section.Misc_VirtualSize:X}")
            break
    
    if rdata_section:
                    
        known_tables = [0xB804, 0xB930, 0xB980]           
        rdata_data = rdata_section.get_data()
        rdata_rva = rdata_section.VirtualAddress
        
        for table_rva in known_tables:
            if rdata_rva <= table_rva < rdata_rva + len(rdata_data):
                offset = table_rva - rdata_rva
                         
                print(f"\n  Check jump table at RVA 0x{table_rva:X}:")
                for j in range(min(4, (len(rdata_data) - offset) // 4)):
                    entry_offset = offset + j * 4
                    if entry_offset + 4 <= len(rdata_data):
                        rel = struct.unpack_from('<i', rdata_data, entry_offset)[0]
                        target = table_rva + rel
                                       
                        in_text = (text_rva <= target < text_rva + text_section.Misc_VirtualSize)
                        status = "✓" if in_text else "⚠"
                        print(f"    Entry {j}: offset=0x{rel:08X} -> target=0x{target:X} {status}")
    
        
    print(f"\n{'='*60}")
    print("Summary:")
    print(f"{'='*60}")
    
    if issues:
        print(f"\n❌ Found {len(issues)} issue(s):")
        for issue in issues:
            print(f"  - {issue}")
    
    if warnings:
        print(f"\n⚠ Found {len(warnings)} warning(s):")
        for warning in warnings:
            print(f"  - {warning}")
    
    if not issues and not warnings:
        print("\n✓ No obvious issues found")
    
    return len(issues) == 0

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("Usage: python3 diagnose.py <pe_file>")
        sys.exit(1)
    
    filepath = sys.argv[1]
    success = diagnose_pe(filepath)
    sys.exit(0 if success else 1)