#!/usr/bin/env python3
"""Extract bitmap images from .egf files (Win32 PE resource containers).

EGF files are PE binaries whose images Resource Hacker shows as RT_BITMAP
(type 2) resources, stored as raw DIBs (BITMAPINFOHEADER, no 14-byte BM
file header). This walks the PE resource directory, pulls every bitmap, and
writes a proper .bmp by prepending a reconstructed BITMAPFILEHEADER.

Usage:
    python3 extract_egf.py [file_or_dir ...]
No args -> all *.egf in current directory.
"""
import struct, sys, os, glob

RT_BITMAP = 2

def u16(d, o): return struct.unpack_from('<H', d, o)[0]
def u32(d, o): return struct.unpack_from('<I', d, o)[0]

def get_sections(d):
    e_lfanew = u32(d, 0x3c)
    if d[e_lfanew:e_lfanew+4] != b'PE\x00\x00':
        raise ValueError('not a PE file')
    coff = e_lfanew + 4
    nsec = u16(d, coff + 2)
    sizeopt = u16(d, coff + 16)
    secoff = coff + 20 + sizeopt
    secs = []
    for i in range(nsec):
        o = secoff + i * 40
        vaddr = u32(d, o + 12)
        vsize = u32(d, o + 8)
        raw = u32(d, o + 20)
        rsize = u32(d, o + 16)
        secs.append((vaddr, vsize, raw, rsize))
    return secs

def make_rva2off(secs):
    def rva2off(rva):
        for vaddr, vsize, raw, rsize in secs:
            if vaddr <= rva < vaddr + max(vsize, rsize):
                return rva - vaddr + raw
        raise ValueError('rva %x not in any section' % rva)
    return rva2off

def dib_to_bmp(dib):
    """Prepend BITMAPFILEHEADER. Compute pixel-data offset from DIB header."""
    biSize = u32(dib, 0)
    biBitCount = u16(dib, 14)
    biClrUsed = u32(dib, 32)
    # palette size
    if biBitCount <= 8:
        n = biClrUsed if biClrUsed else (1 << biBitCount)
        palbytes = n * 4
    else:
        palbytes = 0
        # BI_BITFIELDS (compression==3) adds 3 (or 4) DWORD masks
        biCompression = u32(dib, 16)
        if biCompression == 3:
            palbytes = 3 * 4
    pixoff = 14 + biSize + palbytes
    fileheader = b'BM' + struct.pack('<IHHI', 14 + len(dib), 0, 0, pixoff)
    return fileheader + dib

def parse_dir(d, rsrc_base, off):
    nnamed = u16(d, off + 12)
    nid = u16(d, off + 14)
    base = off + 16
    out = []
    for i in range(nnamed + nid):
        e = base + i * 8
        nameid = u32(d, e)
        offfield = u32(d, e + 4)
        out.append((nameid, offfield))
    return out

def resname(d, rsrc_base, nameid):
    if nameid & 0x80000000:
        no = rsrc_base + (nameid & 0x7fffffff)
        ln = u16(d, no)
        return d[no+2:no+2+ln*2].decode('utf-16le', 'replace')
    return str(nameid)

def extract(path, outdir):
    d = open(path, 'rb').read()
    secs = get_sections(d)
    rva2off = make_rva2off(secs)
    rsrc_va = next(v for (v, vs, r, rs) in secs)  # placeholder, fixed below
    # find .rsrc by RVA range that holds resource dir: use section flagged via data dir? simplest: locate .rsrc by name
    # reparse to grab .rsrc section base
    e_lfanew = u32(d, 0x3c); coff = e_lfanew + 4
    nsec = u16(d, coff + 2); sizeopt = u16(d, coff + 16)
    secoff = coff + 20 + sizeopt
    rsrc_va = rsrc_raw = None
    for i in range(nsec):
        o = secoff + i * 40
        if d[o:o+8].rstrip(b'\x00') == b'.rsrc':
            rsrc_va = u32(d, o + 12); rsrc_raw = u32(d, o + 20)
    if rsrc_va is None:
        print('  no .rsrc in %s' % path); return 0
    rsrc_base = rsrc_raw  # offset of resource dir start (= .rsrc raw, since dir at section start)

    count = 0
    top = parse_dir(d, rsrc_base, rsrc_base)
    for tid, toff in top:
        if (tid & 0x7fffffff) != RT_BITMAP or not (toff & 0x80000000):
            continue
        namedir = rsrc_base + (toff & 0x7fffffff)
        for nid, noff in parse_dir(d, rsrc_base, namedir):
            name = resname(d, rsrc_base, nid)
            if not (noff & 0x80000000):
                continue
            langdir = rsrc_base + (noff & 0x7fffffff)
            for lid, loff in parse_dir(d, rsrc_base, langdir):
                if loff & 0x80000000:
                    continue  # expect leaf
                leaf = rsrc_base + (loff & 0x7fffffff)
                data_rva = u32(d, leaf)
                size = u32(d, leaf + 4)
                doff = rva2off(data_rva)
                dib = d[doff:doff+size]
                bmp = dib_to_bmp(dib)
                base = os.path.splitext(os.path.basename(path))[0]
                safe = name.replace('/', '_').replace('\\', '_')
                fn = '%s_%s.bmp' % (base, safe)
                open(os.path.join(outdir, fn), 'wb').write(bmp)
                count += 1
    return count

def main():
    args = sys.argv[1:]
    files = []
    if not args:
        files = sorted(glob.glob('*.egf'))
    for a in args:
        if os.path.isdir(a):
            files += sorted(glob.glob(os.path.join(a, '*.egf')))
        else:
            files.append(a)
    if not files:
        print('no .egf files'); return
    total = 0
    for f in files:
        outdir = os.path.splitext(f)[0] + '_extracted'
        os.makedirs(outdir, exist_ok=True)
        try:
            n = extract(f, outdir)
        except Exception as ex:
            print('  ERROR %s: %s' % (f, ex)); continue
        print('%s -> %d images (%s)' % (f, n, outdir))
        total += n
    print('TOTAL %d images' % total)

if __name__ == '__main__':
    main()