# The MIT License (MIT)
#
# Copyright (c) 2018 Thomas Barabosch of Fraunhofer FKIE
# <thomas.barabosch (at) fkie.fraunhofer.de>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
# the Software, and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
try:
    import gdb
except ImportError as e:
    raise ImportError("This script must be run in GDB: ", str(e))

import os
import re
import warnings

from lib.gdb_wrapper.elf import read_elf_sects
from lib.tools import AttrDict, memoized
from lib.gdb_wrapper.x86 import Target, Operand, Instruction


class MipsInstruction(Instruction):

    def __init__(self, gdbstr):
        Instruction.__init__(self, gdbstr)
        self.gdbstr = gdbstr

        # strip indicator
        if self.gdbstr[0:2] == "=>":
            self.gdbstr = self.gdbstr[2:]

        # strip hints
        c = 0
        s = self.gdbstr
        self.gdbstr = ""
        for i in range(0, len(s)):
            if s[i] == "<":
                c += 1
                continue
            elif s[i] == ">":
                c -= 1
                continue
            if c > 0:
                continue
            self.gdbstr += s[i]

        toks = []
        for t in self.gdbstr.split():
            t = t.strip().strip(":")
            if t:
                toks.append(t)
        self.addr = int(toks[0], 16)

        self.mnemonic = toks[1]
        self.inst = " ".join(toks[2:])

        if self.inst == "":  # handle "ret", "iret", et al.
            return

        # get operands
        self.operands = [MipsOperand(o) for o in self.inst.split(",")]
        for a, o in zip(self.__get_operand_order(self.mnemonic), self.operands):
            setattr(self, a, o)

    def __get_operand_order(self, mnemonic):
        # TODO: add more instructions
        # http://www.mrc.uidaho.edu/mrc/people/jff/digital/MIPSir.html

        # possbible orders
        # TODO: "aux" needs to be in orders?
        dst = ["dest", "source", "source2"]  # e.g. add
        dsi = ["dest", "source", "immediate"]  # e.g. addi
        std = ["source", "source2", "dest"]  # e.g. beq
        sd = ["source", "dest"]  # e.g. bgez
        ss = ["source", "source2"]  # e.g. div
        d = ["dest"]  # e.g. j
        ds = ["dest", "source"]  # e.g. lb
        di = ["dest", "immediate"]  # e.g. lui

        # mapping from instruction to operand order
        operand_dict = {
            "add": dst,
            "addi": dsi,
            "addiu": dsi,
            "addu": dst,
            "and": dst,
            "andi": dsi,
            "b": d,
            "bal": d,
            "beq": std,
            "bgez": sd,
            "bgezal": sd,
            "bgtz": sd,
            "blez": sd,
            "bltz": sd,
            "bltzal": sd,
            "bne": std,
            "div": ss,
            "divu": ss,
            "j": d,
            "jal": d,
            "jar": d,
            "lb": ds,
            "lbu": ds,
            "lui": di,
            "lw": ds,
            "mfhi": d,
            "mflo": d,
            "mult": ss,
            "multu": ss,
            "or": dst,
            "ori": dsi,
            "sb": sd,
            "sll": dsi,
            "sllv": dst,
            "slt": dst,
            "slti": dsi,
            "sltiu": dsi,
            "sltu": dst,
            "sra": dsi,
            "srl": dsi,
            "srlv": dst,
            "sub": dst,
            "subu": dst,
            "sw": sd,
            "xor": dst,
            "xori": dsi
        }

        if mnemonic in operand_dict:
            return operand_dict[mnemonic]
        else:
            print("[MipsOperand] Unknown mnemonic %s. Please update order dict!" % mnemonic)
            return []


class MipsOperand(Operand):
    '''
    A disassembled instruction operand that can be evaluated. Notable
    attributes include:

    gdbstr - the original string of disassembled operand that was used to
        populate this instance
    is_pointer - True if this operand represents a pointer, False otherwise.
        For example, "DWORD PTR [eax]" is considered a pointer, "eax" is not.
    regs - a list of strings representing the registers used in this Operand
   '''

    _re_hex = re.compile(r'^(-?0x[0-9a-fA-F]+)$')
    _re_addr = re.compile(r""" (-?0?x?[A-Fa-f0-9]*)\(([a-z]{1,5}\d{0,2}[a-z]{0,1})\)$    # "0x10($sp) or -0x7fa0($gp)"
                           """, re.VERBOSE)
    _re_regstrs = re.compile(
        r"""([a-z]{1,5}\d{0,2}[a-z]{0,1})                               # "$ra or $t9"
                 """, re.VERBOSE)
    _re_p_regs = re.compile(r"zero")

    def __init__(self, gdbstr):
        '''
        Parses gdbstr to populate self.  ex gdbstr="BYTE PTR es:[edi]"
        '''
        self.gdbstr = gdbstr
        gdbstr = gdbstr.split("<")[0]  # get rid of "<addr resolutions>"

        # check if hexadecimal immediate value
        imm = self._re_hex.match(gdbstr)
        if imm:
            self.is_pointer = False
            expr = [m for m in imm.groups() if m][0]

        # get addr
        addr = self._re_addr.match(gdbstr)  # ignores segment regs
        if addr:
            self.is_pointer = True
            expr = "+".join([m for m in addr.groups() if m])
        else:
            self.is_pointer = False
            expr = gdbstr

        # the MIPS register $zero/$0 evaluates always to 0
        expr = self._re_p_regs.sub("0", expr)

        self.regs = [t.group() for t in self._re_regstrs.finditer(expr)]
        # prep for GDB evaluation. ex: "edx+0x12" becomes "$edx+0x12"
        if not imm:
            self.expr = self._re_regstrs.sub(
                (lambda mo: "${}".format([mg for mg in mo.groups() if mg][0])), expr)
        else:
            self.expr = expr


class MipsTarget(Target):
    '''
    A wrapper for an ARM Linux GDB Inferior.
    '''

    def _getInstruction(self, gdbstr):
        return MipsInstruction(gdbstr)
