# -*- coding:utf-8 -*- 
# Syscall module: building ropchains performing syscalls
from ropgenerator.IO import string_special, banner, string_bold, error
from enum import Enum
import ropgenerator.exploit.syscalls.Linux32 as Linux32
import ropgenerator.exploit.syscalls.Linux64 as Linux64
import ropgenerator.Architecture as Arch
from ropgenerator.Constraints import Constraint, Assertion, BadBytes, RegsNotModified
from ropgenerator.Load import loadedBinary

#####################
# Available systems #
#####################
sysLinux32 = "LINUX32"
sysLinux64 = "LINUX64"

availableSystems = [sysLinux32, sysLinux64]

###################
# SYSCALL 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_KEEP_REGS = '--keep-regs'
OPTION_KEEP_REGS_SHORT = '-k'

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

OPTION_FUNCTION = "--call"
OPTION_FUNCTION_SHORT = "-c" 

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

CMD_SYSCALL_HELP =  banner([string_bold("'syscall' command"),\
                    string_special("(Call system functions with ROPChains)")])
CMD_SYSCALL_HELP += "\n\n\t"+string_bold("Usage:")+\
"\n\t\tsyscall [OPTIONS]"
CMD_SYSCALL_HELP += "\n\n\t"+string_bold("Options")+":"
CMD_SYSCALL_HELP += "\n\t\t"+string_special(OPTION_FUNCTION_SHORT)+","+\
    string_special(OPTION_FUNCTION)+" <function>\t Call a system function"
CMD_SYSCALL_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_SYSCALL_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_SYSCALL_HELP += "\n\n\t\t"+string_special(OPTION_LIST_SHORT)+","+\
    string_special(OPTION_LIST)+" [<system>]\t List supported functions"
CMD_SYSCALL_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_SYSCALL_HELP += "\n\n\t\t"+string_special(OPTION_HELP_SHORT)+","+string_special(OPTION_HELP)+"\t\t Show this help"
CMD_SYSCALL_HELP += "\n\n\t"+string_bold("Supported systems")+": "+', '.join([string_special(s) for s in availableSystems])
CMD_SYSCALL_HELP += "\n\n\t"+string_bold("Function format")+": "+\
    string_special("function")+"( "+string_special("arg1")+","+string_special(" ...")+\
    "," + string_special(" argN")+")"
CMD_SYSCALL_HELP += "\n\n\t"+string_bold("Examples")+": "+\
    "\n\t\tsyscall -f python -c mprotect32(0x123456, 200, 7)\n\t\tsyscall -l LINUX64"
def print_help():
    print(CMD_SYSCALL_HELP)


def syscall(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
    
    constraint = Constraint()
    assertion = Assertion()
    while i < len(args):
        if( args[i] in [OPTION_LIST, OPTION_LIST_SHORT]):
            listSyscalls(args[1:])
            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_FUNCTION, OPTION_FUNCTION_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_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:
        call(funcName, funcArgs, constraint, assertion)

def call(funcName, parsedArgs, constraint, assertion):
    # Get target system
    if( Arch.currentBinType == Arch.BinaryType.X86_ELF ):
        syscall = Linux32.available.get(funcName)
        system = sysLinux32
    elif( Arch.currentBinType == Arch.BinaryType.X64_ELF ):
        syscall = Linux64.available.get(funcName)
        system = sysLinux64
    else:
        error("Binary type '{}' not supported yet".format(Arch.currentBinType))
        return 
    
    if( not syscall ):
        error("Syscall '{}' not supported for system '{}'".format(\
        funcName, system))
        return 
    
    if( len(parsedArgs) != len(syscall.args)):
        error("Error. Wrong number of arguments")
        return 
        
    # Build syscall
    res = _build_syscall(syscall.buildFunc, parsedArgs, constraint, assertion, optimizeLen = True)
    # Print result 
    if( not res ):
        print(string_bold("\n\tNo matching ROPChain found"))
    else:
        print(string_bold("\n\tFound matching ROPChain\n"))
        badBytes = constraint.getBadBytes()
        if( OUTPUT == OUTPUT_CONSOLE ):
            print(res.strConsole(Arch.bits(), badBytes))
        elif( OUTPUT == OUTPUT_PYTHON ):
            print(res.strPython(Arch.bits(), badBytes))
    
    
def _build_syscall(funcPointer, args , constraint, assertion, optimizeLen=False):
    if( len(args) == 0 ):
        return funcPointer(constraint, assertion, optimizeLen=optimizeLen)
    elif( len(args) == 1 ):
        return funcPointer(args[0], constraint, assertion,optimizeLen=optimizeLen)
    elif( len(args) == 2 ):
        return funcPointer(args[0], args[1], constraint, assertion, optimizeLen=optimizeLen)
    elif( len(args) == 3 ):
        return funcPointer(args[0], args[1], args[2], constraint, assertion, optimizeLen=optimizeLen)
    elif( len(args) == 4 ):
        return funcPointer(args[0], args[1], args[2], args[3], constraint, assertion, optimizeLen=optimizeLen)
    elif( len(args) == 5 ):
        return funcPointer(args[0], args[1], args[2], args[3], args[4], constraint, assertion, optimizeLen=optimizeLen)
    else:
        error("{}-arguments calls not supported".format(len(args)))
        return []

def listSyscalls(args):
    if( not args ):
        systems = availableSystems
    else:
        systems = []
        for arg in args:
            if( arg in availableSystems ):
                systems.append(arg)
            else:
                error("Unknown system: '{}'".format(arg))
                return 
    
    for system in list(set(systems)):
        if( system == sysLinux32 ):
            Linux32.print_available()
        elif( system == sysLinux64 ):
            Linux64.print_available()


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))

