# -*- coding:utf-8 -*- 
# Utils module: useful functions to build exploits
from ropgenerator.semantic.Engine import search, _basic, LMAX, search_optimize_len
from ropgenerator.Constraints import Constraint, RegsNotModified, Assertion, Chainable, StackPointerIncrement
from ropgenerator.semantic.ROPChains import ROPChain
from ropgenerator.Database import QueryType
from ropgenerator.exploit.Scanner import getFunctionAddress, findBytes
from ropgenerator.IO import verbose, string_bold, string_special, string_ropg
import itertools
import ropgenerator.Architecture as Arch 

#### Pop values into registers 

def popMultiple(args, constraint=None, assertion=None, clmax=None, optimizeLen=False):
    """
    args is a list of pairs (reg, value)
    reg is a reg UID
    value is an int
    Creates a chain that pops values into regs
    """
    if( clmax <= 0 ):
        return None
    
    if( constraint is None ):
        constr = Constraint()
    else:
        constr = constraint
    if( assertion is None ):
        a = Assertion()
    else:
        a = assertion
        
    perms = itertools.permutations(args)
    for perm in perms:
        clmax_tmp = clmax 
        res = ROPChain()
        constr_tmp = constr
        for arg in perm:
            if( optimizeLen ):
                pop = search_optimize_len(QueryType.CSTtoREG, arg[0], arg[1], constr_tmp, a, n=1, clmax=clmax_tmp)
            else:
                pop = search(QueryType.CSTtoREG, arg[0], arg[1], constr_tmp, a, n=1, clmax=clmax_tmp)
            if( not pop ):
                break
            else:
                clmax_tmp -= len(pop[0])
                # If Reached max length, exit  
                if( clmax_tmp < 0 ):
                    pop = None
                    break
                else:
                    res.addChain(pop[0])
                    constr_tmp = constr_tmp.add(RegsNotModified([arg[0]]))
        if( pop ):
            return res
    return None


#### Put strings in memory

STR_TO_MEM_LMAX = 6000

def STRtoMEM(string, addr, constraint, assertion, lmax=STR_TO_MEM_LMAX ,addr_str=None, hex_info=False, optimizeLen=False):
    """
    Put a string into memory 
    """
    res = STRtoMEM_memcpy(string, addr, constraint, assertion, lmax, addr_str, hex_info)
    if( (not res) or optimizeLen ):
        res2 = STRtoMEM_strcpy(string, addr, constraint, assertion, lmax, addr_str, hex_info)
        if( not res ):
            res = res2
        elif (res2 and ( len(res2) < len(res))):
                res = res2
    return res 

def STRtoMEM_strcpy(string, addr, constraint, assertion, lmax=STR_TO_MEM_LMAX , addr_str=None, hex_info=False):
    """
    STRCPY STRATEGY
    Copy the string using strcpy function 
    """
    if( not addr_str ):
        addr_str = "0x"+format(addr, '0'+str(Arch.octets()*2)+'x')
    
    # Getting strcpy function 
    (func_name, func_addr ) = getFunctionAddress('strcpy')
    if( not func_addr ):
        verbose('Could not find strcpy function')
        return None
    elif( not constraint.badBytes.verifyAddress(func_addr)):
        verbose("strcpy address ({}) contains bad bytes".format(hex(func_addr)))
        return None
    
    # We decompose the string in substrings to be copied
    substrings_addr = findBytes(string, badBytes = constraint.getBadBytes(), add_null=True)
    if( not substrings_addr ):
        return None
    elif( len(substrings_addr)*4 > lmax ):
        verbose("Strcpy: ROP-Chain too long (length: {}, available bytes: {}) ".format(len(substrings_addr)*4*Arch.octets(), lmax*Arch.octets()))
        return None
        
    # Get a pop-pop-ret gadget 
    ppr_chains = _basic(QueryType.MEMtoREG, Arch.ipNum(), [Arch.spNum(), 2*Arch.octets()], constraint.add(StackPointerIncrement(3*Arch.octets())), assertion, clmax=1, noPadding=True)
    if( not ppr_chains ):
        verbose("Strcpy: Could not find suitable pop-pop-ret gadget")
        return None
    ppr_gadget = ppr_chains[0].chain[0] # Get the first gadget 

    # Build chain 
    res = ROPChain()
    offset = 0
    custom_stack = addr
    for (substring_addr,substring_str) in substrings_addr:
        if( hex_info ):
            substring_info = '\\x'+'\\x'.join(["%02x"%ord(c) for c in substring_str])
        else:
            substring_info = substring_str
        
        res.addPadding(func_addr, comment=string_ropg(func_name))
        res.addGadget(ppr_gadget)
        res.addPadding(substring_addr, comment="Arg1: " + string_ropg(substring_info))
        res.addPadding(custom_stack, comment="Arg2: " + string_ropg("{} + {}".format(addr_str, offset)))
    
        # Adjust
        custom_stack = custom_stack + len(substring_str)
        offset = offset + len(substring_str) - 1 # Because strcpy has a null byte :/ 

    return res

def STRtoMEM_memcpy(string, addr, constraint, assertion, lmax=STR_TO_MEM_LMAX , addr_str=None, hex_info=False):
    """
    MEMCPY STRATEGY
    Copy the string using memcpy function 
    """
    if( not addr_str ):
        addr_str = "0x"+format(addr, '0'+str(Arch.octets()*2)+'x')
    
    # Getting strcpy function 
    (func_name, func_addr ) = getFunctionAddress('memcpy')
    if( not func_addr ):
        verbose('Could not find memcpy function')
        return None
    elif( not constraint.badBytes.verifyAddress(func_addr)):
        verbose("memcpy address ({}) contains bad bytes".format(hex(func_addr)))
        return None
    
    # We decompose the string in substrings to be copied
    substrings_addr = findBytes(string, badBytes = constraint.getBadBytes())
    if( not substrings_addr ):
        return None
    elif( len(substrings_addr)*5 > lmax ):
        verbose("Memcpy: ROP-Chain too long (length: {}, available bytes: {}) ".format(len(substrings_addr)*5*Arch.octets(), lmax*Arch.octets()))
        return None
        
    # Get a pop-pop-pop-ret gadget 
    pppr_chains = _basic(QueryType.MEMtoREG, Arch.ipNum(), [Arch.spNum(), 3*Arch.octets()], constraint.add(StackPointerIncrement(4*Arch.octets())), assertion, clmax=1, noPadding=True)
    if( not pppr_chains ):
        verbose("Memcpy: Could not find suitable pop-pop-pop-ret gadget")
        return None
    pppr_gadget = pppr_chains[0].chain[0] # Get the first gadget 

    # Build chain 
    res = ROPChain()
    offset = 0 
    custom_stack = addr
    for (substring_addr,substring_str) in substrings_addr:
        if( hex_info ):
            substring_info = "'"+'\\x'+'\\x'.join(["%02x"%ord(c) for c in substring_str])+"'"
        else:
            substring_info = "'"+substring_str+"'"
        
        res.addPadding(func_addr, comment=string_ropg(func_name))
        res.addGadget(pppr_gadget)
        res.addPadding(len(substring_str), comment="Arg3: " + string_ropg(str(len(substring_str))))
        res.addPadding(substring_addr, comment="Arg2: " + string_ropg(substring_info))
        res.addPadding(custom_stack, comment="Arg1: " + string_ropg("{} + {}".format(addr_str, offset)))
    
        # Adjust
        custom_stack = custom_stack + len(substring_str)
        offset = offset + len(substring_str)

    return res
