Memoization solution using qiling for VolgaCTF 2020 - f-hash task

qiling_bf.py

from qiling import *
from unicorn import *

import struct

BASE = 0x555555554000
REC_START = BASE + 0x13b0
REC_END = BASE + 0x1424

INLAST = False

# Dict used for memoization (idx, v1, v2) -> (r1, r2, r3)
MEMO = {}
TOSAVE = []

def hook_recursive(ql, address, size):
    global TOSAVE
    global MEMO

    if address == REC_START:
        result_struct = ql.uc.reg_read(UC_X86_REG_RDI)
        counter = ql.uc.reg_read(UC_X86_REG_RSI)
        v1 = ql.uc.reg_read(UC_X86_REG_RDX)
        v2 = ql.uc.reg_read(UC_X86_REG_RCX)

        # If we already saw this call we just reuse the result
        if (counter, v1, v2) in MEMO:
            r1, r2, r3 = MEMO[(counter, v1, v2)]
            e1 = struct.pack("<Q", r1)
            e2 = struct.pack("<Q", r2)
            e3 = struct.pack("<Q", r3)

            # Write in passed pointer
            ql.mem_write(result_struct, e1 + e2 + e3)

            # Write return value in rax
            ql.uc.reg_write(UC_X86_REG_RAX, result_struct)

            # Get from return address and jump
            ql.pc = ql.stack_pop()

            print("MEMO")

            return

        TOSAVE.append((counter, v1, v2))

    if address == REC_END:
        counter, v1, v2 = TOSAVE.pop()

        rax = ql.uc.reg_read(UC_X86_REG_RAX)
        r1 = ql.unpack64(ql.mem_read(rax, 8))
        r2 = ql.unpack64(ql.mem_read(rax + 0x8, 8))
        r3 = ql.unpack64(ql.mem_read(rax + 0x10, 8))

        print(f"Saving: {counter} 0x{v1:08x} 0x{v2:08x} (values: 0x{r1:08x} 0x{r2:08x} 0x{r3:08x})")

        MEMO[(counter, v1, v2)] = (r1, r2, r3)

if __name__ == '__main__':
    ql = Qiling(["./f-hash"], "/")
    ql.show_map_info()
    ql.hook_code(hook_recursive)
    ql.run()