# -*- coding:utf-8 -*- 
# Utils module: useful functions to build exploits (more advanced features
# than Utils.py)

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
from ropgenerator.exploit.Call import build_call 
import itertools
import ropgenerator.Architecture as Arch 


#### Put strings in memory

STR_TO_MEM_LMAX = 6000

def STRtoMEM(string, address, constraint, assertion, limit=None, lmax=STR_TO_MEM_LMAX\
        ,addr_str=None, hex_info=False, optimizeLen=False):
    """
    Put a string into memory 
    
    limit : if (int) then the max address where to write in memory 
            if None then string should be written ONLY at 'addr' (no adjust possible)
            
    return value: 
            if limit is not None then a pair (address, ROPChain) or (None, None)
            if limit is None then a ROPChain or None 
        
    """    
    (addr,chain) = STRtoMEM_memcpy(string, address, constraint, assertion, limit, lmax, addr_str, hex_info)
    res = (addr,chain)
    if( (not chain) or optimizeLen ):
        (addr2,chain2) = STRtoMEM_strcpy(string, address, constraint, assertion, limit, lmax, addr_str, hex_info)
        if( not res[1] ):
            res = (addr2,chain2)
        elif (chain2 and ( len(chain2) < len(res[1]))):
                res = (addr2,chain2)
    # If possible adjust, return address and chain
    if( limit is None  ):
        return res[1]
    # Otherwise return only chain
    else:
        return res

def STRtoMEM_strcpy(string, addr, constraint, assertion, limit=None, 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 = hex(addr)
    
    # Getting strcpy function 
    (func_name, func_addr ) = getFunctionAddress('strcpy')
    if( not func_addr ):
        verbose('Could not find strcpy function')
        return (None,None)
    elif( not constraint.badBytes.verifyAddress(func_addr)):
        verbose("strcpy address ({}) contains bad bytes".format(hex(func_addr)))
        return (None,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,None)

    # Find delivery address
    substr_lengthes = [len(substr[1])-1 for substr in substrings_addr]# -1 becasue strcpy 
    substr_lengthes[-1] += 1
    if( not limit is None ):
        custom_stack = find_closest_base_fake_stack_address(addr, limit, substr_lengthes, constraint)
        if( custom_stack is None ):
            verbose("Couldn't write string in memory because of bad bytes")
            return (None,None)
    else:
        custom_stack = find_closest_base_fake_stack_address(addr, addr+sum(substr_lengthes), substr_lengthes, constraint)
        if( custom_stack is None ):
            verbose("Couldn't write string in memory because of bad bytes")
            return (None,None)
    if( custom_stack != addr ):
        addr_str = hex(custom_stack)

    # Build chain 
    res = ROPChain()
    offset = 0
    saved_custom_stack = custom_stack
    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
        commentStack="Arg2: " + string_ropg("{} + {}".format(addr_str, offset))
        commentSubStr="Arg1: " + string_ropg(substring_info)
        func_call = build_call(func_name, [substring_addr, custom_stack], constraint, assertion, [commentSubStr, commentStack], optimizeLen=True)
        if( isinstance(func_call, str)):
            verbose("strcpy: " + func_call)
            return (None,None)
        else:
            res.addChain(func_call)
            if( len(res) > lmax ):
                return (None,None)
        # Adjust
        # -1 Because strcpy has a null byte :/
        # Except when we INTEND to write a null byte 
        if( substring_str == '\x00' ):
            dec = 0
        else:
            dec = 1
        custom_stack = custom_stack + len(substring_str) -dec
        offset = offset + len(substring_str) - dec

    return (saved_custom_stack, res)

def STRtoMEM_memcpy(string, addr, constraint, assertion, limit=None, 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 = hex(addr)
    
    # Getting strcpy function 
    (func_name, func_addr ) = getFunctionAddress('memcpy')
    if( not func_addr ):
        verbose('Could not find memcpy function')
        return (None,None)
    elif( not constraint.badBytes.verifyAddress(func_addr)):
        verbose("memcpy address ({}) contains bad bytes".format(hex(func_addr)))
        return (None,None)
    
    # We decompose the string in substrings to be copied
    substrings_addr = findBytes(string, badBytes = constraint.getBadBytes())
    if( not substrings_addr ):
        return (None,None)

    # Find delivery address
    substr_lengthes = [len(substr[1]) for substr in substrings_addr]
    if( not limit is None ):
        custom_stack = find_closest_base_fake_stack_address(addr, limit, substr_lengthes, constraint)
        if( custom_stack is None ):
            verbose("Couldn't write string in memory because of bad bytes")
            return (None,None)
    else:
        custom_stack = find_closest_base_fake_stack_address(addr, addr+sum(substr_lengthes), substr_lengthes, constraint)
        if( custom_stack is None ):
            verbose("Couldn't write string in memory because of bad bytes")
            return (None,None)
    if( custom_stack != addr ):
        addr_str = hex(custom_stack)
        
    # Build chain 
    res = ROPChain()
    offset = 0 
    saved_custom_stack = custom_stack
    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+"'"
        comment3 ="Arg3: " + string_ropg(str(len(substring_str)))
        comment2 ="Arg2: " + string_ropg(substring_info)
        comment1 ="Arg1: " + string_ropg("{} + {}".format(addr_str, offset))
        
        func_call = build_call(func_name, [custom_stack, substring_addr, len(substring_str)],\
                    constraint, assertion, [comment1, comment2, comment3], optimizeLen=True)
        
        if( isinstance(func_call, str) ):
            verbose("memcpy: " + func_call)
            return (None,None)
    
        res.addChain(func_call)
        if( len(res) > lmax ):
            return (None,None)
    
        # Adjust
        custom_stack = custom_stack + len(substring_str)
        offset = offset + len(substring_str)

    return (saved_custom_stack,res)

# Util function
def find_closest_base_fake_stack_address(base, limit, substr_lengthes, constraint):
    """
    When writing substrings, a bad address might occur for some of them 
    BASE <- SUB1
    BASE + LEN(SUB1) <- SUB2
    BASE + LEN(SUB1) + LEN(SUB2) -- BAD BYTE IN IT !! :O 
    
    So find another base address that works in the range [lower_addr..upper_addr]
    """
    # Compute the list of addresses that will be used with base and the substring lengthes 
    def get_addr_list(base, substr_lengthes):
        inc = 0
        res = [base]
        for l in substr_lengthes[:-1]: # Don't take the last one because we don't write after 
            inc += l
            res.append(base + inc)
        return res
            
    address = base
    total_length = sum(substr_lengthes)
    while(address + total_length <= limit):
        addr_list = get_addr_list(address, substr_lengthes)
        for addr in addr_list:
            index = constraint.badBytes.findIndex(addr)
            if( index >= 0 ):
                # Bad byte found 
                # If we tried everything for this byte return 
                if( (address & (0xff << index*8)) == (0xff << index*8) ):
                    return None
                # Other wise add 1 and retry 
                address += (0x1 << index*8)
                break
        # No bad bytes found in addresses, return result :) 
        # Else we keep looping 
        if( index == -1 ):
            return address
            
    # We reached upper limit to write without finding a good address
    return None
