#!/usr/bin/env python3
"""
pefparse.py -- Minimal parser/disassembler for classic Mac OS CFM
Preferred Executable Format (PEF) PowerPC containers ('Joy!peff pwpc').

Written for the VWE Tesla restoration project to decompile "Console 4.10".
Big-endian PowerPC. Pure-python; only external dep is capstone (optional,
for the `dis` command).

PEF reference: Apple "Mac OS Runtime Architectures", ch. Preferred
Executable Format. Container header (40 bytes), then N section headers
(28 bytes each), then the section name table, then section bodies.

Usage:
    py -3.13 pefparse.py <file> header      # container + section headers
    py -3.13 pefparse.py <file> loader      # loader section: imports/exports/entry
    py -3.13 pefparse.py <file> strings     # printable strings in code+data
    py -3.13 pefparse.py <file> dis [secidx] [start] [count]  # disassemble
    py -3.13 pefparse.py <file> dump <secidx> <off> <len>     # hex dump section body
"""
import sys, struct, string

def u32(b, o): return struct.unpack_from('>I', b, o)[0]
def u16(b, o): return struct.unpack_from('>H', b, o)[0]
def s32(b, o): return struct.unpack_from('>i', b, o)[0]
def u8(b, o):  return b[o]

SECTION_KINDS = {
    0: 'Code', 1: 'UnpackedData', 2: 'PatternInitData', 3: 'Constant',
    4: 'Loader', 5: 'Debug', 6: 'ExecutableData', 7: 'Exception', 8: 'Traceback',
}

class Section:
    def __init__(self, idx, hdr, container):
        (self.name_off, self.default_addr, self.total_size, self.unpacked_size,
         self.packed_size, self.container_off, self.kind, self.share,
         self.align, self.rsvd) = hdr
        self.idx = idx
        self.name = None
        self._container = container

    def raw(self):
        """Raw (possibly packed) bytes of the section body as stored."""
        return self._container[self.container_off:self.container_off + self.packed_size]

    def data(self):
        """Unpacked section body. PatternInitData(2) is RLE-decompressed."""
        raw = self.raw()
        if self.kind == 2:
            return unpack_pidata(raw, self.unpacked_size)
        return raw

class PEF:
    def __init__(self, blob):
        self.blob = blob
        assert blob[0:4] == b'Joy!', "not a PEF (missing Joy!)"
        assert blob[4:8] == b'peff', "not a PEF (missing peff)"
        self.arch = blob[8:12]
        self.format_version = u32(blob, 12)
        self.timestamp = u32(blob, 16)
        self.old_def_ver = u32(blob, 20)
        self.old_imp_ver = u32(blob, 24)
        self.current_ver = u32(blob, 28)
        self.section_count = u16(blob, 32)
        self.inst_section_count = u16(blob, 34)
        self.sections = []
        o = 40
        for i in range(self.section_count):
            hdr = struct.unpack_from('>iIIIIIBBBB', blob, o)
            self.sections.append(Section(i, hdr, blob))
            o += 28
        # section name table follows the section headers
        name_tbl_off = o
        for s in self.sections:
            if s.name_off >= 0:
                s.name = read_cstr(blob, name_tbl_off + s.name_off)

    def loader(self):
        for s in self.sections:
            if s.kind == 4:
                return s
        return None

def read_cstr(b, o):
    e = o
    while e < len(b) and b[e] != 0:
        e += 1
    return b[o:e].decode('mac_roman', 'replace')

def unpack_pidata(raw, unpacked_size):
    """Decompress PEF pattern-initialized data (RLE opcode stream)."""
    out = bytearray()
    i = 0
    n = len(raw)
    def argcount():
        nonlocal i
        val = 0
        while True:
            b = raw[i]; i += 1
            val = (val << 7) | (b & 0x7f)
            if not (b & 0x80):
                break
        return val
    while i < n:
        op = raw[i]; i += 1
        opcode = op >> 5
        count5 = op & 0x1f
        cnt = count5 if count5 else argcount()
        if opcode == 0:  # Zero
            out.extend(b'\x00' * cnt)
        elif opcode == 1:  # blockCopy
            out.extend(raw[i:i+cnt]); i += cnt
        elif opcode == 2:  # repeatedBlock
            rpt = argcount()
            block = raw[i:i+cnt]; i += cnt
            for _ in range(rpt + 1):
                out.extend(block)
        elif opcode == 3:  # interleaveRepeatBlockWithBlockCopy
            cs = argcount()      # custom size
            rc = argcount()      # repeat count
            common = raw[i:i+cnt]; i += cnt
            for _ in range(rc):
                out.extend(common)
                out.extend(raw[i:i+cs]); i += cs
            out.extend(common)
        elif opcode == 4:  # interleaveRepeatBlockWithZero
            cs = argcount()
            rc = argcount()
            for _ in range(rc):
                out.extend(b'\x00' * cnt)
                out.extend(raw[i:i+cs]); i += cs
            out.extend(b'\x00' * cnt)
        else:
            raise ValueError("bad PIDATA opcode %d at %d" % (opcode, i))
    return bytes(out[:unpacked_size])

# ---- Loader section parsing -------------------------------------------------

def parse_loader(sec):
    b = sec.data()
    (main_section, main_offset, init_section, init_offset,
     term_section, term_offset, imported_lib_count, total_imported_sym_count,
     reloc_section_count, reloc_instr_offset, loader_strings_offset,
     export_hash_offset, export_hash_table_power, exported_symbol_count) = \
        struct.unpack_from('>iIiIiIIIIIIIII', b, 0)
    info = dict(main_section=main_section, main_offset=main_offset,
                init_section=init_section, init_offset=init_offset,
                term_section=term_section, term_offset=term_offset,
                imported_lib_count=imported_lib_count,
                total_imported_sym_count=total_imported_sym_count,
                reloc_section_count=reloc_section_count,
                reloc_instr_offset=reloc_instr_offset,
                loader_strings_offset=loader_strings_offset,
                export_hash_offset=export_hash_offset,
                export_hash_table_power=export_hash_table_power,
                exported_symbol_count=exported_symbol_count)
    # Imported library table: starts at offset 56
    libs = []
    o = 56
    for _ in range(imported_lib_count):
        (name_off, old_impl_ver, curr_ver, imp_sym_count,
         first_imp_sym, options, rsvdA, rsvdB) = struct.unpack_from('>IIIIIBBH', b, o)
        libs.append(dict(name_off=name_off, imp_sym_count=imp_sym_count,
                         first_imp_sym=first_imp_sym, options=options))
        o += 24
    # Imported symbol table follows lib table
    imp_syms = []
    for _ in range(total_imported_sym_count):
        val = u32(b, o); o += 4
        sym_class = val >> 24
        name_off = val & 0x00FFFFFF
        imp_syms.append((sym_class, name_off))
    strtab = info['loader_strings_offset']
    for lib in libs:
        lib['name'] = read_cstr(b, strtab + lib['name_off'])
        syms = []
        for k in range(lib['first_imp_sym'], lib['first_imp_sym'] + lib['imp_sym_count']):
            cls, noff = imp_syms[k]
            syms.append((cls, read_cstr(b, strtab + noff)))
        lib['symbols'] = syms
    # Exported symbols (loader export hash). Layout:
    #   export hash table (2^power u32 slots) at export_hash_offset
    #   then key table (u32 per exported symbol)
    #   then exported symbol table (10 bytes each: class/name u32, value u32, sect u16)
    exports = []
    hash_slots = 1 << export_hash_table_power
    key_tab_off = export_hash_offset + hash_slots * 4
    sym_tab_off = key_tab_off + exported_symbol_count * 4
    for k in range(exported_symbol_count):
        eo = sym_tab_off + k * 10
        clsname = u32(b, eo)
        value = u32(b, eo + 4)
        sect = u16(b, eo + 8)
        sym_class = clsname >> 24
        name_off = clsname & 0x00FFFFFF
        name = read_cstr(b, strtab + name_off)
        exports.append(dict(name=name, sym_class=sym_class, value=value, section=sect))
    return info, libs, exports

SYM_CLASS = {0: 'code(TVect)', 1: 'data', 2: 'TVector', 3: 'TOC/Glue',
             4: 'linker/glue', 0x40: 'weak?'}

# ---- commands ---------------------------------------------------------------

def cmd_header(pef):
    print("PEF container: arch=%s formatVersion=%d timestamp=0x%08x currentVer=0x%08x"
          % (pef.arch.decode('ascii','replace'), pef.format_version, pef.timestamp, pef.current_ver))
    print("sectionCount=%d instantiatedSectionCount=%d" % (pef.section_count, pef.inst_section_count))
    print()
    print("idx name           kind                default_addr  total     unpacked  packed    fileoff   align")
    for s in pef.sections:
        print("%3d %-14s %-19s 0x%08x    %-9d %-9d %-9d 0x%06x  %d" % (
            s.idx, (s.name or '-'), "%d:%s" % (s.kind, SECTION_KINDS.get(s.kind, '?')),
            s.default_addr, s.total_size, s.unpacked_size, s.packed_size,
            s.container_off, s.align))

def cmd_loader(pef):
    lsec = pef.loader()
    if not lsec:
        print("no loader section"); return
    info, libs, exports = parse_loader(lsec)
    print("=== Loader header ===")
    for k, v in info.items():
        if isinstance(v, int) and v > 9 and k.endswith(('offset','section','count')) is False:
            print("  %-28s 0x%x (%d)" % (k, v, v))
        else:
            print("  %-28s %s" % (k, v))
    ms = info['main_section']
    print("  ENTRY main:  section %d offset 0x%x" % (info['main_section'], info['main_offset']) if ms>=0 else "  (no main)")
    print("  ENTRY init:  section %d offset 0x%x" % (info['init_section'], info['init_offset']))
    print("  ENTRY term:  section %d offset 0x%x" % (info['term_section'], info['term_offset']))
    print()
    print("=== Imported libraries (%d) ===" % len(libs))
    for lib in libs:
        print("\n-- %s  (%d symbols) --" % (lib['name'], lib['imp_sym_count']))
        for j,(cls,name) in enumerate(lib['symbols']):
            print("   [%4d] cls=%s %s" % (lib['first_imp_sym']+j, SYM_CLASS.get(cls, hex(cls)), name))
    print()
    print("=== Exported symbols (%d) ===" % len(exports))
    for e in exports:
        print("   %-40s cls=%s value=0x%x sect=%d" % (e['name'], SYM_CLASS.get(e['sym_class'], hex(e['sym_class'])), e['value'], e['section']))

def iter_strings(data, minlen=4):
    cur = bytearray()
    start = 0
    printable = set(bytes(string.printable, 'ascii')) - set(b'\t\n\r\x0b\x0c')
    printable |= set(b' ')
    for i, ch in enumerate(data):
        if ch in printable:
            if not cur:
                start = i
            cur.append(ch)
        else:
            if len(cur) >= minlen:
                yield start, cur.decode('ascii', 'replace')
            cur = bytearray()
    if len(cur) >= minlen:
        yield start, cur.decode('ascii', 'replace')

def cmd_strings(pef, minlen=4):
    for s in pef.sections:
        if s.kind in (0,1,2,3):
            data = s.data()
            print("### section %d (%s) base=0x%x ###" % (s.idx, SECTION_KINDS.get(s.kind), s.default_addr))
            for off, st in iter_strings(data, minlen):
                print("0x%08x  %s" % (s.default_addr + off, st))

def cmd_dump(pef, secidx, off, length):
    data = pef.sections[secidx].data()
    chunk = data[off:off+length]
    for i in range(0, len(chunk), 16):
        row = chunk[i:i+16]
        hexs = ' '.join('%02x' % c for c in row)
        asci = ''.join(chr(c) if 32 <= c < 127 else '.' for c in row)
        print('%08x  %-48s  %s' % (off+i, hexs, asci))

def cmd_dis(pef, secidx=0, start=0, count=200):
    import capstone
    sec = pef.sections[secidx]
    data = sec.data()
    md = capstone.Cs(capstone.CS_ARCH_PPC, capstone.CS_MODE_32 | capstone.CS_MODE_BIG_ENDIAN)
    md.detail = False
    base = sec.default_addr
    for insn in md.disasm(data[start:start+count*4], base+start):
        print("0x%08x  %-8s %s" % (insn.address, insn.mnemonic, insn.op_str))

def main():
    path = sys.argv[1]
    cmd = sys.argv[2] if len(sys.argv) > 2 else 'header'
    with open(path, 'rb') as f:
        blob = f.read()
    pef = PEF(blob)
    if cmd == 'header': cmd_header(pef)
    elif cmd == 'loader': cmd_loader(pef)
    elif cmd == 'strings':
        minlen = int(sys.argv[3]) if len(sys.argv) > 3 else 4
        cmd_strings(pef, minlen)
    elif cmd == 'dump':
        cmd_dump(pef, int(sys.argv[3]), int(sys.argv[4],0), int(sys.argv[5],0))
    elif cmd == 'dis':
        secidx = int(sys.argv[3]) if len(sys.argv) > 3 else 0
        start = int(sys.argv[4],0) if len(sys.argv) > 4 else 0
        count = int(sys.argv[5],0) if len(sys.argv) > 5 else 200
        cmd_dis(pef, secidx, start, count)
    else:
        print("unknown cmd", cmd)

if __name__ == '__main__':
    main()