# coding=utf-8
#
# Copyright 2014 Sascha Schirra
#
# This file is part of Ropper.
#
# Ropper is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ropper is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
from ropper.gadget import Category
from ropper.common.error import *
from ropper.common.utils import *
from ropper.rop import Ropper
from ropper.arch import x86
from ropper.ropchain.ropchain import *
from ropper.loaders.loader import Type
from ropper.gadget import Gadget
from re import match
from filebytes.pe import ImageDirectoryEntry
import itertools
import math
import sys

if sys.version_info.major == 2:
    range = xrange

class RopChainX86(RopChain):

    MAX_QUALI = 7

    def _printHeader(self):
        toReturn = ''
        toReturn += ('#!/usr/bin/env python\n')
        toReturn += ('# Generated by ropper ropchain generator #\n')
        toReturn += ('from struct import pack\n')
        toReturn += ('\n')
        toReturn += ('p = lambda x : pack(\'I\', x)\n')

        toReturn += ('\n')

        return toReturn

    def _printRebase(self):
        toReturn = ''
        for binary,section in self._usedBinaries:
            imageBase = Gadget.IMAGE_BASES[binary]
            toReturn += ('IMAGE_BASE_%d = %s # %s\n' % (self._usedBinaries.index((binary, section)),toHex(imageBase , 4), binary))
            toReturn += ('rebase_%d = lambda x : p(x + IMAGE_BASE_%d)\n\n'% (self._usedBinaries.index((binary, section)),self._usedBinaries.index((binary, section))))
        return toReturn

    @classmethod
    def name(cls):
        return ''

    @classmethod
    def availableGenerators(cls):
        return [RopChainX86System, RopChainX86Mprotect, RopChainX86VirtualProtect]

    @classmethod
    def archs(self):
        return [x86]

    def _createDependenceChain(self, gadgets):
        """
        gadgets - list with tuples

        tuple contains:
        - method to create chaingadget
        - list with arguments
        - dict with named arguments
        - list with registers which are not allowed to override in the gadget
        """
        failed = []
        cur_len = 0
        cur_chain = ''
        counter = 0
        failed_chains={}
        max_perm = math.factorial(len(gadgets))
        for x in itertools.permutations(gadgets):
            counter += 1
            self._printMessage('[*] Try permuation %d / %d' % (counter, max_perm))
            found = False
            for y in failed:

                if x[:len(y)] == y:
                    found = True
                    break
            if found:
                continue
            try:
                fail = []
                chain2 = ''
                dontModify = []
                badRegs = []
                c = 0
                for idx in range(len(x)):
                    g = x[idx]
                    if idx != 0:
                        badRegs.extend(x[idx-1][3])

                    dontModify.extend(g[3])
                    fail.append(g)
                    chain2 += g[0](*g[1], badRegs=badRegs, dontModify=dontModify,**g[2])[0]


                cur_chain += chain2
                break

            except RopChainError as e:
                failed_chains[chain2] = fail

            failed.append(tuple(fail))
        else:
            self._printMessage('Cannot create chain which fills all registers')
            fail_tmp = None
            fail_max = []
            chain_tmp = None
            for chain,fail in failed_chains.items():
                if len(fail) > len(fail_max):
                    fail_max = fail
                    chain_tmp = chain


            cur_chain = '# Filled registers: '
            for fa in fail_max[:-1]:

                cur_chain += (fa[2]['reg']) + ', '
            cur_chain += '\n'
            cur_chain += chain_tmp

        #    print('Impossible to create complete chain')
        return cur_chain

    def _isModifiedOrDereferencedAccess(self, gadget, dontModify):

        regs = []
        for line in gadget.lines[1:]:
            line = line[1]
            if '[' in line:
                return True
            if dontModify:
                m = match('[a-z]+ (e?[abcds][ixlh]),?.*', line)
                if m and m.group(1) in dontModify:
                    return True

        return False



    def _paddingNeededFor(self, gadget):
        regs = []
        for idx in range(1,len(gadget.lines)):
            line = gadget.lines[idx][1]
            matched = match('^pop (...)$', line)
            if matched:
                regs.append(matched.group(1))
        return regs


    def _printRopInstruction(self, gadget, padding=True, number=None):
        toReturn = ('rop += rebase_%d(%s) # %s\n' % (self._usedBinaries.index((gadget.fileName, gadget.section)),toHex(gadget.lines[0][0],4), gadget.simpleString()))
        if number is not None:
            toReturn +=self._printPaddingInstruction(number)        
        if padding:
            regs = self._paddingNeededFor(gadget)
            for i in range(len(regs)):
                toReturn +=self._printPaddingInstruction()
        return toReturn

    def _printAddString(self, string):
        return ('rop += \'%s\'\n' % string)

    def _printRebasedAddress(self, addr, comment='', idx=0):
        return ('rop += rebase_%d(%s)\n' % (idx,addr))

    def _printPaddingInstruction(self, addr='0xdeadbeef'):
        return ('rop += p(%s)\n' % addr)

    def _containsZeroByte(self, addr):
        return self.containsBadbytes(addr)

    def _createZeroByteFillerForSub(self, number):
        start = 0x01010101
        for i in range(start, 0x0f0f0f0f):
            if not self._containsZeroByte(i) and not self._containsZeroByte(i+number):
                return i

        raise RopChainError("Could not create Number for substract gadget")

    def _createZeroByteFillerForAdd(self, number):
        start = 0x01010101
        for i in range(start, 0x0f0f0f0f):
            if not self._containsZeroByte(i) and not self._containsZeroByte(number-i):
                return i

        raise RopChainError("Could not create Number for addition gadget")

    def _find(self, category, reg=None, srcdst='dst', badDst=[], badSrc=None, dontModify=None, srcEqDst=False, switchRegs=False ):
        quali = 1
        while quali < RopChainX86System.MAX_QUALI:
            for binary in self._binaries:
                for gadget in self._gadgets[binary]:
                    if gadget.category[0] == category and gadget.category[1] == quali:
                        if badSrc and gadget.category[2]['src'] in badSrc:
                            continue
                        if badDst and gadget.category[2]['dst'] in badDst:
                            continue
                        if not gadget.lines[len(gadget.lines)-1][1].strip().endswith('ret') or 'esp' in gadget.simpleString():
                            continue
                        if srcEqDst and (not (gadget.category[2]['dst'] == gadget.category[2]['src'])):
                            continue
                        elif not srcEqDst and 'src' in gadget.category[2] and (gadget.category[2]['dst'] == gadget.category[2]['src']):
                            continue
                        if self._isModifiedOrDereferencedAccess(gadget, dontModify):
                            continue
                        if reg:
                            if gadget.category[2][srcdst] == reg:
                                if (gadget.fileName, gadget.section) not in self._usedBinaries:
                                    self._usedBinaries.append((gadget.fileName, gadget.section))
                                return gadget
                            elif switchRegs:
                                other = 'src' if srcdst == 'dst' else 'dst'
                                if gadget.category[2][other] == reg:
                                    if (gadget.fileName, gadget.section) not in self._usedBinaries:
                                        self._usedBinaries.append((gadget.fileName, gadget.section))
                                    return gadget
                        else:
                            if (gadget.fileName, gadget.section) not in self._usedBinaries:
                                self._usedBinaries.append((gadget.fileName, gadget.section))
                            return gadget

            quali += 1


    def _createWriteStringWhere(self, what, where, reg=None, dontModify=[], idx=0):
        badRegs = []
        badDst = []
        while True:
            popReg = self._find(Category.LOAD_REG, reg=reg, badDst=badRegs, dontModify=dontModify)
            if not popReg:
                raise RopChainError('Cannot build writewhatwhere gadget!')
            write4 = self._find(Category.WRITE_MEM, reg=popReg.category[2]['dst'],  badDst=
            badDst, srcdst='src')
            if not write4:
                badRegs.append(popReg.category[2]['dst'])
                continue
            else:
                popReg2 = self._find(Category.LOAD_REG, reg=write4.category[2]['dst'], dontModify=[popReg.category[2]['dst']]+dontModify)
                if not popReg2:
                    badDst.append(write4.category[2]['dst'])
                    continue
                else:
                    break;

        if len(what) % 4 > 0:
            what += ' ' * (4 - len(what) % 4)
        toReturn = ''
        for index in range(0,len(what),4):
            part = what[index:index+4]

            toReturn += self._printRopInstruction(popReg,False)
            toReturn += self._printAddString(part)
            regs = self._paddingNeededFor(popReg)
            for i in range(len(regs)):
                toReturn +=self._printPaddingInstruction()
            toReturn += self._printRopInstruction(popReg2, False)

            toReturn += self._printRebasedAddress(toHex(where+index,4), idx=idx)
            regs = self._paddingNeededFor(popReg2)
            for i in range(len(regs)):
                toReturn +=self._printPaddingInstruction()
            toReturn += self._printRopInstruction(write4)
        return (toReturn,popReg.category[2]['dst'], popReg2.category[2]['dst'])


    def _createWriteRegValueWhere(self, what, where, dontModify=[], idx=0):
        badRegs = []
        badDst = []
        while True:


            write4 = self._find(Category.WRITE_MEM, reg=what,  badDst=badDst, dontModify=dontModify, srcdst='src')
            if not write4:
                raise RopChainError('Cannot build writewhatwhere gadget!')
            else:
                popReg2 = self._find(Category.LOAD_REG, reg=write4.category[2]['dst'], dontModify=[what]+dontModify)
                if not popReg2:
                    badDst.append(write4.category[2]['dst'])
                    continue
                else:
                    break;

        toReturn = self._printRopInstruction(popReg2, False)
        toReturn += self._printRebasedAddress(toHex(where,4), idx=idx)
        regs = self._paddingNeededFor(popReg2)
        for i in range(len(regs)):
            toReturn +=self._printPaddingInstruction()
        toReturn += self._printRopInstruction(write4)

        return (toReturn,what, popReg2.category[2]['dst'])

    def _createLoadRegValueFrom(self, what, from_reg, dontModify=[], idx=0):
        try:
            return self._createLoadRegValueFromMov(what, from_reg, dontModify, idx)
        except RopChainError:
            return self._createLoadRegValueFromXchg(what, from_reg, dontModify, idx)

    def _createLoadRegValueFromMov(self, what, from_reg, dontModify=[], idx=0):
        badRegs = []
        badDst = []
        while True:


            load4 = self._find(Category.LOAD_MEM, reg=what,  badDst=badDst, dontModify=dontModify, srcdst='dst')
            if not load4:
                raise RopChainError('Cannot build loadwhere gadget!')
            else:
                popReg2 = self._find(Category.LOAD_REG, reg=load4.category[2]['src'], dontModify=[what,load4.category[2]['src']]+dontModify)
                if not popReg2:
                    badDst.append(load4.category[2]['src'])
                    continue
                else:
                    break;

        toReturn = self._printRopInstruction(popReg2, False)
        toReturn += self._printPaddingInstruction(toHex(from_reg,4))
        regs = self._paddingNeededFor(popReg2)
        for i in range(len(regs)):
            toReturn +=self._printPaddingInstruction()
        toReturn += self._printRopInstruction(load4)

        return (toReturn,what, popReg2.category[2]['dst'])

    def _createLoadRegValueFromXchg(self, what, from_reg, dontModify=[], idx=0):
        badRegs = []
        badDst = []
        while True:


            load4 = self._find(Category.XCHG_REG, reg=what,  badDst=badDst, dontModify=dontModify, srcdst='src')
            if not load4:
                raise RopChainError('Cannot build loadwhere gadget!')
            else:
                mov = self._find(Category.LOAD_MEM, reg=load4.category[2]['dst'],  badDst=badDst, dontModify=[load4.category[2]['dst']]+dontModify, srcdst='dst')
                if not mov:
                    badDst.append(load4.category[2]['dst'])
                    continue

                popReg2 = self._find(Category.LOAD_REG, reg=mov.category[2]['src'], dontModify=[what,load4.category[2]['src']]+dontModify)
                if not popReg2:
                    badDst.append(load4.category[2]['src'])
                    continue
                else:
                    break;



        toReturn = self._printRopInstruction(popReg2, False)
        toReturn += self._printPaddingInstruction(toHex(from_reg,4))
        regs = self._paddingNeededFor(popReg2)
        for i in range(len(regs)):
            toReturn +=self._printPaddingInstruction()

        toReturn += self._printRopInstruction(mov)

        toReturn += self._printRopInstruction(load4)

        return (toReturn,what, popReg2.category[2]['dst'])

    def _createNumberSubtract(self, number, reg=None, badRegs=None, dontModify=None):
        if not badRegs:
            badRegs=[]
        while True:
            sub = self._find(Category.SUB_REG, reg=reg, badDst=badRegs, badSrc=badRegs, dontModify=dontModify)
            if not sub:
                raise RopChainError('Cannot build number with subtract gadget for reg %s!' % reg)
            popSrc = self._find(Category.LOAD_REG, reg=sub.category[2]['src'], dontModify=dontModify)
            if not popSrc:
                badRegs.append=[sub.category[2]['src']]
                continue
            popDst = self._find(Category.LOAD_REG, reg=sub.category[2]['dst'], dontModify=[sub.category[2]['src']]+dontModify)
            if not popDst:
                badRegs.append=[sub.category[2]['dst']]
                continue
            else:
                break;

        filler = self._createZeroByteFillerForSub(number)

        toReturn = self._printRopInstruction(popSrc, False)
        toReturn += self._printPaddingInstruction(toHex(filler,4))
        regs = self._paddingNeededFor(popSrc)
        for i in range(len(regs)):
            toReturn += self._printPaddingInstruction()
        toReturn += self._printRopInstruction(popDst, False)
        toReturn += self._printPaddingInstruction(toHex(filler+number,4))
        regs = self._paddingNeededFor(popDst)
        for i in range(len(regs)):
            toReturn += self._printPaddingInstruction()
        toReturn += self._printRopInstruction(sub)
        return (toReturn, popDst.category[2]['dst'],popSrc.category[2]['dst'])

    def _createNumberAddition(self, number, reg=None, badRegs=None, dontModify=None):
        if not badRegs:
            badRegs=[]
        while True:
            sub = self._find(Category.ADD_REG, reg=reg, badDst=badRegs, badSrc=badRegs, dontModify=dontModify)
            if not sub:
                raise RopChainError('Cannot build number with addition gadget for reg %s!' % reg)
            popSrc = self._find(Category.LOAD_REG, reg=sub.category[2]['src'], dontModify=dontModify)
            if not popSrc:
                badRegs.append=[sub.category[2]['src']]
                continue
            popDst = self._find(Category.LOAD_REG, reg=sub.category[2]['dst'], dontModify=[sub.category[2]['src']]+dontModify)
            if not popDst:
                badRegs.append(sub.category[2]['dst'])
                continue
            else:
                break;

        filler = self._createZeroByteFillerForAdd(number)

        toReturn = self._printRopInstruction(popSrc, False)
        toReturn += self._printPaddingInstruction(toHex(filler,4))
        regs = self._paddingNeededFor(popSrc)
        for i in range(len(regs)):
            toReturn += self._printPaddingInstruction()
        toReturn += self._printRopInstruction(popDst, False)
        toReturn += self._printPaddingInstruction(toHex(number - filler,4))
        regs = self._paddingNeededFor(popDst)
        for i in range(len(regs)):
            toReturn += self._printPaddingInstruction()
        toReturn += self._printRopInstruction(sub)

        return (toReturn, popDst.category[2]['dst'],popSrc.category[2]['dst'])

    def _createNumberPop(self, number, reg=None, badRegs=None, dontModify=None):
        if self._containsZeroByte(0xffffffff):
            raise RopChainError("Cannot write value with pop -1 and inc gadgets, because there are badbytes in the negated number")
        while True:
            popReg = self._find(Category.LOAD_REG, reg=reg, badDst=badRegs,dontModify=dontModify)
            if not popReg:
                raise RopChainError('Cannot build number with xor gadget!')
            incReg = self._find(Category.INC_REG, reg=popReg.category[2]['dst'], dontModify=dontModify)
            if not incReg:
                if not badRegs:
                    badRegs = []
                badRegs.append(popReg.category[2]['dst'])
            else:
                break

        toReturn = self._printRopInstruction(popReg)
        toReturn += self._printPaddingInstruction(toHex(0xffffffff,4))
        for i in range(number+1):
            toReturn += self._printRopInstruction(incReg)

        return (toReturn ,popReg.category[2]['dst'],)


    def _createNumberXOR(self, number, reg=None, badRegs=None, dontModify=None):
        while True:
            clearReg = self._find(Category.CLEAR_REG, reg=reg, badDst=badRegs, badSrc=badRegs,dontModify=dontModify, srcEqDst=True)
            if not clearReg:
                raise RopChainError('Cannot build number with xor gadget!')
            if number > 0:
                incReg = self._find(Category.INC_REG, reg=clearReg.category[2]['src'], dontModify=dontModify)
                if not incReg:
                    if not badRegs:
                        badRegs = []
                    badRegs.append(clearReg.category[2]['src'])
                else:
                    break
            else:
                break

        toReturn = self._printRopInstruction(clearReg)
        for i in range(number):
            toReturn += self._printRopInstruction(incReg)

        return (toReturn, clearReg.category[2]['dst'],)

    def _createNumberXchg(self, number, reg=None, badRegs=None, dontModify=None):
        xchg = self._find(Category.XCHG_REG, reg=reg, badDst=badRegs, dontModify=dontModify)
        if not xchg:
            raise RopChainError('Cannot build number gadget with xchg!')

        other = xchg.category[2]['src'] if xchg.category[2]['dst'] else xchg.category[2]['dst']

        toReturn = self._createNumber(number, other, badRegs, dontModify)[0]

        toReturn += self._printRopInstruction(xchg)
        return (toReturn, reg, other)

    def _createNumberNeg(self, number, reg=None, badRegs=None, dontModify=None):
        if number == 0:
            raise RopChainError('Cannot build number gadget with neg if number is 0!')
        if self._containsZeroByte((~number)+1):
            raise RopChainError("Cannot use neg gadget, because there are badbytes in the negated number")
        neg = self._find(Category.NEG_REG, reg=reg, badDst=badRegs, dontModify=dontModify)
        if not neg:
            raise RopChainError('Cannot build number gadget with neg!')

        pop = self._find(Category.LOAD_REG, reg=reg, badDst=badRegs, dontModify=dontModify)
        if not pop:
            raise RopChainError('Cannot build number gadget with neg!')

        toReturn = self._printRopInstruction(pop)


        toReturn += self._printPaddingInstruction(toHex((~number)+1)) # two's complement
        toReturn += self._printRopInstruction(neg)
        return (toReturn, reg,)

    def _createNumber(self, number, reg=None, badRegs=None, dontModify=None, xchg=True):
        try:
            if self._containsZeroByte(number):
                try:
                    return self._createNumberNeg(number, reg, badRegs,dontModify)
                except RopChainError as e:

                    if number < 0x50:
                        try:
                            return self._createNumberXOR(number, reg, badRegs,dontModify)
                        except RopChainError:
                            try:
                                return self._createNumberPop(number, reg, badRegs,dontModify)
                            except RopChainError:
                                try:
                                    return self._createNumberSubtract(number, reg, badRegs,dontModify)
                                except RopChainError:
                                    return self._createNumberAddition(number, reg, badRegs,dontModify)

                    else :
                        try:
                            return self._createNumberSubtract(number, reg, badRegs,dontModify)
                        except RopChainError:
                            return self._createNumberAddition(number, reg, badRegs,dontModify)
            else:
                popReg =self._find(Category.LOAD_REG, reg=reg, badDst=badRegs,dontModify=dontModify)
                if not popReg:
                    raise RopChainError('Cannot build number gadget!')
                toReturn = self._printRopInstruction(popReg, padding=True, number=toHex(number,4))
                return (toReturn , popReg.category[2]['dst'])
        except RopChainError as e:
            return self._createNumberXchg(number, reg, badRegs, dontModify)

    def _createAddress(self, address, reg=None, badRegs=None, dontModify=None):
        popReg = self._find(Category.LOAD_REG, reg=reg, badDst=badRegs,dontModify=dontModify)
        if not popReg:
            raise RopChainError('Cannot build address gadget!')

        toReturn = ''

        toReturn += self._printRopInstruction(popReg,False)
        toReturn += self._printRebasedAddress(toHex(address, 4), idx=self._usedBinaries.index((popReg.fileName, popReg.section)))
        regs = self._paddingNeededFor(popReg)
        for i in range(len(regs)):
            toReturn +=self._printPaddingInstruction()

        return (toReturn,popReg.category[2]['dst'])

    def _createSyscall(self, reg=None, badRegs=None, dontModify=None):
        syscall = self._find(Category.SYSCALL, reg=None, badDst=None, dontModify=dontModify)
        if not syscall:
            raise RopChainError('Cannot build syscall gadget!')

        toReturn = ''

        toReturn += self._printRopInstruction(syscall)

        return (toReturn,)

    def _createOpcode(self, opcode):

        return self._printRopInstruction(self._searchOpcode(opcode))


    def _searchOpcode(self, opcode):
        r = Ropper()
        gadgets = []
        for binary in self._binaries:
            for section in binary.executableSections:
                vaddr = section.virtualAddress
                gadgets.extend(r.searchOpcode(binary,opcode=opcode,disass=True))

        if len(gadgets) > 0:
            for gadget in gadgets:
                if not gadget:
                    continue
                if not self.containsBadbytes(gadget.IMAGE_BASES.get(gadget.fileName,0) + gadget.lines[0][0]):
                    if (gadget.fileName, gadget._section) not in self._usedBinaries:
                        self._usedBinaries.append((gadget.fileName, gadget._section))
                    return gadget
        else:
            raise RopChainError('Cannot create gadget for opcode: %s' % opcode)

    def create(self):
        pass


class RopChainX86System(RopChainX86):


    @classmethod
    def name(cls):
        return 'execve'

    def _createCommand(self, what, where, reg=None, dontModify=[], idx=0):
        if len(what) % 4 > 0:
            what = '/' * (4 - len(what) % 4) + what
        return self._createWriteStringWhere(what,where, idx=idx)

    def create(self, options={}):
        cmd = options.get('cmd')
        address = options.get('address')
        if not cmd:
            cmd = '/bin/sh'
        if len(cmd.split(' ')) > 1:
            raise RopChainError('No argument support for execve commands')

        self._printMessage('ROPchain Generator for syscall execve:\n')
        self._printMessage('\nwrite command into data section\neax 0xb\nebx address to cmd\necx address to null\nedx address to null\n')
        chain = self._printHeader()
        gadgets = []
        can_create_command = False
        chain_tmp = '\n'
        if address is None:
            section = self._binaries[0].getSection('.data')
            length = math.ceil(float(len(cmd))/4) * 4
            
            try:
                chain_tmp += self._createCommand(cmd,section.offset+0x1000)[0]
                can_create_command = True
                
            except RopChainError as e:
                self._printMessage('Cannot create gadget: writewhatwhere')
                self._printMessage('Use 0x41414141 as command address. Please replace that value.')
                address = 0x41414141
            if can_create_command:
                
                badregs = []
                while True:

                    ret = self._createNumber(0x0, badRegs=badregs)
                    chain_tmp += ret[0]
                    try:
                        chain_tmp += self._createWriteRegValueWhere(ret[1], section.offset+0x1000+length)[0]
                        break
                    except BaseException as e:
                        #raise e
                        badregs.append(ret[1])

                gadgets.append((self._createAddress, [section.offset+0x1000],{'reg':'ebx'},['ebx', 'bx', 'bl', 'bh']))
                gadgets.append((self._createAddress, [section.offset+0x1000+length],{'reg':'ecx'},['ecx', 'cx', 'cl', 'ch']))
                gadgets.append((self._createAddress, [section.offset+0x1000+length],{'reg':'edx'},['edx', 'dx', 'dl', 'dh']))
                gadgets.append((self._createNumber, [0xb],{'reg':'eax'},['eax', 'ax', 'al', 'ah']))
        if address is not None and not can_create_command:
            if type(address) is str:
                address = int(address, 16)
            nulladdress = options.get('nulladdress')
            if nulladdress is None:
                self._printMessage('No address to a null bytes was given, 0x42424242 is used instead.')
                self._printMessage('Please replace that value.')
                nulladdress = 0x42424242
            elif type(nulladdress) is str:
                nulladdress = int(nulladdress,16)
            gadgets.append((self._createNumber, [address],{'reg':'ebx'},['ebx', 'bx', 'bl', 'bh']))
            gadgets.append((self._createNumber, [nulladdress],{'reg':'ecx'},['ecx', 'cx', 'cl', 'ch']))
            gadgets.append((self._createNumber, [nulladdress],{'reg':'edx'},['edx', 'dx', 'dl', 'dh']))
            gadgets.append((self._createNumber, [0xb],{'reg':'eax'},['eax', 'ax', 'al', 'ah']))

        self._printMessage('Try to create chain which fills registers without delete content of previous filled registers')
        chain_tmp += self._createDependenceChain(gadgets)
        try:
            self._printMessage('Look for syscall gadget')
            chain_tmp += self._createSyscall()[0]
            self._printMessage('syscall gadget found')

        except RopChainError:
            try:
                self._printMessage('No syscall gadget found!')
                self._printMessage('Look for int 0x80 opcode')

                chain_tmp += self._createOpcode('cd80')
                self._printMessage('int 0x80 opcode found')

            except RopChainError:
                try:
                    self._printMessage('No int 0x80 opcode found')
                    self._printMessage('Look for call gs:[0x10] opcode')
                    chain_tmp += self._createOpcode('65ff1510000000')
                    self._printMessage('call gs:[0x10] found')
                except RopChainError:
                    chain_tmp += '# INSERT SYSCALL GADGET HERE\n'

                    self._printMessage('No call gs:[0x10] opcode found')


        chain += self._printRebase()
        chain += 'rop = \'\'\n'

        chain += chain_tmp
        chain += 'print rop'
        return chain


class RopChainX86Mprotect(RopChainX86):
    """
    Builds a ropchain for mprotect syscall
    eax 0x7b
    ebx address
    ecx size
    edx 0x7 -> RWE
    """

    @classmethod
    def name(cls):
        return 'mprotect'

    def _createJmp(self, reg=['esp']):
        r = Ropper()
        gadgets = []
        for binary in self._binaries:
            for section in binary.executableSections:
                vaddr = section.virtualAddress
                gadgets.extend(
                    r.searchJmpReg(binary,reg))



        if len(gadgets) > 0:
            if (gadgets[0].fileName, gadgets[0].section) not in self._usedBinaries:
                self._usedBinaries.append((gadgets[0].fileName, gadgets[0].section))
            return self._printRopInstruction(gadgets[0])
        else:
            return None


    def create(self, options={}):
        address = options.get('address')
        size = options.get('size')
        if not address:
            raise RopChainError('Missing parameter: address')
        if not size:
            raise RopChainError('Missing parameter: size')

        if not match('0x[0-9a-fA-F]{1,8}', address):
            raise RopChainError('Parameter address have to have the following format: <hexnumber>')

        if not match('0x[0-9a-fA-F]+', size):
            raise RopChainError('Parameter size have to have the following format: <hexnumber>')

        address = int(address, 16)
        size = int(size, 16)

        self._printMessage('ROPchain Generator for syscall mprotect:\n')
        self._printMessage('eax 0x7b\nebx address\necx size\nedx 0x7 -> RWE\n')

        chain = self._printHeader()

        chain += 'shellcode = \'\\xcc\'*100\n\n'

        gadgets = []
        gadgets.append((self._createNumber, [address],{'reg':'ebx'},['ebx', 'bx', 'bl', 'bh']))
        gadgets.append((self._createNumber, [size],{'reg':'ecx'},['ecx', 'cx', 'cl', 'ch']))
        gadgets.append((self._createNumber, [0x7],{'reg':'edx'},['edx', 'dx', 'dl', 'dh']))
        gadgets.append((self._createNumber, [0x7d],{'reg':'eax'},['eax', 'ax', 'al', 'ah']))

        self._printMessage('Try to create chain which fills registers without delete content of previous filled registers')
        chain_tmp = ''
        chain_tmp += self._createDependenceChain(gadgets)
        try:
            self._printMessage('Look for syscall gadget')
            chain_tmp += self._createSyscall()[0]
            self._printMessage('syscall gadget found')
        except RopChainError:
            chain_tmp += '\n# ADD HERE SYSCALL GADGET\n\n'
            self._printMessage('No syscall gadget found!')

        self._printMessage('Look for jmp esp')
        jmp_esp = self._createJmp()
        if jmp_esp:
            self._printMessage('jmp esp found')
            chain_tmp += jmp_esp
        else:
            self._printMessage('no jmp esp found')
            chain_tmp += '\n# ADD HERE JMP ESP\n\n'

        chain += self._printRebase()
        chain += '\nrop = \'\'\n'
        chain += chain_tmp
        chain += 'rop += shellcode\n\n'
        chain += 'print(rop)\n'

        return chain

class RopChainX86VirtualProtect(RopChainX86):
    """
    Builds a ropchain for a VirtualProtect call using pushad
    eax 0x90909090
    ecx old protection (writable addr)
    edx 0x40 (RWE)
    ebx size
    esp address
    ebp return address (jmp esp)
    esi pointer to VirtualProtect
    edi ret (rop nop)
    """


    @classmethod
    def name(cls):
        return 'virtualprotect'


    def _createPushad(self):
        pushad = self._find(Category.PUSHAD)
        if pushad:
            return self._printRopInstruction(pushad)
        else:
            self._printMessage('No pushad found!')
            return '# Add here PUSHAD gadget!'


    def _createJmp(self, reg=['esp']):
        r = Ropper()
        gadgets = []
        for binary in self._binaries:
            for section in binary.executableSections:
                vaddr = section.offset
                gadgets.extend(
                    r.searchJmpReg(binary,reg))



        if len(gadgets) > 0:
            if (gadgets[0].fileName, gadgets[0].section) not in self._usedBinaries:
                self._usedBinaries.append((gadgets[0].fileName, gadgets[0].section))
            return gadgets[0]
        else:
            return None

    def __extract(self, param):
        if (not match('0x[0-9a-fA-F]{1,8},0x[0-9a-fA-F]+', param)) and (not match('0x[0-9a-fA-F]+', param)):
            raise RopChainError('Parameter have to have the following format: <hexnumber>')

        return (None, int(param, 16))

    def __getVirtualProtectEntry(self):
        for binary in self._binaries:
            if binary.type == Type.PE:
                imports = binary._binary.dataDirectory[ImageDirectoryEntry.IMPORT]
                if not imports:
                    return None
                for descriptorData in imports:
                    for thunk in descriptorData.importAddressTable:

                        if thunk.importByName and thunk.importByName.name == 'VirtualProtect':
                            return thunk.rva, binary.imageBase
            else:
                self._printMessage('File is not a PE file.')
        return None



    def create(self, options={}):

        self._printMessage('Ropchain Generator for VirtualProtect:\n')
        self._printMessage('eax 0x90909090\necx old protection (writable addr)\nedx 0x40 (RWE)\nebx size\nesp address\nebp return address (jmp esp)\nesi pointer to VirtualProtect\nedi ret (rop nop)\n')
        
        image_base = 0
        address = options.get('address')      
        given = False
        if not address:
            address, image_base = self.__getVirtualProtectEntry()
            if not address:
                self._printMessage('No IAT-Entry for VirtualProtect found!')
                raise RopChainError('No IAT-Entry for VirtualProtect found and no address is given')
        else:
            if address:
                if not match('0x[0-9a-fA-F]{1,8}', address):
                    raise RopChainError('Parameter address have to have the following format: <hexnumber>')
                address = int(address, 16)
            given = True


        writeable_ptr = self._binaries[0].getWriteableSection().offset

        for i in range(0,0x10000,4):
            if not self.containsBadbytes((writeable_ptr + i) & 0xffff,2):
                writeable_ptr += i
                break
        jmp_esp = self._createJmp()
        ret_addr = self._searchOpcode('c3')

        chain = self._printHeader()

        chain += '\n\nshellcode = \'\\xcc\'*100\n\n'
        gadgets = []
        to_extend = []
        chain_tmp = ''
        got_jmp_esp = False
        try:
            self._printMessage('Try to create gadget to fill esi with content of IAT address: 0x%x' % (address + image_base))
            chain_tmp += self._createLoadRegValueFrom('esi', address+image_base)[0]
            gadgets.append((self._createNumber, [0x90909090],{'reg':'eax'},['eax', 'ax', 'ah', 'al','esi','si']))

            to_extend = ['esi','si']
            if jmp_esp:
                gadgets.append((self._createAddress, [jmp_esp.lines[0][0]],{'reg':'ebp'},['ebp', 'bp']+to_extend))
            got_jmp_esp = True
        except RopChainError:
            self._printMessage('Cannot create fill esi gadget!')
            self._printMessage('Try to create this chain:\n')
            self._printMessage('eax Pointer to VirtualProtect\necx old protection (writable addr)\nedx 0x40 (RWE)\nebx size\nesp address\nebp return address (pop ebp;ret)\nesi pointer to jmp [eax]\nedi ret (rop nop)\n')

            jmp_eax = self._searchOpcode('ff20') # jmp [eax]
            gadgets.append((self._createAddress, [jmp_eax.lines[0][0]],{'reg':'esi'},['esi','si']))
            
            gadgets.append((self._createNumber, [address],{'reg':'eax'},['eax', 'ax', 'ah', 'al']))
            
            pop_ebp = self._searchOpcode('5dc3')
            if pop_ebp:
                gadgets.append((self._createAddress, [pop_ebp.lines[0][0]],{'reg':'ebp'},['ebp', 'bp']+to_extend))


        gadgets.append((self._createNumber, [0x1],{'reg':'ebx'},['ebx', 'bx', 'bl', 'bh']+to_extend))
        gadgets.append((self._createAddress, [writeable_ptr],{'reg':'ecx'},['ecx', 'cx', 'cl', 'ch']+to_extend))

        gadgets.append((self._createNumber, [0x40],{'reg':'edx'},['edx', 'dx', 'dh', 'dl']+to_extend))

        gadgets.append((self._createAddress, [ret_addr.lines[0][0]],{'reg':'edi'},['edi', 'di']+to_extend))

        self._printMessage('Try to create chain which fills registers without delete content of previous filled registers')
        chain_tmp += self._createDependenceChain(gadgets)

        self._printMessage('Look for pushad gadget')
        chain_tmp += self._createPushad()

        if not got_jmp_esp and jmp_esp:
            chain_tmp += self._printRopInstruction(jmp_esp)


        chain += self._printRebase()
        chain += 'rop = \'\'\n'
        chain += chain_tmp
        chain += 'rop += shellcode\n\n'
        chain += 'print(rop)\n'

        return chain

# class RopChainX86VirtualAlloc(RopChainX86):
#     """
#     Builds a ropchain for a VirtualProtect call using pushad
#     eax 0x90909090
#     ecx old protection (writable addr)
#     edx 0x40 (RWE)
#     ebx size
#     esp address
#     ebp return address (jmp esp)
#     esi pointer to VirtualProtect
#     edi ret (rop nop)
#     """


#     @classmethod
#     def name(cls):
#         return 'virtualalloc'


#     def _createPushad(self):
#         pushad = self._find(Category.PUSHAD)
#         if pushad:
#             return self._printRopInstruction(pushad)
#         else:
#             self._printer.printInfo('No pushad found!')
#             return '# Add here PUSHAD gadget!'


#     def _createJmp(self, reg=['esp']):
#         r = Ropper()
#         gadgets = []
#         for binary in self._binaries:
#             for section in binary.executableSections:
#                 vaddr = section.offset
#                 gadgets.extend(
#                     r.searchJmpReg(self._binaries[0],reg))



#         if len(gadgets) > 0:
#             if (gadgets[0]._binary, gadgets[0]._section) not in self._usedBinaries:
#                 self._usedBinaries.append((gadgets[0]._binary, gadgets[0]._section))
#             return gadgets[0]
#         else:
#             return None

#     def __extract(self, param):
#         if (not match('0x[0-9a-fA-F]{1,8},0x[0-9a-fA-F]+', param)) and (not match('0x[0-9a-fA-F]+', param)):
#             raise RopChainError('Parameter have to have the following format: <hexnumber>,<hexnumber> or <hexnumber>')

#         split = param.split(',')
#         if len(split) == 2:
#             if isHex(split[1]):
#                 return (int(split[0], 16), int(split[1], 16))
#         else:
#             return (None, int(split[0], 16))

#     def __getVirtualProtectEntry(self):
#         for binary in self._binaries:
#             if binary.type == Type.PE:
#                 s = binary._binary.dataDirectory[ImageDirectoryEntry.IMPORT]
#                 for thunk in s.importNameTable:
#                     if thunk.importByName.name == 'VirtualAlloc':
#                             return thunk.rva + binary.imageBase
#             else:
#                 self._printer.printError('File is not a PE file.')
#         return None



#     def create(self, param=None):
#         if not param:
#             raise RopChainError('Missing parameter: address,size or size')

#         self._printer.printInfo('Ropchain Generator for VirtualProtect:\n')
#         self._printer.println('eax 0x90909090\necx old protection (writable addr)\nedx 0x40 (RWE)\nebx size\nesp address\nebp return address (jmp esp)\nesi pointer to VirtualProtect\nedi ret (rop nop)\n')

#         address, size = self.__extract(param)
#         given = False
#         if not address:
#             address = self.__getVirtualProtectEntry()
#             if not address:
#                 self._printer.printError('No IAT-Entry for VirtualProtect found!')
#                 raise RopChainError('No IAT-Entry for VirtualProtect found and no address is given')
#         else:
#             given = True

#         jmp_esp = self._createJmp()
#         ret_addr = self._searchOpcode('c3')

#         chain = self._printHeader()

#         chain += '\n\nshellcode = \'\\xcc\'*100\n\n'
#         gadgets = []
#         to_extend = []
#         chain_tmp = ''
#         try:
#             self._printer.printInfo('Try to create gadget to fill esi with content of IAT address: %s' % address)
#             chain_tmp += self._createLoadRegValueFrom('esi', address)[0]
#             if given:
#                 gadgets.append((self._createNumber, [address],{'reg':'eax'},['eax', 'ax', 'ah', 'al','esi','si']))
#             else:
#                 gadgets.append((self._createAddress, [address],{'reg':'eax'},['eax', 'ax', 'ah', 'al','esi','si']))
#             to_extend = ['esi','si']
#         except:
#             self._printer.printInfo('Cannot create fill esi gadget!')
#             self._printer.printInfo('Try to create this chain:\n')
#             self._printer.println('eax Pointer to VirtualProtect\necx old protection (writable addr)\nedx 0x40 (RWE)\nebx size\nesp address\nebp return address (jmp esp)\nesi pointer to jmp [eax]\nedi ret (rop nop)\n')

#             jmp_eax = self._searchOpcode('ff20') # jmp [eax]
#             gadgets.append((self._createAddress, [jmp_eax.lines[0][0]],{'reg':'esi'},['esi','si']))
#             if given:
#                 gadgets.append((self._createNumber, [address],{'reg':'eax'},['eax', 'ax', 'ah', 'al']))
#             else:
#                 gadgets.append((self._createAddress, [address],{'reg':'eax'},['eax', 'ax', 'ah', 'al']))



#         gadgets.append((self._createNumber, [size],{'reg':'ebx'},['ebx', 'bx', 'bl', 'bh']+to_extend))
#         gadgets.append((self._createNumber, [0x40],{'reg':'ecx'},['ecx', 'cx', 'cl', 'ch']+to_extend))
#         if jmp_esp:
#             gadgets.append((self._createAddress, [jmp_esp.lines[0][0]],{'reg':'ebp'},['ebp', 'bp']+to_extend))
#         gadgets.append((self._createNumber, [0x1000],{'reg':'edx'},['edx', 'dx', 'dh', 'dl']+to_extend))

#         gadgets.append((self._createAddress, [ret_addr.lines[0][0]],{'reg':'edi'},['edi', 'di']+to_extend))

#         self._printer.printInfo('Try to create chain which fills registers without delete content of previous filled registers')
#         chain_tmp += self._createDependenceChain(gadgets)

#         self._printer.printInfo('Look for pushad gadget')
#         chain_tmp += self._createPushad()


#         chain += self._printRebase()
#         chain += 'rop = \'\'\n'
#         chain += chain_tmp
#         chain += 'rop += shellcode\n\n'
#         chain += 'print(rop)\n'

#         return chain
