solver for nightmare (ICMTC 2023)

nightmare.py

import sys
sys.path.append("/data/opt/triton/lib/python3.11/site-packages/")
from triton import *


BASE_ARGV  = 0x20000000
BASE_STACK = 0x9ffffff0

FLAG_LEN = 0x30

conditions = [
    0x140002566,
    0x14000263f,
    0x14000279a,
    0x14000281d,
    0x140002ade,
    0x140002d18,
    0x140002e4c,
    0x140002efe,
    0x1400030e1,
    0x14000321b,
    0x1400032b3,
    0x140003350,
    0x14000355f,
    0x140003652,
    0x140003866,
    0x140003901,
    0x140003a46,
    0x140003a6b,
    0x140003b13,
    0x140003bc1,
    0x140003ce1,
    0x140003d12,
    0x140003e2e,
    0x140003f5d,
    0x14000414e,
    0x14000419a,
    0x1400047bc,
    0x14000484b,
    0x140004865,
    0x140004971,
    0x140004b70,
    0x140004d94,
    0x140004d4e,
    0x140004d57,
    0x140004d76,
    0x140004d7f,
    0x140004d8a,
    0x140004d94
]

def get_code():
    """ load main function bytes
    """
    main_offset = [0x470, 0x4204]
    with open("nightmare.exe", "rb") as fp:
        fp.seek(main_offset[0])
        return fp.read(main_offset[1] - main_offset[0])

def emulate(ctx, pc):
    while pc:
        opcode = ctx.getConcreteMemoryAreaValue(pc, 16)
        instruction = Instruction(pc, opcode)
        instruct_len = instruction.getOpcode()
        addr = instruction.getAddress()
        ctx.processing(instruction)
        #print(instruction)


        # win
        if addr == 0x140004da2:
            flag_base = ctx.getConcreteRegisterValue(ctx.registers.r8)
            flag = ''
            for x in range(FLAG_LEN):
                flag += chr(ctx.getConcreteMemoryValue(flag_base + x))
                
            print(flag)
            sys.exit(0)
        
        if addr in conditions:
            print("solving @ 0x%x"%addr)
            zf = ctx.getSymbolicRegister(ctx.registers.zf).getAst()
            ctx.pushPathConstraint(zf == 1)
            mod = ctx.getModel(ctx.getPathPredicate())
            for k,v in list(mod.items()):
                ctx.setConcreteVariableValue(ctx.getSymbolicVariable(v.getId()), v.getValue())

        # skip calls
        if instruction.getType() == OPCODE.X86.CALL:
            break
        pc = ctx.getConcreteRegisterValue(ctx.registers.rip)



def main():
    ctx = TritonContext(ARCH.X86_64)
    # Set a symbolic optimization mode
    ctx.setMode(MODE.ALIGNED_MEMORY, True)
    ctx.setMode(MODE.ONLY_ON_SYMBOLIZED, True)
    
    ast = ctx.getAstContext()

    ctx.setConcreteMemoryAreaValue(0x140001070, get_code())

    # argc
    ctx.setConcreteRegisterValue(ctx.registers.ecx, 2) 

    # argv
    argvs = [
            b'lol.exe',
            b'b'*FLAG_LEN
        ]
    base  = BASE_ARGV
    addrs = list()

    for argv in argvs:
        addrs.append(base)
        ctx.setConcreteMemoryAreaValue(base, argv + b'\x00')
        base += len(argv) + 1

    argv = base
    for addr in addrs:
        ctx.setConcreteMemoryValue(MemoryAccess(base, CPUSIZE.QWORD), addr)
        base += CPUSIZE.QWORD

    ctx.setConcreteRegisterValue(ctx.registers.rdx, argv)

    # symbolize argv[1]
    argv1 = ctx.getConcreteMemoryValue(MemoryAccess(ctx.getConcreteRegisterValue(ctx.registers.rdx) + 8, CPUSIZE.QWORD))
    for index in range(len(argvs[1])):
        var = ctx.symbolizeMemory(MemoryAccess(argv1+index, CPUSIZE.BYTE))
        vast = ast.variable(var)
        #ctx.pushPathConstraint(ast.land([vast >= ord(b' '), vast <= ord(b'~')]))

    # stack
    ctx.setConcreteRegisterValue(ctx.registers.rbp, BASE_STACK)
    ctx.setConcreteRegisterValue(ctx.registers.rsp, BASE_STACK)

    emulate(ctx, 0x140001070)


if __name__ == "__main__":
    main()