unpacker for nanobits (pwnme 2023)

nanofix.py

#!/usr/bin/env python3
# unpacker for nanobits (PWNME 2023)

import io
import struct
import sys
import zlib
import enum
import os

import pefile
from keystone import *

DEBUG = False

def log(s):
    if DEBUG:
        print(s)


class OP(enum.IntEnum):
    JE = 0x1
    JNE = 0x2
    JMP = 0x3
    JG = 0x4
    JGE = 0x5
    JA = 0x6
    JAE = 0x7
    JL = 0x8
    JLE = 0x9
    JB = 0xa
    JBE = 0xb
    CALL = 0xc


def p64(bla):
    return struct.pack("<Q", bla)


def extract_rsrc(pe):
    """ stolen from binref, thx Jesko
    """
    def _search(pe: pefile.PE, directory, level=0, *parts):
        for entry in directory.entries:
            if entry.name:
                identifier = str(entry.name)
            elif entry.id is not None:
                identifier = entry.id
            else:
                identifier = "lol"

            if entry.struct.DataIsDirectory:
                yield from _search(pe, entry.directory, level + 1, *parts, identifier)
            else:
                rva = entry.data.struct.OffsetToData
                size = entry.data.struct.Size
                path = '/'.join(str(p) for p in (*parts, identifier))
                extract = None
                if extract is None:
                    def extract(pe=pe):
                        return pe.get_data(rva, size)
                yield {
                        'path': path,
                        'name': path.split("/")[1], 
                        'extract': extract, 
                        'offset': pe.get_offset_from_rva(rva)
                        }

    pe.parse_data_directories(directories=pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_RESOURCE'])
    rsrc = pe.DIRECTORY_ENTRY_RESOURCE
    return _search(pe, rsrc)


def fix_nanomites(child_name, nanomites):
    # load nanomites
    loaded_addr = {}
    nano_table = []
    while True:
        data = nanomites.read(6*4)
        if len(data) != 6*4:
            break

        data = struct.unpack("<QIIII", data)
        # only load first addr match
        if data[0] not in loaded_addr:
            nano_table.append(data)
            loaded_addr[data[0]] = 1

    pe = pefile.PE(child_name)
    ks = Ks(KS_ARCH_X86, KS_MODE_32)    # assembler

    # base address
    base = pe.OPTIONAL_HEADER.ImageBase

    # 
    for nano in nano_table:
        eip, op, _, target, _ = nano
        #print([hex(_) for _ in nano])

        # fix addresses
        target = base + target
        eip = base + eip

        if op == OP.CALL:
            asm = "call 0x%x"%target
        elif op == OP.JMP:
            asm = "jmp 0x%x"%target
        elif op == OP.JLE:
            asm = "jle 0x%x"%target
        elif op == OP.JE:
            asm = "je 0x%x"%target
        elif op == OP.JNE:
            asm = "jne 0x%x"%target
        elif op == OP.JG:
            asm = "jg 0x%x"%target
        elif op == OP.JGE:
            asm = "jge 0x%x"%target
        elif op == OP.JA:
            asm = "ja 0x%x"%target
        elif op == OP.JAE:
            asm = "jae 0x%x"%target
        elif op == OP.JL:
            asm = "jl 0x%x"%target
        elif op == OP.JB:
            asm = "jb 0x%x"%target
        elif op == OP.JBE:
            asm = "jbe 0x%x"%target
        elif not op:
            pass
        else:
            print("unsupported op %d"%op)
            raise

        opcodes = ks.asm(asm, addr=eip, as_bytes=True)[0]
        log([hex(eip), asm, op, opcodes])

        patch(pe, eip-base, opcodes)
        #patch(pe, eip, opcodes)

    pe.write("lolz.exe")
    os.unlink(child_name)
    os.rename("lolz.exe", child_name)
    print("extracted/fixed bin as %s"%child_name)

def patch(pe, eip, opcodes):
    data = p64(pe.get_qword_at_rva(eip))
    #print(data)
    if data[0] != 0xcc:
        log("err: 0x%x not a CC"%eip)
        return

    noplen = len(data) - len(data[1:].lstrip(b'\x90'))
    if noplen < len(opcodes):
        log("err: 0x%x not enough nops (%d) for opcodes (%d)"%(eip, noplen, len(opcodes)))
        return

    pe.set_bytes_at_rva(eip, opcodes)


def main(argv):
    pe = pefile.PE(argv[1])
    child_name = 'child_' + argv[1]

    got_pe = False
    got_nanomites = False

    for rsrc in extract_rsrc(pe):
        if rsrc['name'] == '999':
            child_pe = zlib.decompress(rsrc['extract']()[16:])
            with open(child_name, "wb") as fp:
                fp.write(child_pe)
            got_pe = True
        elif rsrc['name'] == '777':
            nanomites = zlib.decompress(rsrc['extract']()[16:])
            got_nanomites = True

    # got both pe and nanomites ?
    assert got_pe and got_nanomites == True

    fix_nanomites(child_name, io.BytesIO(nanomites))


if __name__ == "__main__":
    main(sys.argv)