solver for astral (ICMTC 2023)

astral.py

from capstone import *
from unicorn import *
from unicorn.x86_const import *
from z3 import *



def get_code():
    with open("Astral.exe", "rb") as fp:
        fp.seek(0x122a)
        return fp.read(0x40b5 - 0x122a)

def disasm(code):
    md = Cs(CS_ARCH_X86, CS_MODE_64)
    for i in md.disasm(code, 0x1000):
        #return ("0x%x:\t%s\t%s" %(i.address, i.mnemonic, i.op_str))
        return (i.address, i.mnemonic, i.op_str)


"""
unsure why it's broken
def hook_intr(mu, intno, data):
    rip = mu.reg_read(UC_X86_REG_RIP)
    mu.reg_write(UC_X86_REG_RIP, rip + 2)
    mu.reg_write(UC_X86_REG_EFLAGS, 0x2)
    print("interrupt 0x%x"%intno)
    return
"""

class VM:

    def __init__(self):
        # vm state
        self.eax_is_user = False    # is it moving user input ?
        self.table = [0, ]*8
        self.state = {}

        # symbolic flag
        self.input = []
        for x in range(0x48):
            f = BitVec("i%d"%x, 8)
            self.input.append(f)

    def hook_code(self, mu, address, size, user_data):
        instruction = mu.mem_read(address, size)
        instruction_str = ''.join('{:02x} '.format(x) for x in instruction)

        addr, mnemonic, op_str = disasm(instruction)
        
        #print('# 0x%x, size = 0x%x: %s\t%s' % (address, size, mnemonic, op_str))

        # idiv esi
        #if instruction == b'\xf7\xfe':
        # division by zero, triggers VEH / VM handler
        if mnemonic == "idiv" and op_str == "esi":
            rip = mu.reg_read(UC_X86_REG_RIP)
            mu.reg_write(UC_X86_REG_RIP, rip + 2)
            #print("division by zero, handle vm instruction...")
            self.vm_handler(mu)

        # losy detection of when it loads user input data to eax
        elif mnemonic == "mov" and "eax" in op_str:
            if "rbp" not in op_str:
                self.eax_is_user = False
            else:
                self.eax_is_user = True
                # get stack string index and save symolic data
                index = 0x48 - int(op_str.split()[-1][2:].strip("]"), 16)
                self.state["eax"] = Concat(self.input[index + 3], self.input[index + 2], self.input[index + 1], self.input[index])


    def emu(self, code):
        # flag starts @ rbp-0x48
        ADDR_TEXT = 0x1000000
        ADDR_STACK = 0x7000000

        mu = Uc(UC_ARCH_X86, UC_MODE_64)

        mu.mem_map(ADDR_TEXT, 0x3000)
        mu.mem_map(ADDR_STACK, 0x4000)

        mu.mem_write(ADDR_TEXT, code)

        rsp = ADDR_STACK + 0x400
        rbp = ADDR_STACK + 0x1000
        mu.reg_write(UC_X86_REG_RBP, rbp)
        mu.reg_write(UC_X86_REG_RSP, rsp)

        #mu.mem_write(rbp - 0x48, b"789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`abcdefghijklmnopqrstuvwxyz")
        mu.mem_write(rbp - 0x48, b"AAAABBBBCCCCDDDDEEEEFFFFGGGGHHHHIIIIJJJJKKKKLLLLMMMMNNNNOOOOPPPPQQQQR")

        mu.hook_add(UC_HOOK_CODE, self.hook_code)
        #mu.hook_add(UC_HOOK_INTR, hook_intr)
        mu.emu_start(ADDR_TEXT, ADDR_TEXT + len(code))
        print("end_emu")


    def vm_handler(self, mu):
        """
        only these opcodes are used:
         0x10 0x12 0x13 0x15 0x16 0xa
        """
        ecx = mu.reg_read(UC_X86_REG_ECX)
        eax = mu.reg_read(UC_X86_REG_EAX)
        opcode = ecx & 0xff
        ecx_10 = ecx >> 0x10
        ecx_8 = ecx >> 8 & 7

        #print("----------- VM")
        #print(" opcode: 0x%x / 0x%x"%(opcode, ecx_10))

        match opcode:
            case 0xa:
                # mov array, imm
                print("vm_mov table[%d] , 0x%x"%(ecx_8, eax))
                if self.eax_is_user:
                    self.table[ecx_8] = self.state["eax"]
                    #self.table[ecx_8] = eax
                else:
                    self.table[ecx_8] = eax

            case 0x10:
                # xor array, array
                print("vm_xor table[%d] , table[%d]"%(ecx_8, ecx_10))
                self.table[ecx_8] ^= self.table[ecx_10]

            case 0x12:
                # or array, array
                print("vm_or table[%d] , table[%d]"%(ecx_8, ecx_10))
                self.table[ecx_8] |= self.table[ecx_10]
            
            case 0x13:
                # if table[7] == 0 -> flag OK
                print("vm_check_result")
                s = Solver()
                s.add(self.table[7] == 0)
                assert s.check() == sat

                model = s.model()
                flag = ''
                for x in self.input:
                    if model[x] is not None:
                        flag += chr(model[x].as_long())
                print(flag)
        
            case 0x15:
                # zero * at the beginning
                print("vm_init_zero_mem")

            case 0x16:
                # mov array, array
                print("vm_mov table[%d] , table[%d]"%(ecx_8, ecx_10))
                self.table[ecx_8] = self.table[ecx_10]

            case _:
                print("unknown opcode 0x%x"%opcode)
                

def main():
    vm = VM()
    vm.emu(get_code())


if __name__ == "__main__":
    main()