### 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.

These objects are designed to be used when an AddressSanitizer log is passed 
as an argument to the 'exploitable' command. They are used by a specified user, 
who graciously contributed the code, but have not been tested.

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

from lib.tools import AttrDict, memoized

from lib.gdb_wrapper.x86 import Backtrace, Frame, Target, ProcMaps

class ASanFrame(Frame):
    '''
    Wrapper for gdb.Frame. Content for this frame is generated from 
    AddressSanitizer log data rather than the GDB API.
    '''
    def __init__(self, target, position=None):
        self.target = target
        self.position = position
        self.blacklisted = False
        if position is not None:
            self._pc = target.asan_stack[position].addr
            self._name = target.asan_stack[position].name
            self.offset = target.asan_stack[position].off
            self.mapped_region = target.procmaps().findByAddr(self._pc)

    @classmethod
    def create(cls, target, addr, name, offset):
        o = cls(target)
        o._pc = addr
        o._name = name
        o.offset = offset
        o.mapped_region = target.procmaps().findByAddr(addr)
        return o

    def name(self):
        return self._name

    def pc(self):
        return self._pc

    def older(self):
        try:
            return self.target.asan_stack[self.position+1]
        except IndexError:
            raise RuntimeError()

    def unwind_stop_reason(self):
        return 0

    def type(self):
        return 0

    def __getattribute__(self, name):
        return object.__getattribute__(self, name)

class ASanBacktrace(Backtrace):
    '''
    A backtrace composed of a list of GDB Frames (ordered from innermost to
    outermost) and other attributes generated from an ASan log.
    '''

    def _next_frame(target=None, frame=None, i=None):
        '''
        Gets the next frame (the frame after 'frame').
        If called without params (or if target is None, at least) then True
        is returned (so that the backtrace loop in the calling function is 
        entered).
        '''
        if not target:
            return True
        if i >= len(target.asan_stack):
            return None
        return ASanFrame(target, i)

class ASanProcMaps(ProcMaps):
    '''
    A list of process address mappings. This object should be instantiated 
    when the inferior is not running, but an ASan log file has been
    specified. 
    '''

    def __init__(self): 
        self._common_init()
        self._add_by_target()

    def add_file(self, filepath, offset):
        '''
        For the ELF file at filepath, adds a mapping to self and loads debug 
        symbols into GDB. offset is added to the start and end addresses for
        the mapping (offset is the offset that the ELF is loaded to)
        ''' 
        bn = os.path.basename(filepath)
        if bn in self._files:
            return
        elf_sects = read_elf_sects(filepath)
        sects = {}
        for sect, (start, size) in elf_sects.items():
            start += offset
            sects[sect] = start
            ad = AttrDict(start=start, end=start + size, size=size, offset=0, name=bn, sect=sect)
            self.append(ad)
        cmd = "add-symbol-file '{}' {:#x} {}".format(filepath, sects[".text"],
                                                     " ".join("-s {} {:#x}".format(nm, st)
                                                              for (nm, st) in sects.items()
                                                              if nm != ".text"))
        gdb.execute(cmd, False, True)
        self._files.add(bn)

    def _add_by_target(self, num_initial_stacks=0):
        '''
        Searches for thread stacks and adds corresponding ranges to self.
        Also sets self.tgt_img to the name associated with the first thread 
        stack.

        The first num_initial_stacks sections in the target (such as a core 
        file) are treated as thread stacks. This logic is currently only
        used by the ASan logic, which is used by a specific user and
        largely untested.
        '''

        # Collect thread stack info
        mapstr = str(gdb.execute("info target", False, True))
        ranges = []
        fn = None
        for m in self._re_info_target.finditer(mapstr):
            if m.group("file") is not None:
                fn = os.path.basename(m.group("file"))
            else:
                lib = m.group("lib")
                if lib:
                    lib = os.path.basename(lib)
                ranges.append((int(m.group("start"), 16), int(m.group("end"), 16), fn, lib, m.group("section")))

        # Append thread stack info to self
        for i, (st, en, fn, lib, sect) in enumerate(ranges):
            if i < num_initial_stacks:
                # We could correlate thread stack pointer to one of the ranges,
                # but we'll just assume for now that for num_initial_stacks, 
                # the first num_intial_stacks sections in the core are the stacks
                name = "[stack]"
            elif lib is not None:
                name = lib
            else:
                name = fn
            ad = AttrDict(start=st, end=en, size=(en-st), offset=0, name=name, sect=sect)
            self.append(ad)
            #gdb.write("{}: {:#x}-{:#x} {}".format(file, ad.start, ad.end, sect))
            self._files.add(os.path.basename(fn))
        if ranges:
            self.tgt_img = ranges[0][2]


class ASanTarget(Target):
    '''
    A wrapper for a Linux GDB Inferior. This wrapper enhances the state of 
    the Linux GDB Inferior with output from AddressSanitizer.
    Note: This code was contributed by BlackBerry, and has not been thoroughly
    tested outside of their environment.
    ''' 
    _re_asan_bt = re.compile(r"""^(\033\133[0-9]+m)* # ANSI colour code
                                  (?P<all>\s*\#(?P<frame>[0-9]+)\s*
                                  (?P<addr>0x[A-Fa-f0-9]+)\s*\()
                                  (?P<img>.*?)
                                  \+(?P<offset>0x[A-Fa-f0-9]+)\)$""", re.VERBOSE|re.MULTILINE)
#==19027== ERROR: AddressSanitizer: SEGV on unknown address 0x000000000000 (pc 0x7f0fa35f1bf1 sp 0x7fffbcc88a48 bp 0x7fffbcc89290 T0)
#==16758== ERROR: AddressSanitizer: attempting double-free on 0x60040000dff0:
#==16769== ERROR: AddressSanitizer: heap-use-after-free on address 0x60040000dff0 at pc 0x4025b1 bp 0x7ffff1dbe6c0 sp 0x7ffff1dbe6b8
#READ of size 1 at 0x60040000dff0 thread T0
    _re_asan_fault = re.compile(r"""^(\033\133[0-9]+m)* # ANSI colour code
                                     =+[0-9]+=+\s*ERROR:\s*AddressSanitizer:\s*
                                     (attempting\s)?(?P<desc>[A-Za-z0-9_-]+)\s
                                     on\s((unknown\s)?address\s)?(?P<fault>0x[A-Fa-f0-9]+)
                                     (:$ | \s(\(|at\s)
                                         pc\s(?P<pc>0x[A-Fa-f0-9]+)\s
                                         (?P<bspid1>[bs]p)\s(?P<bsp1>0x[A-Fa-f0-9]+)\s
                                         (?P<bspid2>[bs]p)\s(?P<bsp2>0x[A-Fa-f0-9]+)([\r\n]+
                                     (\033\133[0-9]+m)* # ANSI colour code
                                     (?P<operation>[A-Za-z0-9_-]+))?)""", re.VERBOSE|re.MULTILINE)
    _re_asan_report_sym = re.compile(r"^(__asan_report_error|__asan::ReportDoubleFree)")
    _re_asan_printloop = re.compile(r"^(__sanitizer::RawWrite|__sanitizer::Abort|[A-Za-z_0-9]*puts)")
    _re_symline_trim = re.compile(r" starts at address.*\s*")

    _asan_reasons = ["double-free", "bad-free", "alloc-dealloc-mismatch", "unknown-crash",
                     "heap-buffer-overflow", "global-buffer-overflow", "stack-use-after-scope",
                     "use-after-poison", "stack-use-after-return", "stack-buffer-overflow",
                     "initialization-order-fiasco", "stack-buffer-underflow", "heap-use-after-free", "SEGV"]

    def __init__(self, asan_output, bt_limit=0):
        self.__memo__ = {"isPossibleStackCorruption()": False,
                         "isStackCorruption()": False,
                         "isStackOverflow()": False,
                         "si_signo()": 11}
        if not asan_output:
            raise GdbWrapperError("no ASan data to analyze")

        # symbolize asan_message
        self.asan_stack = []
        out = []
        last = 0
        all_frames = []
        maps = self.procmaps()
        i = 0
        for m in self._re_asan_bt.finditer(asan_output):
            frame, addr, img, offset = m.group("frame", "addr", "img", "offset")
            frame = int(frame)
            addr = int(addr, 16) #+ 1
            if img:
                maps.add_file(img, addr - offset)
            out.append(asan_output[last:m.end("all")])
            all_frames.append((frame, addr, offset, img, len(out)))
            out.append(None)
            last = m.end()
	    
            i += 1
            if i >= bt_limit:
                break

        if not all_frames:
            raise GdbWrapperError("No frames found in address sanitizer log")

        out.append(asan_output[last:])
        frame = -1
        for num, addr, offset, img, outpos in all_frames:
            region = maps.findByAddr(addr)
            symbol = gdb.execute("info symbol {:#x}".format(addr), False, True)
            symline = gdb.execute("info line *{:#x}".format(addr), False, True)
            if symline and symline.startswith("Line"):
                symline = "\n\t{}".format(self._re_symline_trim.sub("", symline))
            else:
                symline = ""
            symbol_m = self._re_gdb_info_sym.search(symbol)
            if img:
                lib = img
            elif region:
                lib = region.name
            else:
                lib = None
            if symbol_m is None:
                sym = None
                off = offset
            else:
                sym = symbol_m.group("sym")
                off = int(symbol_m.group("off"))
            if frame == -1:
                self.asan_pc_img = lib, offset
            if frame is not None and num > frame:
                frame = num
                if lib:
                    lib = os.path.basename(lib)
                self.asan_stack.append(AttrDict(addr=addr,
                                                lib=lib,
                                                off=off,
                                                name=sym))
            else:
                frame = None
            out[outpos] = "{}){}".format(ASanFrame.create(self, addr, sym, off).terse(), symline)
        asan_output = "".join(out)
        gdb.write(asan_output)
        gdb.flush()
        # parse ASAN's analysis
        m = self._re_asan_fault.search(asan_output)
        self.__memo__["si_addr()"] = int(m.group("fault"), 16)
        self.asan_reason = m.group("desc")
        if self.asan_reason == "double-free":
            self.__memo__["pc()"] = self.asan_stack[1].addr
            self.__memo__["stack_pointer()"] = None # what to do? ....
        else:
            self.__memo__["pc()"] = int(m.group("pc"), 16)
            if m.group("bspid1") == "sp":
                self.__memo__["stack_pointer()"] = int(m.group("bsp1"), 16)
            else:
                self.__memo__["stack_pointer()"] = int(m.group("bsp2"), 16)
            if self.asan_reason != "SEGV":
                self.asan_operation = m.group("operation")

    @memoized
    def current_instruction(self):
        img, addr = getattr(self, "asan_pc_img", (None, self.pc()))
        if img and os.path.basename(img) != getattr(self.procmaps(), "tgt_img", None):
            gdb.execute("file {}".format(img), False, True)
        try:
            gdbstr = gdb.execute("x/i {:#x}".format(addr), False, True).splitlines()[0]
            self._getInstruction(gdbstr)
        except RuntimeError:
            return None

    @memoized
    def procmaps(self):
        return ASanProcMaps()
