drgn/tests/elfwriter.py

import struct
from typing import Optional, Sequence

from tests.elf import ET, PT, SHT


class ElfSection:
    def __init__(self, data: bytes,
                 name: Optional[str] = None,
                 sh_type: Optional[SHT] = None,
                 p_type: Optional[PT] = None,
                 vaddr: int = 0,
                 paddr: int = 0,
                 memsz: Optional[int] = None,
                 p_align: int = 0):
        self.data = data
        self.name = name
        self.sh_type = sh_type
        self.p_type = p_type
        self.vaddr = vaddr
        self.paddr = paddr
        self.memsz = memsz
        self.p_align = p_align

        assert (self.name is not None) or (self.p_type is not None)
        assert (self.name is None) == (self.sh_type is None)
        if self.p_type is None:
            assert self.memsz is None
        elif self.memsz is None:
            self.memsz = len(self.data)


def create_elf_file(type: ET, sections: Sequence[ElfSection],
                    little_endian: bool = True, bits: int = 64):
    endian = '<' if little_endian else '>'
    if bits == 64:
        ehdr_struct = struct.Struct(endian + '16BHHIQQQIHHHHHH')
        shdr_struct = struct.Struct(endian + 'IIQQQQIIQQ')
        phdr_struct = struct.Struct(endian + 'IIQQQQQQ')
        e_machine = 62 if little_endian else 43  # EM_X86_64 or EM_SPARCV9
    else:
        assert bits == 32
        ehdr_struct = struct.Struct(endian + '16BHHIIIIIHHHHHH')
        shdr_struct = struct.Struct(endian + '10I')
        phdr_struct = struct.Struct(endian + '8I')
        e_machine = 3 if little_endian else 8  # EM_386 or EM_MIPS

    shstrtab = ElfSection(
        name='.shstrtab',
        sh_type=SHT.STRTAB,
        data=bytearray(1),
    )
    tmp = [shstrtab]
    tmp.extend(sections)
    sections = tmp
    shnum = 1  # One for the SHT_NULL section.
    phnum = 0
    for section in sections:
        if section.name is not None:
            shstrtab.data.extend(section.name.encode())
            shstrtab.data.append(0)
            shnum += 1
        if section.p_type is not None:
            phnum += 1

    shdr_offset = ehdr_struct.size
    phdr_offset = shdr_offset + shdr_struct.size * shnum
    headers_size = phdr_offset + phdr_struct.size * phnum
    buf = bytearray(headers_size)
    ehdr_struct.pack_into(
        buf, 0,
        0x7f,  # ELFMAG0
        ord('E'),  # ELFMAG1
        ord('L'),  # ELFMAG2
        ord('F'),  # ELFMAG3
        2 if bits == 64 else 1,  # EI_CLASS = ELFCLASS64 or ELFCLASS32
        1 if little_endian else 2,  # EI_DATA = ELFDATA2LSB or ELFDATA2MSB
        1,  # EI_VERSION = EV_CURRENT
        0,  # EI_OSABI = ELFOSABI_NONE
        0,  # EI_ABIVERSION
        0, 0, 0, 0, 0, 0, 0,  # EI_PAD
        type,  # e_type
        e_machine,
        1,  # e_version = EV_CURRENT
        0,  # e_entry
        phdr_offset,  # e_phoff
        shdr_offset,  # e_shoff
        0,  # e_flags
        ehdr_struct.size,  # e_ehsize
        phdr_struct.size,  # e_phentsize
        phnum,  # e_phnum
        shdr_struct.size,  # e_shentsize,
        shnum,  # e_shnum,
        1,  # e_shstrndx
    )

    shdr_offset += shdr_struct.size
    for section in sections:
        if section.p_align:
            padding = (section.vaddr % section.p_align -
                       len(buf) % section.p_align)
            buf.extend(bytes(padding))
        if section.name is not None:
            shdr_struct.pack_into(
                buf, shdr_offset,
                shstrtab.data.index(section.name.encode()),  # sh_name
                section.sh_type,  # sh_type
                0,  # sh_flags
                section.vaddr,  # sh_addr
                len(buf),  # sh_offset
                len(section.data),  # sh_size
                0,  # sh_link
                0,  # sh_info
                1 if section.p_type is None else bits // 8,  # sh_addralign
                0,  # sh_entsize
            )
            shdr_offset += shdr_struct.size
        if section.p_type is not None:
            flags = 7   # PF_R | PF_W | PF_X
            if bits == 64:
                phdr_struct.pack_into(
                    buf, phdr_offset,
                    section.p_type,  # p_type
                    flags,  # p_flags
                    len(buf),  # p_offset
                    section.vaddr,  # p_vaddr
                    section.paddr,  # p_paddr
                    len(section.data),  # p_filesz
                    section.memsz,  # p_memsz
                    section.p_align,  # p_align
                )
            else:
                phdr_struct.pack_into(
                    buf, phdr_offset,
                    section.p_type,  # p_type
                    len(buf),  # p_offset
                    section.vaddr,  # p_vaddr
                    section.paddr,  # p_paddr
                    len(section.data),  # p_filesz
                    section.memsz,  # p_memsz
                    flags,  # p_flags
                    section.p_align,  # p_align
                )
            phdr_offset += phdr_struct.size
        buf.extend(section.data)

    return buf