### BEGIN LICENSE ###
### Use of the triage tools and related source code is subject to the terms
### of the license below.
###
### ------------------------------------------------------------------------
### Copyright (C) 2011 Carnegie Mellon University. All Rights Reserved.
### Portions Copyright 2013 BlackBerry Ltd. All Rights Reserved.
### ------------------------------------------------------------------------
### Redistribution and use in source and binary forms, with or without
### modification, are permitted provided that the following conditions are
### met:
###
### 1. Redistributions of source code must retain the above copyright
###    notice, this list of conditions and the following acknowledgments
###    and disclaimers.
###
### 2. Redistributions in binary form must reproduce the above copyright
###    notice, this list of conditions and the following disclaimer in the
###    documentation and/or other materials provided with the distribution.
###
### 3. All advertising materials for third-party software mentioning
###    features or use of this software must display the following
###    disclaimer:
###
###    "Neither Carnegie Mellon University nor its Software Engineering
###     Institute have reviewed or endorsed this software"
###
### 4. The names "Department of Homeland Security," "Carnegie Mellon
###    University," "CERT" and/or "Software Engineering Institute" shall
###    not be used to endorse or promote products derived from this software
###    without prior written permission. For written permission, please
###    contact permission@sei.cmu.edu.
###
### 5. Products derived from this software may not be called "CERT" nor
###    may "CERT" appear in their names without prior written permission of
###    permission@sei.cmu.edu.
###
### 6. Redistributions of any form whatsoever must retain the following
###    acknowledgment:
###
###    "This product includes software developed by CERT with funding
###     and support from the Department of Homeland Security under
###     Contract No. FA 8721-05-C-0003."
###
### THIS SOFTWARE IS PROVIDED BY CARNEGIE MELLON UNIVERSITY ``AS IS'' AND
### CARNEGIE MELLON UNIVERSITY MAKES NO WARRANTIES OF ANY KIND, EITHER
### EXPRESS OR IMPLIED, AS TO ANY MATTER, AND ALL SUCH WARRANTIES, INCLUDING
### WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE, ARE
### EXPRESSLY DISCLAIMED. WITHOUT LIMITING THE GENERALITY OF THE FOREGOING,
### CARNEGIE MELLON UNIVERSITY DOES NOT MAKE ANY WARRANTY OF ANY KIND
### RELATING TO EXCLUSIVITY, INFORMATIONAL CONTENT, ERROR-FREE OPERATION,
### RESULTS TO BE OBTAINED FROM USE, FREEDOM FROM PATENT, TRADEMARK AND
### COPYRIGHT INFRINGEMENT AND/OR FREEDOM FROM THEFT OF TRADE SECRETS.
### END LICENSE ###
'''
A collection of Python objects that wrap and extend the GDB Python API.

See x86.py in this directory for more info.
'''
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 ArmInstruction(Instruction):
    '''
    A disassembled ARM instruction. See Instruction class comments for more 
    details.

    This code was graciously contributed by a user and needs some cleanup -- 
    please help!
    '''
    _re_arm_operands = re.compile(r"""(?P<a>0x[A-Fa-f0-9]+)      | # addr
                                      (?P<b>{)                   | # start brace
                                      (?P<c>[cs]psr)             | # <PSR>
                                      (?P<e>\])                  | # end bracket
                                      _(?P<f>[cfgnqsvxz])        | # <fields|flags>
                                      (?P<g>[aif])               | # <iflags>
                                      \#(?P<i>-?[0-9]+)          | # immediate value (with #)
                                      (?P<l>lsr|lsl|asr|ror|rrx) | # shift/rotate
                                      (?P<n>[bl]e)               | # endianness
                                      (?P<p>[+-])                | # +/-
                                      (?P<q>})                   | # end brace
                                      (?P<r>r\d{1,2}|lr|pc|sp)   | # register
                                      (?P<s>\[)                  | # start bracket
                                      (?P<x>;.*)                   # semicolon annotation
                                   """, re.VERBOSE)
    _re_arm_groups = re.compile(r"[sebq]")

    def __init__(self, gdbstr):
        '''
        Parses gdbstr to populate self.
        '''
        Instruction.__init__(self, gdbstr)
        self.gdbstr = gdbstr

        def split_expr(inst, tok):
            com = inst.index(",", tok.end(0))
            return inst[:com], inst[com:]

        # get instruction addr and instruction string
        inst = [pt for pt in self.gdbstr.split(":\t", 1)[1].strip().split() if pt]
        self.addr = int(self._re_hex_int.search(self.gdbstr).group(), 16)
        self.mnemonic = inst[0]
        inst = " ".join(inst[1:])
        if inst == "": # handle "ret", "iret", et al.
            return

        # tokenize the arguments
        toks = [t for t in self._re_arm_operands.finditer(inst)]
        form = "".join(t.lastgroup for t in toks)
        if form.endswith("x"):
            form = form[:-1]
            toks = toks[:-1]
        form = self._re_arm_groups.sub(lambda c:{"s": "[", "e": "]", "b": "{", "q": "}"}[c.group()], form)
        strtoks = [t.group(0) for t in toks]
        if form in ["rri",           # add, adc, sub, sbc, rsb, rsc, asr, lsl, lsr, ror, and, eor, orr, orn, bic
                    "rrr",           # add, adc, sub, sbc, rsb, rsc, qadd, qdadd, qsub, qdsub,
                                     #   (s|q|sh|u|uq|uh)(add|sub|asx|sax), usad8, mul, smul, smuad, smusd, smmul,
                                     #   [su]div, asr, lsl, lsr, ror, and, eor, orr, orn, bic, pkhbt, pkhtb, [us]xta[hb], sel
                    "rrrli",         # add, adc, sub, sbc, rsb, rsc, and, eor, orr, orn, bic, pkhbt, pkhtb, [us]xta[hb]
                    "rrrlr", "rrrl", # add, adc, sub, sbc, rsb, rsc, and, eor, orr, orn, bic
                    "rir", "rirli",  # [us]sat
                    "ri",   # mov, mvn, cmp, cmn, tst, teq, srs
                    "rr",   # mov, mvn, rrx, clz, cmp, cmn, tst, teq, [us]xt[hb], rbit, rev
                    "rrli", # mov, mvn, cmp, cmn, tst, teq, [us]xt[hb]
                    "rrl", "rlr",  # mov, mvn, cmp, cmn, tst, teq
                    "rii",  # bfc
                    "rrii", # bfi, [us]bfx
                    ]:
            dst, src = split_expr(inst, toks[0])
            self.dest = ArmOperand(dst, strtoks[:1], form[:1])
            self.source = ArmOperand(src, strtoks[1:], form[1:])
        elif form == "ra": # adr, cbz, cbnz, ldr
            short = self.mnemonic[:3]
            if short in ["cbz", "cbn"]:
                self.dest = ArmOperand("", ["pc"], "r")
                self.source = ArmOperand(inst, strtoks, form)
            elif short in ["ldr", "adr"]:
                dst, src = split_expr(inst, toks[0])
                self.dest = ArmOperand(dst, strtoks[:1], form[:1])
                self.source = ArmOperand(src, strtoks[1:], form[1:], is_pointer=(short =="ldr"))
            else:
                raise RuntimeError("unexpected mnemonic for form \"{}\": {}".format(form, self.mnemonic))
        elif form == "rrrr": # usada8, mla, mls, ([us]mull|[us]mlal|umaal|smlsld), smla, smlad, smlald, smlsd, smmla, smmls
            if self.mnemonic in ("umull", "smull", "umlal", "smlal", "umaal", "smlsld"):
                dst, src = split_expr(inst, toks[1])
                dst_r = strtoks[:2]
                dst_f = form[:2]
                if self.mnemonic in ["umlal", "umaal", "smlal", "smlsld"]:
                    src = inst
                    src_r = strtoks
                    src_f = form
                else:
                    src_r = strtoks[2:]
                    src_f = form[2:]
            else:
                dst, src = split_expr(inst, toks[0])
                dst_r = strtoks[:1]
                dst_f = form[:1]
                src_r = strtoks[1:]
                src_f = form[1:]
            self.dest = ArmOperand(dst, dst_r, dst_f)
            self.source = ArmOperand(src, src_r, src_f)
        elif form in ["a", # b, bl, blx, pld, pli
                      "r"]: # bx, blx, bxj, rfe
            if self.mnemonic.startswith("b"):
                self.dest = ArmOperand("", ["pc"], "r") # ignore LR so it evaluates properly
                self.source = ArmOperand(inst, strtoks, form)
            elif self.mnemonic.startswith("r"):
                self.dest = ArmOperand("", ["pc"], "r")
                self.source = ArmOperand(inst, strtoks, form, is_pointer=True)
            else:
                # pli/pld shouldn't fault
                raise RuntimeError("unexpected mnemonic for form \"{}\": {}".format(form, self.mnemonic))
        elif form in ["[rr]",    # tbb, pli, pld
                      "[rrli]"]: # tbh, pli, pld
            if self.mnemonic.startswith("p"):
                # pli/pld shouldn't fault
                raise RuntimeError("unexpected mnemonic for form \"{}\": {}".format(form, self.mnemonic))
            self.dest = ArmOperand("", ["pc"], "r")
            self.source = ArmOperand(inst, strtoks, form, is_pointer=True)
        elif form == "rc": # mrs
            self.dest = ArmOperand(inst, strtoks[:1], form[:1])
        elif form in ["cfr", # msr
                      "cfi"]: # msr
            self.source = ArmOperand(inst, strtoks[-1:], form[-1:])
        elif form in ["g", "gi", # cpsi[de]
                      "i",       # cps, bkpt, smc, svc
                      "n"]:      # setend
            pass
        elif form in ["r[r]", # ldr, str, ldrex
                      "r[ri]", "r[rr]", "r[rpr]", "r[rrli]", "r[rprli]", "r[rrl]", "r[rprl]", # ldr, str
                      "rr[r]", # ldrd, strd, ldrexd, strex, swp
                      "rr[ri]", "rr[rr]", "rr[rpr]", "rr[rrli]", "rr[rprli]", "rr[rrl]", "rr[rprl]", # ldrd, strd
                      "rra", # ldrd
                      "rrr[r]"]: # strexd
            regs = len(form) - len(form.lstrip("r"))
            dst, src = split_expr(inst, toks[regs-1])
            if self.mnemonic.startswith("l"):
                order = ["dest", "source"]
            else:
                order = ["source", "dest"]
            setattr(self, order[0], ArmOperand(dst, strtoks[:regs], form[:regs]))
            setattr(self, order[1], ArmOperand(src, strtoks[regs:], form[regs:], is_pointer=True))
        elif form in ["r[r]i", "r[r]r", "r[r]pr", "r[r]rli", "r[r]prli", "r[r]rl", "r[r]prl", # ldr, str
                      "rr[r]i", # ldrd, strd
                      "rr[r]r", "rr[r]pr", "rr[r]rli", "rr[r]prli", "rr[r]rl", "rr[r]prl"]: # ldrd, strd
            regs = len(form) - len(form.lstrip("r"))
            dst, src = split_expr(inst, toks[regs-1])
            if self.mnemonic.startswith("l"):
                order = ["dest", "source"]
            else:
                order = ["source", "dest"]
            setattr(self, order[0], ArmOperand(dst, strtoks[:regs], form[:regs]))
            setattr(self, order[1], ArmOperand(src, strtoks[regs:regs+3], form[regs:regs+3], is_pointer=True))
        #elif form == "[r]": # pli, pld
        #elif form == "[ri]": # pli, pld
        #elif form == "[rpr]": # pli, pld
        #elif form == "[rprli]": # pli, pld
        #elif form == "[rrl]": # pli, pld
        #elif form == "[rprl]": # pli, pld
        elif form[:2] in ["r{", # ldm, stm
                          "{r"]: # pop, push
            mne = self.mnemonic
            if mne.startswith("p"):
                form = "r" + form
                strtoks.insert(0, "sp")
                if mne[1] == "u": # push
                    mne = "stm"
                else:
                    mne = "ldm"
                dst, src = "", inst
            else:
                dst, src = split_expr(inst, toks[0])
            if mne.startswith("s"):
                order = ["dest", "source"]
            else:
                order = ["source", "dest"]
            setattr(self, order[0], ArmOperand(dst, strtoks[:1], form[:1], is_pointer=True))
            setattr(self, order[1], ArmOperand(src, strtoks[1:], form[1:]))
        elif self.mnemonic in ("cdp", "mrc", "mrrc", "mcr", "ldc", "stc"):
            # coprocessor writes. very special!
            raise RuntimeError("coprocessor instructions are unhandled, and you can too!")
        else:
            raise RuntimeError("unexpected mnemonic for form \"{}\": {}".format(form, self.mnemonic))
        self.operands = [o.regs for o in (self.dest, self.source, self.aux) if o]

class ArmOperand(Operand):
    '''
    A disassembled ARM instruction operand that can be evaluated. See Operand
    class comments for more details. 
    '''
    def __init__(self, gdbstr, toklist, form, is_pointer=False):
        self.gdbstr = gdbstr
        self.regs = []
        self.is_pointer = is_pointer
        self.expr = ""
        assert len(form) == len(toklist), "form=%s, toklist=%s" % (form, toklist)
        eat = False
        for f, t in zip(form, toklist):
            if f in "[]{}cfgn":
                continue
            if f in "ai":
                if not eat and self.expr:
                    self.expr += "+"
                if f == "a":
                    self.expr += "{}".format(int(t, 16))
                else:
                    self.expr += "{}".format(int(t[1:]))
                eat = False
            elif f == "l":
                self.expr += {"lsl": "<<",
                              "lsr": ">>",
                              "asr": ">>",
                              "ror": ">>",
                              "rrx": ">>1"}[t]
                if t != "rrx":
                    eat = True
            elif f == "p" and t == "-":
                self.expr += "-"
                eat = True
            elif f == "r":
                if not eat and self.expr:
                    self.expr += "+"
                self.expr += "$" + t
                eat = False
                self.regs.append(t)

class ArmTarget(Target):
    '''
    A wrapper for an ARM Linux GDB Inferior.
    '''
    def _getInstruction(self, gdbstr):
        return ArmInstruction(gdbstr)
