#!/usr/bin/env python3
# readcoff -- list or dump PE/COFF file sections (such as .efi or .exe files)
import argparse
from enum import IntEnum
from nullroute.io import BinaryReader

class CoffImageFileType(IntEnum):
    Unknown = 0x0
    I386 = 0x14c
    Alpha = 0x184
    Arm = 0x1c0
    Itanium = 0x200
    Alpha64 = 0x284
    Motorola68k = 0x268
    Mips16 = 0x266
    MipsFPU = 0x366
    MipsFPU16 = 0x466
    PowerPC = 0x1f0
    R3000 = 0x162
    R4000 = 0x166
    R10000 = 0x168
    Sh3 = 0x1a2
    Sh4 = 0x1a6
    Thumb = 0x1c2

def br_enum_sections(br):
    br.seek(0)
    # MS-DOS stub and PE signature
    dos_stub = br.read(0x3c)
    if dos_stub[0:2] != b'MZ':
        raise ValueError("File does not start with MS-DOS MZ magic")
    pe_offset = br.read_u16_le()
    br.seek(pe_offset)
    pe_sig = br.read(4)
    if pe_sig != b'PE\0\0':
        raise ValueError("File does not contain PE signature")
    # COFF header
    target_machine = br.read_u16_le()
    num_sections = br.read_u16_le()
    time_date = br.read_u32_le()
    symtab_offset = br.read_u32_le()
    num_symbols = br.read_u32_le()
    opthdr_size = br.read_u16_le()
    characteristics = br.read_u16_le()
    #target_machine = CoffImageFileType(target_machine)
    #print(target_machine)
    # Optional PE32 Header
    _ = br.read(opthdr_size)
    # Section Table
    for i in range(num_sections):
        section_name = br.read(8).rstrip(b"\0")
        virtual_size = br.read_u32_le()
        virtual_addr = br.read_u32_le()
        section_size = br.read_u32_le()
        section_offset = br.read_u32_le()
        relocs_offset = br.read_u32_le()
        linenums_offset = br.read_u32_le()
        num_relocs = br.read_u16_le()
        num_linenums = br.read_u16_le()
        characteristics = br.read_u32_le()
        # Section names longer than 8 bytes are stored as pointers into the
        # symbol table. Note that num_symbols will be 0 as COFF symbols are
        # generally unused for other purposes.
        if symtab_offset and section_name.startswith(b"/"):
            sym_offset = int(section_name[1:])
            prev_pos = br.tell()
            br.seek(symtab_offset + sym_offset)
            # The symbol table is optimized for memory-mapped usage, so it just
            # consists of null-terminated strings.
            section_name = br.read(32).split(b"\0")[0]
            br.seek(prev_pos)
        yield (section_name,
               section_offset,
               min(section_size, virtual_size))

def read_coff_section(file, wanted_section):
    with open(file, "rb") as fh:
        br = BinaryReader(fh)
        sections = {}
        for name, offset, size in br_enum_sections(br):
            name = name.decode()
            sections[name] = (offset, size)
            if not wanted_section:
                print(name)
        if wanted_section:
            if wanted_section in sections:
                offset, size = sections[wanted_section]
                br.seek(offset)
                return br.read(size)
            else:
                raise KeyError("section %r not found" % wanted_section)

parser = argparse.ArgumentParser()
parser.add_argument("-o", "--output", metavar="PATH",
                    help="write section to the specified file if found")
parser.add_argument("path")
parser.add_argument("section", nargs="?")
args = parser.parse_args()

infile = args.path
if args.section:
    section = args.section
    outfile = "/dev/stdout"
    if args.output:
        outfile = args.output
    data = read_coff_section(infile, section)
    with open(outfile, "wb") as fh:
        fh.write(data)
else:
    read_coff_section(infile, None)
