#!/usr/bin/env python3
"""
leimage.py -- minimal loader for DOS/4G "LE" (Linear Executable) images,
as produced by Watcom C/C++32 + the Rational/Tenberry DOS/4G(W) extender.

Path A tooling for dpl3-revive: lets us inspect FLYK.EXE (and the other HPDAVE
binaries) -- object table, names tables, debug-info presence -- and reconstruct
the flat code/data images for capstone disassembly.

Usage:
    python leimage.py header  FLYK.EXE      # LE header fields
    python leimage.py names   FLYK.EXE      # resident + non-resident name tables
    python leimage.py objects FLYK.EXE      # object table + page reconstruction sanity
    python leimage.py strings FLYK.EXE PAT  # locate a symbol/string + file offset
"""
import sys, struct

def rd(f, off, n):
    f.seek(off); return f.read(n)

class LE:
    def __init__(self, path):
        self.path = path
        with open(path, 'rb') as f:
            mz = f.read(0x40)
            assert mz[:2] == b'MZ', "not an MZ file"
            self.lfanew = struct.unpack_from('<I', mz, 0x3c)[0]
            hdr = rd(f, self.lfanew, 0xB0)
        assert hdr[:2] == b'LE', "no LE header (got %r)" % hdr[:2]
        self.base = self.lfanew
        u = lambda o: struct.unpack_from('<I', hdr, o)[0]
        h = lambda o: struct.unpack_from('<H', hdr, o)[0]
        self.cpu            = h(0x08)
        self.os             = h(0x0A)
        self.mod_flags      = u(0x10)
        self.n_pages        = u(0x14)
        self.eip_obj        = u(0x18)
        self.eip            = u(0x1C)
        self.esp_obj        = u(0x20)
        self.esp            = u(0x24)
        self.page_size      = u(0x28)
        self.last_page      = u(0x2C)
        self.fixup_size     = u(0x30)
        self.loader_size    = u(0x38)
        self.obj_tab_off    = u(0x40) + self.base     # relative to header
        self.n_objects      = u(0x44)
        self.obj_page_off   = u(0x48) + self.base
        self.res_names_off  = u(0x58) + self.base
        self.entry_off      = u(0x5C) + self.base
        self.fixup_pg_off   = u(0x68) + self.base
        self.fixup_rec_off  = u(0x6C) + self.base
        self.data_pages_off = u(0x80)                 # absolute from file start
        self.nonres_off     = u(0x88)                 # absolute from file start
        self.nonres_len     = u(0x8C)
        self.debug_off      = u(0x98)                 # absolute; 0 = none
        self.debug_len      = u(0x9C)

    def header(self):
        cpu = {2:'80386',3:'80486',4:'Pentium'}.get(self.cpu, hex(self.cpu))
        print(f"file            {self.path}")
        print(f"MZ e_lfanew     {self.lfanew:#x}  (LE header base)")
        print(f"cpu             {cpu}    os={self.os:#x}   flags={self.mod_flags:#x}")
        print(f"pages           {self.n_pages}   page_size={self.page_size:#x}   last_page_bytes={self.last_page:#x}")
        print(f"entry (CS:EIP)  obj{self.eip_obj}:{self.eip:#x}")
        print(f"stack (SS:ESP)  obj{self.esp_obj}:{self.esp:#x}")
        print(f"objects         {self.n_objects}  @ {self.obj_tab_off:#x}   page_map @ {self.obj_page_off:#x}")
        print(f"loader_size     {self.loader_size:#x}   fixup_size={self.fixup_size:#x}")
        print(f"data_pages      @ {self.data_pages_off:#x}")
        print(f"resident names  @ {self.res_names_off:#x}")
        print(f"nonres  names   @ {self.nonres_off:#x}  len={self.nonres_len:#x}")
        print(f"DEBUG info      @ {self.debug_off:#x}  len={self.debug_len:#x}   -> {'PRESENT' if self.debug_len else 'none'}")

    def objects(self):
        with open(self.path,'rb') as f:
            data = rd(f, self.obj_tab_off, 24*self.n_objects)
        for i in range(self.n_objects):
            vsize, base, flags, pgidx, pgcnt, resv = struct.unpack_from('<6I', data, i*24)
            fl = []
            if flags & 0x1: fl.append('READ')
            if flags & 0x2: fl.append('WRITE')
            if flags & 0x4: fl.append('EXEC')
            if flags & 0x2000: fl.append('32BIT')
            print(f"obj{i+1}  vsize={vsize:#010x}  base={base:#010x}  flags={flags:#06x} [{'|'.join(fl)}]  "
                  f"pagemap[{pgidx}..{pgidx+pgcnt-1}] ({pgcnt} pages)")

    def image(self, objno):
        """Reconstruct the flat linear image of object `objno` (1-based).
        Assumes sequential, non-iterated pages (true for DOS/4GW LE)."""
        with open(self.path, 'rb') as f:
            objs = rd(f, self.obj_tab_off, 24 * self.n_objects)
            vsize, base, flags, pgidx, pgcnt, _ = struct.unpack_from('<6I', objs, (objno-1)*24)
            out = bytearray()
            for k in range(pgidx, pgidx + pgcnt):          # k = global 1-based page number
                length = self.last_page if k == self.n_pages else self.page_size
                out += rd(f, self.data_pages_off + (k-1)*self.page_size, length)
        out = out[:vsize] + bytes(max(0, vsize - len(out)))  # trim / zero-fill to vsize
        return base, bytes(out)

    def _read_names(self, off, hardlen=None):
        """Read a chain of length-prefixed name entries: len(1), chars, ordinal(2)."""
        out = []
        with open(self.path,'rb') as f:
            f.seek(off)
            blob = f.read(hardlen if hardlen else 0x8000)
        p = 0
        while p < len(blob):
            ln = blob[p]; p += 1
            if ln == 0: break
            name = blob[p:p+ln]; p += ln
            if p+2 > len(blob): break
            ordv = struct.unpack_from('<H', blob, p)[0]; p += 2
            out.append((name.decode('latin1'), ordv))
        return out

    def names(self):
        print("== resident names ==")
        for n,o in self._read_names(self.res_names_off):
            print(f"  [{o:5}] {n}")
        print("== non-resident names ==")
        nr = self._read_names(self.nonres_off, self.nonres_len) if self.nonres_len else []
        for n,o in nr:
            print(f"  [{o:5}] {n}")
        if not nr:
            print("  (none)")

if __name__ == '__main__':
    cmd, path = sys.argv[1], sys.argv[2]
    le = LE(path)
    if cmd == 'header':  le.header()
    elif cmd == 'objects': le.objects()
    elif cmd == 'names': le.names()
    else: print("commands: header | objects | names")