# -*- coding:utf-8 -*- 
# Utils module: useful functions to build exploits
from ropgenerator.semantic.Engine import search, LMAX
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 

POP_MULTIPLE_LMAX = 6000

def popMultiple(args, constraint=None, assertion=None, clmax=None, optimizeLen=False):
    """
    args is a list of pairs (reg, value) 
        OR a list of triples (reg, value, comment)
    reg is a reg UID
    value is an int
    Creates a chain that pops values into regs
    """
    if( clmax is None ):
        clmax = POP_MULTIPLE_LMAX
    elif( 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( len(arg) == 3 ):
                comment = arg[2]
            else:
                comment = None
            if( optimizeLen ):
                pop = search(QueryType.CSTtoREG, arg[0], arg[1], constr_tmp, a, n=1, clmax=clmax_tmp, CSTtoREG_comment=comment, optimizeLen=True)
            else:
                pop = search(QueryType.CSTtoREG, arg[0], arg[1], constr_tmp, a, n=1, clmax=clmax_tmp, CSTtoREG_comment=comment)
            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



