import os
import re
import glob
import mmap
import shutil
import signal
import struct
import logging
import tempfile
import subprocess
import contextlib

from io import BytesIO


l = logging.getLogger("archr.analyzers.qemu_tracer")

from . import ContextAnalyzer
from .. import _angr_available

if _angr_available:
    import angr

from ..utils import filter_strace_output, get_file_maps


class QEMUTracerError(BaseException):
    pass


class QEMUBBLTrace:
    _element_size = 8  # bytes
    _mmaped_trace = None
    _trace_file_fh = None
    _trace_len = None

    def __del__(self):
        self._mmaped_trace.close()
        self._trace_file_fh.close()

    def __getitem__(self, index):
        if index < 0:
            index = self._trace_len + index

        return struct.unpack_from("<Q", self._mmaped_trace, index * self._element_size)[0]

    def __init__(self, trace_file, trace_len):
        self._trace_file_fh = open(trace_file, "rb")
        self._mmaped_trace = mmap.mmap(self._trace_file_fh.fileno(), 0, mmap.MAP_PRIVATE, mmap.PROT_READ)
        self._trace_len = trace_len

    def __len__(self):
        return self._trace_len

    def count(self, element):
        count = 0
        for curr_index in range(0, self._trace_len):
            if self[curr_index] == element:
                count = count + 1

        return count

    def index(self, element, start=0, stop=None):
        stop = stop if stop else self._trace_len
        for curr_index in range(start, stop):
            if self[curr_index] == element:
                return curr_index

        raise ValueError(f"{element} is not in trace")


class QemuTraceResult:
    # results
    returncode = None
    signal = None
    crashed = None
    timed_out = None

    # introspection
    trace = None
    mapped_files = None
    crash_address = None
    base_address = None
    image_base = None
    entry_point = None
    magic_contents = None
    halfway_core_path = None
    core_path = None
    taint_fd = None

    def tracer_technique(self, **kwargs):
        return angr.exploration_techniques.Tracer(self.trace, crash_addr=self.crash_address, **kwargs)


_trace_old_re = re.compile(rb"Trace (.*) \[(?P<addr>.*)\].*")
_trace_new_re = re.compile(rb"Trace (.*) \[(?P<something1>.*)\/(?P<addr>.*)\/(?P<flags>.*)\].*")


class QEMUTracerAnalyzer(ContextAnalyzer):
    REQUIRED_IMPLANT = "shellphish_qemu"

    def __init__(
        self,
        target,
        timeout=10,
        ld_linux=None,
        ld_preload=None,
        library_path=None,
        seed=None,
        qemu_args=None,
        shell=False,
        **kwargs,
    ):
        super().__init__(target, **kwargs)
        self.timeout = timeout
        self.ld_linux = ld_linux
        self.ld_preload = ld_preload
        self.library_path = library_path
        self.seed = seed
        self.hackbind = qemu_args.pop("hackbind", False) if qemu_args else False
        self.hackproc = qemu_args.pop("hackproc", False) if qemu_args else False
        self.hacksysinfo = qemu_args.pop("hacksysinfo", False) if qemu_args else False
        self.execve = qemu_args.pop("execve", False) if qemu_args else False
        self.mmap_base = qemu_args.pop("mmap_base", False) if qemu_args else False
        self.shell = shell

    def pickup_env(self):
        for e in self.target.target_env:
            key, value = e.split("=", 1)
            if key == "LD_PRELOAD" and self.ld_preload is None:
                self.ld_preload = value
            if key == "LD_LIBRARY_PATH" and self.library_path is None:
                self.library_path = value

    @contextlib.contextmanager
    def _target_mk_tmpdir(self):
        tmpdir = tempfile.mktemp(prefix="/tmp/tracer_target_")
        self.target.run_command(["mkdir", tmpdir]).wait()
        self.target.run_command(["chmod", "777", tmpdir]).wait()
        try:
            yield tmpdir
        finally:
            self.target.run_command(["rm", "-rf", tmpdir])

    @staticmethod
    @contextlib.contextmanager
    def _local_mk_tmpdir():
        tmpdir = tempfile.mkdtemp(prefix="/tmp/tracer_local_")
        try:
            yield tmpdir
        finally:
            with contextlib.suppress(FileNotFoundError):
                shutil.rmtree(tmpdir)

    @staticmethod
    def line_iter(content):
        io = BytesIO(content)
        for line in io:
            yield line.strip()

    @contextlib.contextmanager
    def fire_context(
        self,
        record_trace=True,
        record_magic=False,
        save_core=False,
        record_file_maps=False,  # pylint: disable=arguments-differ
        crash_addr=None,
        trace_bb_addr=None,
        taint=None,
        **kwargs,
    ):  # pylint:disable=arguments-differ
        with self._target_mk_tmpdir() as tmpdir:
            tmp_prefix = tempfile.mktemp(dir="/tmp", prefix="tracer-")
            target_trace_filename = tmp_prefix + ".trace" if record_trace else None
            target_magic_filename = tmp_prefix + ".magic" if record_magic else None
            local_core_filename = tmp_prefix + ".core" if save_core else None
            local_halfway_core_filename = (
                tmp_prefix + f".halfway_{hex(crash_addr[0])}_{crash_addr[1]}.core" if crash_addr else None
            )

            target_cmd = self._build_command(
                trace_filename=target_trace_filename,
                magic_filename=target_magic_filename,
                coredump_dir=tmpdir,
                crash_addr=crash_addr,
                start_trace_addr=trace_bb_addr,
                taint=taint,
            )
            kwargs["shell"] = self.shell

            l.debug("launch QEMU with command: %s", " ".join(target_cmd))
            r = QemuTraceResult()

            try:
                with self.target.flight_context(target_cmd, timeout=self.timeout, result=r, **kwargs) as flight:
                    yield flight
            except subprocess.TimeoutExpired:
                r.timed_out = True
                self.target.restart()
            else:
                r.timed_out = False
                r.returncode = flight.process.returncode

                # did a crash occur?
                if r.returncode in [139, -11]:
                    r.crashed = True
                    r.signal = signal.SIGSEGV
                elif r.returncode == [132, -9]:
                    r.crashed = True
                    r.signal = signal.SIGILL

            l.debug(
                "Qemu tracer returned with code=%s timed_out=%s crashed=%s signal=%s",
                r.returncode,
                r.timed_out,
                r.crashed,
                r.signal,
            )

            if local_core_filename or crash_addr:
                # choose the correct core dump to retrieve
                with self._local_mk_tmpdir() as local_tmpdir:
                    self.target.retrieve_into(tmpdir, local_tmpdir)
                    target_cores = glob.glob(os.path.join(local_tmpdir, "*", "qemu_*.core"))
                    tmp_crash_core_path = None
                    tmp_halfway_core_path = None

                    for x in target_cores:
                        if "crash" in x.rsplit("_")[-1]:
                            tmp_crash_core_path = x
                        if "coreaddr" in x.rsplit("_")[-1]:
                            tmp_halfway_core_path = x

                    if r.crashed and tmp_crash_core_path is None and len(target_cores) == 1:
                        tmp_crash_core_path = target_cores[0]

                    # sanity check core dumps
                    if save_core and not tmp_crash_core_path:
                        raise QEMUTracerError(
                            "the target didn't crash inside qemu or no corefile was created!"
                            + "Make sure you launch it correctly!\n"
                            + "command: %s" % " ".join(target_cmd)
                        )
                    if crash_addr and not tmp_halfway_core_path:
                        raise QEMUTracerError(
                            "the target didn't generate a halfway core file!" + "command: %s" % " ".join(target_cmd)
                        )

                    if local_core_filename and tmp_crash_core_path:
                        shutil.move(tmp_crash_core_path, local_core_filename)
                    if local_halfway_core_filename and tmp_halfway_core_path:
                        shutil.move(tmp_halfway_core_path, local_halfway_core_filename)
                    r.core_path = local_core_filename
                    r.halfway_core_path = local_halfway_core_filename

            if target_trace_filename:
                temp_trace_file = tempfile.NamedTemporaryFile()
                self.target.copy_file(target_trace_filename, temp_trace_file.name)
                trace_fh = open(temp_trace_file.name, "rb")

                # Find where qemu loaded the binary. Primarily for PIE
                try:
                    # the image base is the first mapped address in the page dump following the log line 'guest_base'
                    for t in trace_fh:
                        if t.startswith(b"guest_base"):
                            # iterate to the appropriate line
                            next(trace_fh)
                            next(trace_fh)
                            t = next(trace_fh)
                            # parse out the first line
                            r.image_base = int(t.split(b"-")[0], 16)
                            break

                    r.base_address = int(
                        next(t.split()[1] for t in trace_fh if t.startswith(b"start_code")), 16
                    )  # pylint:disable=stop-iteration-return

                    # for a dynamically linked binary, the entry point is in the runtime linker
                    # in this case it can be useful to keep track of the entry point
                    r.entry_point = int(next(t.split()[1] for t in trace_fh if t.startswith(b"entry")), 16)
                except StopIteration as e:
                    raise QEMUTracerError(
                        "The trace does not include any data. Did you forget to chmod +x the binary?"
                    ) from e

                # record the trace
                _trace_re = _trace_old_re if self.target.target_os == "cgc" else _trace_new_re
                bbl_trace_file = tempfile.NamedTemporaryFile(dir="/tmp", prefix="bbl-trace-")
                bbl_trace_fh = open(bbl_trace_file.name, "wb")
                bbl_trace_len = 0
                strace_lines = []
                prev_strace_entry = ""
                for entry in trace_fh:
                    if entry.startswith(b"Trace "):
                        bbl_trace_fh.write(struct.pack("<Q", int(_trace_re.match(entry).group("addr"), 16)))
                        bbl_trace_len += 1
                    elif record_file_maps:
                        curr_entry = entry.decode("utf-8").strip()
                        if "page layout changed following target_mmap" in curr_entry:
                            prev_strace_entry = curr_entry.replace("page layout changed following target_mmap", "")
                        else:
                            if re.match("^ = |^= ", curr_entry) and prev_strace_entry:
                                ret_val = filter_strace_output([prev_strace_entry + curr_entry])
                            else:
                                ret_val = filter_strace_output([curr_entry])

                            if ret_val:
                                strace_lines.append(ret_val[-1].strip())
                                prev_strace_entry = ""
                    elif r.crashed and entry.startswith(b"qemu: last read marker was read through fd:"):
                        # grab the taint_fd
                        r.taint_fd = int(re.search(rb"\[(\d+)\]", entry).group(1))
                        l.debug("Detected the tainted fd to be %s", r.taint_fd)

                lastline = entry
                bbl_trace_fh.close()
                trace_fh.close()
                temp_trace_file.close()
                r.trace = QEMUBBLTrace(bbl_trace_file.name, bbl_trace_len)

                if r.crashed:
                    if r.taint_fd is None and self.target.target_os != "cgc":
                        l.error(
                            "Unexpected status line from qemu tracer. Cannot get the last read marker to set taint_fd. "
                            "Please make sure you are using the latest shellphish-qemu."
                        )

                    # grab the faulting address
                    if lastline.startswith(b"Trace") or lastline.find(b"Segmentation") == -1:
                        l.warning(
                            "Trace return code was less than zero, but the last line of the trace does not"
                            "contain the uncaught exception error from qemu."
                            "If using an older version of shellphish_qemu try using 'ulimit -Sc 0' or "
                            "updating to a newer version of shellphish_qemu."
                        )
                    r.crash_address = int(lastline.split(b"[")[1].split(b"]")[0], 16)
                    l.debug("Detected the crashing address at %s", hex(r.crash_address))

                l.debug("Trace consists of %d basic blocks", len(r.trace))

                if record_file_maps:
                    r.mapped_files = get_file_maps(strace_lines)

                # remove the trace file on the target
                self.target.remove_path(target_trace_filename)

            if target_magic_filename:
                r.magic_contents = self.target.retrieve_contents(target_magic_filename)
                if len(r.magic_contents) != 0x1000:
                    raise QEMUTracerError("Magic content read from QEMU improper size, should be a page in length")

                # remove the magic file on the target
                self.target.remove_path(target_magic_filename)

    @staticmethod
    def qemu_variant(target_os, target_arch, record_trace):
        """
        Need to know if we're tracking or not, specifically for what cgc qemu to use.
        """

        if target_os == "cgc":
            suffix = "tracer" if record_trace else "base"
            qemu_variant = "shellphish-qemu-cgc-%s" % suffix
        else:
            qemu_variant = "shellphish-qemu-linux-%s" % target_arch

        return qemu_variant

    def _build_command(
        self,
        trace_filename=None,
        magic_filename=None,
        coredump_dir=None,
        report_bad_args=None,
        crash_addr=None,
        start_trace_addr=None,
        taint=None,
    ):
        """
        Here, we build the tracing command.
        """

        #
        # First, the implant invocation
        #

        qemu_variant = self.qemu_variant(self.target.target_os, self.target.target_arch, trace_filename is not None)
        qemu_path = os.path.join(self.target.tmpwd, "shellphish_qemu", qemu_variant)
        fire_path = os.path.join(self.target.tmpwd, "shellphish_qemu", "fire")
        cmd_args = [fire_path, qemu_path]
        if coredump_dir:
            cmd_args += ["-C", coredump_dir]
        if crash_addr:
            cmd_args += ["-A", "0x{:x}:{}".format(*crash_addr)]
        if start_trace_addr:
            cmd_args += ["-T", "0x{:x}:{}".format(*start_trace_addr)]
        if taint:
            cmd_args += ["-M", taint.hex()]

        #
        # Next, we build QEMU options.
        #

        # hardcode an argv[0]
        # cmd_args += [ "-0", program_args[0] ]

        # record trace
        if trace_filename:
            flags = "nochain,exec,page,strace" if "cgc" not in qemu_variant else "exec"
            cmd_args += ["-d", flags, "-D", trace_filename]
        else:
            if "cgc" in qemu_variant:
                cmd_args += ["-enable_double_empty_exiting"]

        # save CGC magic page
        if magic_filename:
            if "cgc" not in qemu_variant:
                raise QEMUTracerError("Specified magic page dump on non-cgc architecture")
            cmd_args += ["-magicdump", magic_filename]

        if self.seed is not None:
            cmd_args.append("-seed")
            cmd_args.append(str(self.seed))

        if report_bad_args:
            cmd_args += ["-report_bad_args"]

        # Memory limit option is only available in shellphish-qemu-cgc-*
        if "cgc" in qemu_variant:
            cmd_args += ["-m", "8G"]

        if "cgc" not in qemu_variant and "LD_BIND_NOW=1" not in self.target.target_env:
            cmd_args += ["-E", "LD_BIND_NOW=1"]

        if self.ld_preload:
            cmd_args += ["-E", "LD_PRELOAD=" + self.ld_preload]

        if self.library_path and not self.ld_linux:
            cmd_args += ["-E", "LD_LIBRARY_PATH=" + self.library_path]

        # Greenhouse related hacks
        if self.hackbind:
            cmd_args += ["-hackbind"]
        if self.hackproc:
            cmd_args += ["-hackproc"]
        if self.hacksysinfo:
            cmd_args += ["-hacksysinfo"]
        if self.execve:
            cmd_args += ["-execve", self.execve]
        if self.mmap_base:
            cmd_args += ["-S", hex(self.mmap_base)]

        # now set up the loader
        if self.ld_linux:
            cmd_args += [self.ld_linux]
            if self.library_path:
                cmd_args += ["--library-path", self.library_path]

        # Now, we add the program arguments.
        cmd_args += ["--"]  # separate QEMU arguments and target arguments
        cmd_args += self.target.target_args

        return cmd_args
