# -*- coding:utf-8 -*- 
# Call module: call functions with ropchains 

from ropgenerator.IO import string_special, string_bold, banner, error, string_ropg
from ropgenerator.Constraints import Constraint, Assertion, BadBytes, RegsNotModified
from ropgenerator.Load import loadedBinary
from ropgenerator.exploit.Scanner import getFunctionAddress, getAllFunctions
from ropgenerator.semantic.ROPChains import ROPChain, validAddrStr
from ropgenerator.semantic.Engine import search 
from ropgenerator.Database import QueryType
from ropgenerator.exploit.Utils import popMultiple
import ropgenerator.Architecture as Arch

###################
#  CALL COMMAND   # 
###################

OPTION_OUTPUT = '--output-format'
OPTION_OUTPUT_SHORT = '-f'
# Options for output
OUTPUT_CONSOLE = 'console'
OUTPUT_PYTHON = 'python'
OUTPUT_RAW = 'raw'
OUTPUT = None # The one choosen 

OPTION_BAD_BYTES = '--bad-bytes'
OPTION_BAD_BYTES_SHORT = '-b'

OPTION_CALL = "--call"
OPTION_CALL_SHORT = "-c"

OPTION_KEEP_REGS = '--keep-regs'
OPTION_KEEP_REGS_SHORT = '-k'

OPTION_LIST = "--list"
OPTION_LIST_SHORT = "-l"

OPTION_HELP = "--help"
OPTION_HELP_SHORT = "-h"

OPTION_SHORTEST = '--shortest'
OPTION_SHORTEST_SHORT = '-s'

OPTION_LMAX = '--max-length'
OPTION_LMAX_SHORT = '-m'


CMD_CALL_HELP =  banner([string_bold("'call' command"),\
                    string_special("(Call functions with ROPChains)")])
CMD_CALL_HELP += "\n\n\t"+string_bold("Usage:")+\
"\n\t\tcall [OPTIONS]"
CMD_CALL_HELP += "\n\n\t"+string_bold("Options")+":"
CMD_CALL_HELP += "\n\t\t"+string_special(OPTION_CALL_SHORT)+","+\
    string_special(OPTION_CALL)+" <function>\t Call a function"
CMD_CALL_HELP += "\n\n\t\t"+string_special(OPTION_BAD_BYTES_SHORT)+","+string_special(OPTION_BAD_BYTES)+" <bytes>\t Bad bytes for payload.\n\t\t\t\t\t Expected format is a list of bytes \n\t\t\t\t\t separated by comas (e.g '-b 0A,0B,2F')"
CMD_CALL_HELP += "\n\n\t\t"+string_special(OPTION_KEEP_REGS_SHORT)+","+string_special(OPTION_KEEP_REGS)+" <regs>\t Registers that shouldn't be modified.\n\t\t\t\t\t Expected format is a list of registers \n\t\t\t\t\t separated by comas (e.g '-k edi,eax')"
CMD_CALL_HELP += "\n\n\t\t"+string_special(OPTION_LMAX_SHORT)+","+string_special(OPTION_LMAX)+" <int>\t Max length of the ROPChain in bytes"
CMD_CALL_HELP += "\n\n\t\t"+string_special(OPTION_SHORTEST_SHORT)+","+string_special(OPTION_SHORTEST)+"\t\t Find the shortest matching ROP-Chains"
CMD_CALL_HELP += "\n\n\t\t"+string_special(OPTION_LIST_SHORT)+","+\
    string_special(OPTION_LIST)+"\t\t List available functions"
CMD_CALL_HELP += "\n\n\t\t"+string_special(OPTION_OUTPUT_SHORT)+","+\
    string_special(OPTION_OUTPUT)+\
    " <fmt> Output format for ropchains.\n\t\t\t\t\t Expected format is one of the\n\t\t\t\t\t following: "+\
    string_special(OUTPUT_CONSOLE)+','+string_special(OUTPUT_PYTHON)
    
CMD_CALL_HELP += "\n\n\t\t"+string_special(OPTION_HELP_SHORT)+","+string_special(OPTION_HELP)+"\t\t Show this help"
CMD_CALL_HELP += "\n\n\t"+string_bold("Function format")+": "+\
    string_special("function")+"( "+string_special("arg1")+","+string_special(" ...")+\
    "," + string_special(" argN")+")"
CMD_CALL_HELP += "\n\n\t"+string_bold("Examples")+": "+\
    "\n\t\tcall strcpy(0x123456, 0x987654)"
    
def print_help():
    print(CMD_CALL_HELP)


def call(args):
    global OUTPUT, OUTPUT_CONSOLE, OUTPUT_PYTHON
    # Parsing arguments
    if( not args):
        print_help()
        return
    OUTPUT = OUTPUT_CONSOLE
    funcName = None
    i = 0
    seenOutput = False
    seenFunction = False
    seenBadBytes = False
    seenKeepRegs = False
    seenShortest = False
    seenLmax = False
    clmax = None
    
    constraint = Constraint()
    assertion = Assertion()
    while i < len(args):
        if( args[i] in [OPTION_LIST, OPTION_LIST_SHORT]):
            func_list = getAllFunctions()
            print_functions(func_list)
            return 
        if( args[i] in [OPTION_BAD_BYTES, OPTION_BAD_BYTES_SHORT]):
            if( seenBadBytes ):
                error("Error. '" + args[i] + "' option should be used only once")
                return 
            if( i+1 >= len(args)):
                error("Error. Missing bad bytes after option '"+args[i]+"'")
                return 
            seenBadBytes = True
            (success, res) = parse_bad_bytes(args[i+1])
            if( not success ):
                error(res)
                return
            i = i+2
            constraint = constraint.add(BadBytes(res))
        elif( args[i] in [OPTION_KEEP_REGS, OPTION_KEEP_REGS_SHORT]):
            if( seenKeepRegs ):
                error("Error. '" + args[i] + "' option should be used only once")
                return 
            if( i+1 >= len(args)):
                error("Error. Missing register after option '"+args[i]+"'")
                return 
            seenKeepRegs = True
            (success, res) = parse_keep_regs(args[i+1])
            if( not success ):
                error(res)
                return
            i = i+2
            constraint = constraint.add(RegsNotModified(res))
        elif( args[i] in [OPTION_CALL, OPTION_CALL_SHORT] ):
            if( not loadedBinary() ):
                error("Error. You should load a binary before building ROPChains")
                return 
            elif( seenFunction ):
                error("Option '{}' should be used only once".format(args[i]))
                return  
            userInput = ''
            i +=1
            while( i < len(args) and args[i][0] != "-"):
                userInput += args[i]
                i += 1
            (funcName, funcArgs ) = parseFunction(userInput)
            if( not funcName):
                return 
            seenFunction = True
        elif( args[i] in [OPTION_OUTPUT, OPTION_OUTPUT_SHORT]):
            if( seenOutput ):
                error("Option '{}' should be used only once".format(args[i]))
                return 
            if( i+1 >= len(args)):
                error("Error. Missing output format after option '"+args[i]+"'")
                return 
            if( args[i+1] in [OUTPUT_CONSOLE, OUTPUT_PYTHON]):
                OUTPUT = args[i+1]
                seenOutput = True
                i += 2
            else:
                error("Error. Unknown output format: {}".format(args[i+1]))
                return 
        elif( args[i] in [OPTION_SHORTEST, OPTION_SHORTEST_SHORT]):
            if( seenShortest ):
                error("Option '{}' should be used only once".format(args[i]))
                return 
            seenShortest = True
            i += 1
        elif( args[i] == OPTION_LMAX or args[i] == OPTION_LMAX_SHORT ):
            if( seenLmax ):
                error("Option '{}' should be used only once".format(args[i]))
                return 
            if( i+1 >= len(args)):
                error("Error. Missing length after option '"+args[i]+"'")
                return 
            try:
                clmax = int(args[i+1])
                if( clmax < Arch.octets() ):
                    raise Exception()
                # Convert number of bytes into number of ropchain elements
                clmax /= Arch.octets()
            except:
                error("Error. '" + args[i+1] +"' bytes is not valid")
                return 
            i += 2 
            seenLmax = True
        elif( args[i] in [OPTION_HELP, OPTION_HELP_SHORT]):
            print_help()
            return 
        else:
            error("Error. Unknown option '{}'".format(args[i]))
            return 
    
    if( not funcName ):
        error("Missing function to call")
    else:
        res = build_call(funcName, funcArgs, constraint, assertion, clmax=clmax, optimizeLen=seenShortest)
        if( isinstance(res, str) ):
            error(res)
        else:
            print_chains([res], "Built matching ROPChain", constraint.getBadBytes())

def build_call(funcName, funcArgs, constraint, assertion, argsDescription=None, clmax=None, optimizeLen=False):
    """
    funcArgs : list of pairs (arg_value, arg_description)
    """
    # Merge description and args 
    if( argsDescription ):
        funcArgs = zip(funcArgs, argsDescription)
    else:
        funcArgs = [(arg,) for arg in funcArgs]
    
    if( Arch.currentBinType == Arch.BinaryType.X86_ELF ):
        return build_call_linux86(funcName, funcArgs, constraint, assertion, clmax, optimizeLen)
    elif( Arch.currentBinType == Arch.BinaryType.X64_ELF ):
        return build_call_linux64(funcName, funcArgs, constraint, assertion, clmax, optimizeLen)
    return []
        
def build_call_linux64(funcName, funcArgs, constraint, assertion, clmax=None, optimizeLen=False):
    # Arguments registers 
    # (Args should go in these registers for x64)
    argsRegsNames = ['rdi','rsi','rdx','rcx', 'r8', 'r9']
    argsRegs = [Arch.n2r(name) for name in argsRegsNames]
    # Find the address of the fonction 
    (funcName2, funcAddr) = getFunctionAddress(funcName)
    if( funcName2 is None ):
        return "Couldn't find function '{}' in the binary".format(funcName)
    
    # Check if bad bytes in function address 
    if( not constraint.badBytes.verifyAddress(funcAddr) ):
        return "'{}' address ({}) contains bad bytes".format(funcName2, string_special('0x'+format(funcAddr, '0'+str(Arch.octets()*2)+'x')))
    
    # Check how many arguments 
    if( len(funcArgs) > 6 ):
        return "Doesn't support function call with more than 6 arguments with Linux X64 calling convention :("
        
    # Find a gadget for the fake return address
    if( funcArgs ):
        # Build the ropchain with the arguments
        args_chain = popMultiple(map(lambda x,y:(x,)+y,  argsRegs[:len(funcArgs)], funcArgs), constraint, assertion, clmax=clmax, optimizeLen=optimizeLen)
        if( not args_chain):
            return "Couldn't load arguments in registers"
    else:
        # No arguments 
        args_chain = ROPChain()
    
    # Build call chain (function address + fake return address)
    return args_chain.addPadding(funcAddr, comment=string_ropg(funcName2))
    

def build_call_linux86(funcName, funcArgs, constraint, assertion, clmax=None, optimizeLen=False):
    # Find the address of the fonction 
    (funcName2, funcAddr) = getFunctionAddress(funcName)
    if( funcName2 is None ):
        return "Couldn't find function '{}' in the binary".format(funcName)
    
    # Check if bad bytes in function address 
    if( not constraint.badBytes.verifyAddress(funcAddr) ):
        return "'{}' address ({}) contains bad bytes".format(funcName2, string_special('0x'+format(funcAddr, '0'+str(Arch.octets()*2)+'x')))
    
    # Check if lmax too small
    if( (1 + len(funcArgs) + (lambda x: 1 if len(x)>0 else 0)(funcArgs)) > clmax ):
        return "Not enough bytes to call function '{}'".format(funcName)
    
    # Find a gadget for the fake return address
    if( funcArgs ):
        offset = (len(funcArgs)-1)*Arch.octets() # Because we do +octets() at the beginning of the loop
        skip_args_chains = []
        i = 4 # Try 4 more maximum 
        while( i > 0 and (not skip_args_chains)):
            offset += Arch.octets() 
            skip_args_chains = search(QueryType.MEMtoREG, Arch.ipNum(), \
                        (Arch.spNum(),offset), constraint, assertion, n=1, optimizeLen=optimizeLen)
            i -= 1
            
        if( not skip_args_chains ):
            return "Couldn't build ROP-Chain"
        skip_args_chain = skip_args_chains[0]
    else:
        # No arguments 
        skip_args_chain = None
    
    # Build the ropchain with the arguments 
    args_chain = ROPChain()
    arg_n = len(funcArgs)
    for arg in reversed(funcArgs):
        if( isinstance(arg, int) ):
            args_chain.addPadding(arg, comment="Arg{}: {}".format(arg_n, string_ropg(hex(arg))))
            arg_n -= 1
        else:
            return "Type of argument '{}' not supported yet :'(".format(arg)
    
    # Build call chain (function address + fake return address)
    call_chain = ROPChain()
    call_chain.addPadding(funcAddr, comment=string_ropg(funcName2))
    if( funcArgs ):
        skip_args_addr = int( validAddrStr(skip_args_chain.chain[0], constraint.getBadBytes(), Arch.bits())  ,16)
        call_chain.addPadding(skip_args_addr, comment="Address of: "+string_bold(str(skip_args_chain.chain[0])))
    
    return call_chain.addChain(args_chain)
    
def parseFunction(string):
    def seek(char, string):
        for i in range(0, len(string)):
            if string[i] == char:
                return (string[:i], i)
        return ([],-1)
        
    if( not string ):
        error("Missing fuction to call")
        return (None, None)
    
    # COmpress the string
    string = "".join(string.split())
    
    # Get the function name 
    (funcName, index) = seek("(", string)
    if( not funcName ):
        error("Invalid function call")
        return (None, None)
    rest = string[index+1:]
    args = []
    arg = ''
    i = 0
    end = False
    while(i < len(rest)):
        c = rest[i]
        # No args
        if( c == ")" and not args):
            end = True
            i += 1
        # String
        elif( c == '"' or c == "'" ):
            (s, index)= seek(c, rest[i+1:])
            if( not s ):
                error("Missing closing {} for string".format(c))
                return (None, None)
            args.append(s)
            i += index +2
            if( i >= len(rest)):
                error("Error. Missing ')'")
                return (None, None)
            elif( rest[i] == ')' ):
                end = True
                i += 1
            elif( rest[i] == "," ):
                i += 1
        # Constant
        else:
            # Get the constant 
            arg = ''
            ok = False
            for j in range(i, len(rest)):
                if( rest[j] == ")" ):
                    end = True
                    ok = True
                    break
                elif( rest[j] == ','):
                    ok = True
                    break
                else:
                    arg += rest[j]
            if( not ok ):
                error("Missing ')' after argument")
                return (None, None)
            if( (not arg) and args):
                error("Missing argument")
                return (None, None)
            # Convert to int 
            try:
                value = int(arg)
            except:
                try:
                    value = int(arg, 16)
                except:
                    try:
                        value = int(arg, 2)
                    except:
                        error("Invalid operand: " + arg )
                        return (None, None)
            args.append(value)
            i = j+1
        if( end):
            break
    
    if( not end ):
        error("Error. Missing ')'")
        return    (None, None)     
    if( i < len(rest)):
        error("Error. Extra argument: {}".format(rest[i:]))
        return (None, None)

    return (funcName, args)

## ------------
## Parsing from ropgenerator.semantic.Find
def parse_bad_bytes(string):
    """
    Parses a bad bytes string into a list of bad bytes
    Input: a string of format like "00,0A,FF,32,C7"
    Ouput if valid string : (True, list) where list = 
        ['00', '0a', 'ff', '32', 'c7'] (separate them in individual strings
        and force lower case)
    Output if invalid string (False, error_message)
    """
    hex_chars = '0123456789abcdefABCDEF'
    i = 0
    bad_bytes = []
    user_bad_bytes = [b.lower() for b in string.split(',')]
    for user_bad_byte in user_bad_bytes:
        if( not user_bad_byte ):
            return (False, "Error. Missing bad byte after ','")
        elif( len(user_bad_byte) != 2 ):
            return (False, "Error. '{}' is not a valid byte".format(user_bad_byte))
        elif( not ((user_bad_byte[i] in hex_chars) and (user_bad_byte[i+1] in hex_chars))):
            return (False, "Error. '{}' is not a valid byte".format(user_bad_byte))
        else:
            bad_bytes.append(user_bad_byte)
    return (True, bad_bytes)
    
def parse_keep_regs(string):
    """
    Parses a 'keep registers' string into a list of register uids
    Input: a string of format like "rax,rcx,rdi"
    Output if valid string (True, list) where list = 
        [1, 3, 4] (R1 is rax, R3 is RCX, ... )
    Output if invalid string (False, error_message)
    """
    user_keep_regs = string.split(',')
    keep_regs = set()
    for reg in user_keep_regs:
        if( reg in Arch.regNameToNum ):
            keep_regs.add(Arch.n2r(reg))
        else:
            return (False, "Error. '{}' is not a valid register".format(reg))
    return (True, list(keep_regs))

##########################
# Pretty print functions #
##########################

def print_chains(chainList, msg, badBytes=[]):
    global OUTPUT
    sep = "------------------"
    if( chainList):
        print(string_bold('\n\t'+msg))
        if( OUTPUT == OUTPUT_CONSOLE ):
            print("\n"+chainList[0].strConsole(Arch.currentArch.bits, badBytes))
        elif( OUTPUT == OUTPUT_PYTHON ):
            print('\n' + chainList[0].strPython(Arch.currentArch.bits, badBytes))
        for chain in chainList[1:]:
            if( OUTPUT == OUTPUT_CONSOLE ):
                print('\t'+sep + "\n"+ chain.strConsole(Arch.currentArch.bits, badBytes))
            elif( OUTPUT == OUTPUT_PYTHON ):
                print('\t'+sep + '\n' + chain.strPython(Arch.currentArch.bits, badBytes))
    else:
        print(string_bold("\n\tNo matching ROPChain found"))

def print_functions(func_list):
    """
    func_list - list of pairs (str, int) = (funcName, funcAddress)
    """
    space = 28
    print(banner([string_bold("Available functions")]))
    print("\tFunction" + " "*(space-8) + "Address")
    print("\t------------------------------------")
    for (funcName, funcAddr) in sorted(func_list, key=lambda x:x[0] ):
        space2 = space - len(funcName)
        if( space2 < 0 ):
            space2 = 2
        print("\t"+string_special(funcName)+ " "*space2 + hex(funcAddr))
    print("")
    
