#!/usr/bin/env python3
"""Convert a flat ELF64 to a minimal PE32+ UEFI application."""

import struct, sys

def read_elf_sections(path):
    with open(path, "rb") as f:
        data = f.read()
    
    # ELF header
    assert data[:4] == b'\x7fELF', "Not ELF"
    assert data[4] == 2, "Not 64-bit"  # EI_CLASS=2 for 64-bit
    e_phoff = struct.unpack_from('<Q', data, 0x20)[0]
    e_phentsize = struct.unpack_from('<H', data, 0x36)[0]
    e_phnum = struct.unpack_from('<H', data, 0x38)[0]
    
    segments = []
    for i in range(e_phnum):
        off = e_phoff + i * e_phentsize
        p_type = struct.unpack_from('<I', data, off)[0]
        p_flags = struct.unpack_from('<I', data, off + 4)[0]
        p_offset = struct.unpack_from('<Q', data, off + 8)[0]
        p_vaddr = struct.unpack_from('<Q', data, off + 16)[0]
        p_paddr = struct.unpack_from('<Q', data, off + 24)[0]
        p_filesz = struct.unpack_from('<Q', data, off + 32)[0]
        p_memsz = struct.unpack_from('<Q', data, off + 40)[0]
        if p_type == 1:  # PT_LOAD
            segments.append({
                'flags': p_flags,
                'offset': p_offset,
                'vaddr': p_vaddr,
                'paddr': p_paddr,
                'filesz': p_filesz,
                'memsz': p_memsz,
                'data': data[p_offset:p_offset + p_filesz]
            })
    
    return segments, data

def make_pe(segments, elf_data):
    """
    Create PE32+ from ELF64 LOAD segments.
    All segments start at vaddr=0 (from our linker script), so section offsets are relative.
    We map: R segment -> .text, RW segment -> .data, etc.
    """
    # Find the entry point from ELF header
    e_entry = struct.unpack_from('<Q', elf_data, 0x18)[0]
    
    # Sort segments by vaddr
    segments.sort(key=lambda s: s['vaddr'])
    
    # For our linker script (. = 0):
    #   .text at 0x0000 (R+X)  -> PE .text
    #   .rodata at 0x1000 (R)  -> PE .rdata 
    #   .data at 0x2000 (RW)   -> PE .data
    #   .bss at 0x3000 (RW, no file data)
    
    # But we need to look at actual vaddrs to figure out sections
    text_data = b''
    rdata_data = b''
    data_data = b''
    bss_size = 0
    
    for seg in segments:
        if seg['filesz'] == 0:
            # BSS segment
            bss_size = seg['memsz']
            continue
        if seg['flags'] == 5:  # PF_R | PF_X
            text_data = seg['data']
        elif seg['flags'] == 4:  # PF_R
            rdata_data = seg['data']
        elif seg['flags'] == 6:  # PF_R | PF_W
            data_data = seg['data']
    
    # PE constants
    IMAGE_BASE = 0x10000
    SECTION_ALIGN = 0x1000
    FILE_ALIGN = 0x200
    
    # Headers size: DOS(64) + PE_sig(4) + COFF(20) + Optional(240) + section_headers
    num_sections = 3  # .text, .rdata, .data (+ .bss as empty)
    headers_size = 64 + 4 + 20 + 240 + num_sections * 40
    headers_size = (headers_size + FILE_ALIGN - 1) // FILE_ALIGN * FILE_ALIGN  # align to 0x200
    
    # Section RVAs
    text_rva = SECTION_ALIGN
    rdata_rva = text_rva + (len(text_data) + SECTION_ALIGN - 1) // SECTION_ALIGN * SECTION_ALIGN
    data_rva = rdata_rva + (len(rdata_data) + SECTION_ALIGN - 1) // SECTION_ALIGN * SECTION_ALIGN
    bss_rva = data_rva + (len(data_data) + SECTION_ALIGN - 1) // SECTION_ALIGN * SECTION_ALIGN
    
    # Section file offsets
    text_file_off = headers_size
    rdata_file_off = text_file_off + (len(text_data) + FILE_ALIGN - 1) // FILE_ALIGN * FILE_ALIGN
    data_file_off = rdata_file_off + (len(rdata_data) + FILE_ALIGN - 1) // FILE_ALIGN * FILE_ALIGN
    
    size_of_image = bss_rva + bss_size
    size_of_image = (size_of_image + SECTION_ALIGN - 1) // SECTION_ALIGN * SECTION_ALIGN
    
    # Build the image
    pe = bytearray()
    
    # DOS header
    dos = bytearray(64)
    dos[0:2] = b'MZ'
    struct.pack_into('<I', dos, 0x3C, 64)  # PE offset
    pe += dos
    
    # PE signature
    pe += b'PE\x00\x00'
    
    # COFF header
    coff = bytearray(20)
    struct.pack_into('<H', coff, 0, 0x8664)   # Machine: AMD64
    struct.pack_into('<H', coff, 2, num_sections)
    struct.pack_into('<H', coff, 16, 240)      # SizeOfOptionalHeader
    struct.pack_into('<H', coff, 18, 0x22)     # Characteristics: EXECUTABLE_IMAGE | LARGE_ADDRESS_AWARE
    pe += coff
    
    # Optional header (PE32+)
    opt = bytearray(240)
    struct.pack_into('<H', opt, 0, 0x20b)      # Magic: PE32+
    struct.pack_into('<B', opt, 2, 0)           # MajorLinkerVersion
    struct.pack_into('<B', opt, 3, 0)           # MinorLinkerVersion
    struct.pack_into('<I', opt, 4, len(text_data))   # SizeOfCode
    struct.pack_into('<I', opt, 8, len(rdata_data) + len(data_data))  # SizeOfInitializedData
    struct.pack_into('<I', opt, 12, bss_size)   # SizeOfUninitializedData
    struct.pack_into('<I', opt, 16, IMAGE_BASE + text_rva)  # AddressOfEntryPoint
    struct.pack_into('<I', opt, 20, IMAGE_BASE + text_rva)  # BaseOfCode
    struct.pack_into('<Q', opt, 24, IMAGE_BASE)             # ImageBase
    struct.pack_into('<I', opt, 32, SECTION_ALIGN)          # SectionAlignment
    struct.pack_into('<I', opt, 36, FILE_ALIGN)             # FileAlignment
    struct.pack_into('<H', opt, 40, 6)        # MajorOperatingSystemVersion
    struct.pack_into('<H', opt, 42, 0)        # MinorOperatingSystemVersion
    struct.pack_into('<H', opt, 44, 0)        # MajorImageVersion
    struct.pack_into('<H', opt, 46, 0)        # MinorImageVersion
    struct.pack_into('<H', opt, 48, 6)        # MajorSubsystemVersion (UEFI requires >= 6)
    struct.pack_into('<H', opt, 50, 0)        # MinorSubsystemVersion
    struct.pack_into('<I', opt, 52, 0)        # Win32VersionValue
    struct.pack_into('<I', opt, 56, size_of_image)  # SizeOfImage
    struct.pack_into('<I', opt, 60, headers_size)   # SizeOfHeaders
    struct.pack_into('<I', opt, 64, 0)        # CheckSum
    struct.pack_into('<H', opt, 68, 10)       # Subsystem: EFI APPLICATION
    struct.pack_into('<H', opt, 70, 0x8560)   # DllCharacteristics: DYNAMIC_BASE|NX_COMPAT|NO_SEH|TERMINAL_SERVER_AWARE|HIGH_ENTROPY_VA
    struct.pack_into('<Q', opt, 72, 0x100000)  # SizeOfStackReserve
    struct.pack_into('<Q', opt, 80, 0x1000)    # SizeOfStackCommit
    struct.pack_into('<Q', opt, 88, 0x100000)  # SizeOfHeapReserve
    struct.pack_into('<Q', opt, 96, 0x1000)    # SizeOfHeapCommit
    struct.pack_into('<I', opt, 104, 0)       # LoaderFlags
    struct.pack_into('<I', opt, 108, 16)      # NumberOfRvaAndSizes
    pe += opt
    
    # Section headers
    def make_section(name, vsize, vaddr, raw_size, raw_off, chars):
        s = bytearray(40)
        s[0:8] = name.ljust(8, b'\x00')[:8]
        struct.pack_into('<I', s, 8, vsize)
        struct.pack_into('<I', s, 12, vaddr)
        struct.pack_into('<I', s, 16, raw_size)
        struct.pack_into('<I', s, 20, raw_off)
        struct.pack_into('<I', s, 24, 0)  # PointerToRelocations
        struct.pack_into('<I', s, 28, 0)  # PointerToLinenumbers
        struct.pack_into('<H', s, 32, 0)  # NumberOfRelocations
        struct.pack_into('<H', s, 34, 0)  # NumberOfLinenumbers
        struct.pack_into('<I', s, 36, chars)
        return s
    
    TEXT_CHARS  = 0x60000020  # CNT_CODE | MEM_READ | MEM_EXECUTE
    RDATA_CHARS = 0x40000040  # CNT_INITIALIZED_DATA | MEM_READ
    DATA_CHARS  = 0xC0000040  # CNT_INITIALIZED_DATA | MEM_READ | MEM_WRITE
    
    pe += make_section(b'.text', len(text_data), text_rva, len(text_data), text_file_off, TEXT_CHARS)
    pe += make_section(b'.rdata', len(rdata_data), rdata_rva, len(rdata_data), rdata_file_off, RDATA_CHARS)
    pe += make_section(b'.data', len(data_data), data_rva, len(data_data), data_file_off, DATA_CHARS)
    
    # Pad to headers_size
    while len(pe) < headers_size:
        pe += b'\x00'
    
    # Section data
    # .text
    while len(pe) < text_file_off:
        pe += b'\x00'
    pe += text_data
    while len(pe) < text_file_off + len(text_data):
        pe += b'\x00'
    
    # .rdata
    while len(pe) < rdata_file_off:
        pe += b'\x00'
    pe += rdata_data
    while len(pe) < rdata_file_off + len(rdata_data):
        pe += b'\x00'
    
    # .data
    while len(pe) < data_file_off:
        pe += b'\x00'
    pe += data_data
    while len(pe) < data_file_off + len(data_data):
        pe += b'\x00'
    
    return bytes(pe)

segments, elf_data = read_elf_sections("bootx64.so")
pe_data = make_pe(segments, elf_data)
with open("bootx64.efi", "wb") as f:
    f.write(pe_data)
print(f"PE32+ created: {len(pe_data)} bytes, {len(segments)} segments")

# Verify entry point
e_entry = struct.unpack_from('<Q', elf_data, 0x18)[0]
print(f"ELF entry: {hex(e_entry)}")
