import subprocess
import struct
import re
import logging
import cle
import io

from . import strace_parser

l = logging.getLogger("archr.utils")


def parse_ldd(mem_map_str):
    entries = [l.strip() for l in mem_map_str.decode("utf-8").splitlines()]
    parsed = {}
    for entry in entries:
        if "=>" in entry:
            libname, paren_addr = entry.split("=>")[1].split()
        else:
            libname, paren_addr = entry.split()
        libaddr = int(paren_addr.strip("()"), 16)
        parsed[libname] = libaddr
    return parsed


def parse_proc_maps(proc_str):
    entries = [l.strip() for l in proc_str.splitlines()]
    parsed = {}
    for entry in entries:
        what = entry.split()[-1].decode("utf-8")
        addr_range = entry.split()[0]
        start, end = addr_range.split(b"-")
        if what in parsed:
            continue
        if what.startswith("/"):
            parsed[what] = int(start, 16)
        elif what.startswith("["):
            parsed[what] = int(start, 16)
            parsed[what.rstrip("]") + "-end]"] = int(end, 16)
    return parsed


def lib_dependencies(filepath):
    mem_map_str, _ = subprocess.Popen(["ldd", filepath], stdout=subprocess.PIPE).communicate()
    return [lib for lib in parse_ldd(mem_map_str) if lib != "linux-vdso.so.1"]


def hook_entry(binary, asm_code=None, bin_code=None):
    main_bin = io.BytesIO(binary)
    b = cle.Loader(main_bin, auto_load_libs=False, perform_relocations=False)
    start_addr = b.main_object.addr_to_offset(b.main_object.entry)
    arch = b.main_object.arch
    if arch.name in ("ARMHF", "ARMEL") and arch.is_thumb(start_addr):  # OMG, thumb mode is a disaster
        start_addr &= ~1  # recover the real address
        main_bin.seek(start_addr)
        padding = (4 - (start_addr + 8) % 4) % 4  # we hardcode the shellcode so that its length is 8

        # we can' use arch.asm here because the shellcode THUMB, 8+padding-4 because the shellcode has length 8+padding,
        # we also need to take into account that in arm, pc points to two instructions ahead, which is 4 bytes in thumb mode
        main_bin.write(b"xF\x00\xf1" + struct.pack("<H", 8 + padding - 4) + b"\x00G" + b"A" * padding)

        # now place our payload after this mini shellcode
        start_addr += 8 + padding
    main_bin.seek(start_addr)
    main_bin.write(b.main_object.arch.asm(asm_code) if asm_code else bin_code)
    main_bin.seek(0)
    return main_bin.read()


def filter_strace_output(lines):
    """
    a function to filter QEMU logs returning only the strace entries

    Parameters
    ----------
    lines : list
        a list of strings representing the lines from a QEMU log/trace.

    Returns
    -------
    list
        a list of strings representing only the strace log entries
        the entries will also be cleaned up if a page dump occurs in the middle of them
    """

    # we only want the strace lines, so remove/ignore lines that start with the following:
    line_starts = [
        r"^[\d,a-f]{16}-",  # pylint: disable=anomalous-backslash-in-string
        r"^[\d,a-f]{8}-",  # pylint: disable=anomalous-backslash-in-string
        "^page",
        "^start",
        "^host",
        "^Locating",
        "^guest_base",
        "^end_",
        "^brk",
        "^entry",
        "^argv_",
        "^env_",
        "^auxv_",
        "^Trace",
        "^--- SIGSEGV",
        "^qemu",
    ]
    filter_string = "|".join(line_starts)

    filtered = []
    prev_line = ""
    for line in lines:
        if re.match(filter_string, line):
            continue
        # workaround for https://gitlab.com/qemu-project/qemu/-/issues/654
        if re.search("page layout changed following target_mmap", line):
            prev_line = line.replace("page layout changed following target_mmap", "")
            continue
        # The strace_parser does not properly handle _newselect strace lines (It's format differs significantly).
        # so, workaround by simply removing those lines. We currently do not use them.
        if re.search("newselect", line):
            continue  # ignore _newselect syscalls
        if re.match("^ = |^= ", line):
            line = prev_line + line

        filtered.append(line)
    return filtered


def get_file_maps(strace_log_lines):
    """
    a function to return a mapping of filenames and associated mmaped addressed for process under QEMU
    this is accomplished by tracking the filenames through file-desciptors across mmap() calls.

    Parameters
    ----------
    strace_log_lines : list
        a list of strings representing only the strace info logged by QEMU

    Returns
    -------
    dict
        a dictionary of filenames and mmapped addresses associated with each file

    """
    files = {"open": {}, "closed": {}, "sockets": []}

    entries = strace_parser.parse(strace_log_lines)
    entries = [entry for entry in entries if entry.syscall in ("openat", "open", "socket", "mmap", "mmap2", "close")]

    for entry in entries:
        # for an openat, create a dict entry for the file descriptor
        # the entry should be a tuple of the filename, and mmaps (initially empty)
        if entry.syscall == "openat":
            fd = entry.syscall.result
            # only care about file descriptors other than STDIN,STDOUT,STDERR
            # also ignore errors
            if fd >= 3:
                # use only the base filename
                filename = entry.syscall.args[1].split("/")[-1]
                # tracking if an executable page was ever mapped from the file descriptor
                files["open"][fd] = [filename, []]

        elif entry.syscall == "open":
            fd = entry.syscall.result
            # only care about file descriptors other than STDIN,STDOUT,STDERR
            # also ignore errors
            if fd >= 3:
                # use only the base filename
                filename = entry.syscall.args[0].split("/")[-1]
                # tracking if an executable page was ever mapped from the file descriptor
                files["open"][fd] = [filename, []]

        elif entry.syscall == "socket":
            fd = entry.syscall.result
            if fd >= 3:
                files["sockets"].append(fd)

        # if a file descriptor is closed, we need to remove it from the open files dictionary
        # we want to track the mmaps, so move it to 'closed' by file name since the file descriptor can be re-used.
        elif entry.syscall == "close":
            fd = entry.syscall.args[0]
            # only care about file descriptors other than STDIN,STDOUT,STDERR
            if fd >= 3:
                if fd in files["sockets"]:
                    files["sockets"].remove(fd)
                else:
                    filename = files["open"][fd][0]
                    mmaps = files["open"][fd][1]

                    # if we never mapped any pages, then we don't care about it.
                    if mmaps:
                        # otherwise move to 'closed'
                        files["closed"][filename] = mmaps

                    del files["open"][fd]

        # we can use the file descriptor to look up the dict entry to update the mmaps
        # TODO: track sizes which should be capturable from the mmap arguments
        elif entry.syscall in ("mmap", "mmap2"):
            # only care about valid file descriptors
            fd = entry.syscall.args[4]
            if fd >= 3:
                files["open"][fd][1].append(entry.syscall.result)

    # lets "close" everything that never got closed
    for fd, (filename, mmaps) in files["open"].items():
        files["closed"][filename] = mmaps

    return files["closed"]


def hook_addr(binary, addr, asm_code=None, bin_code=b""):
    main_bin = io.BytesIO(binary)
    loader = cle.Loader(main_bin, auto_load_libs=False, perform_relocations=False)
    offset = loader.main_object.addr_to_offset(addr)
    main_bin.seek(offset)
    main_bin.write(loader.main_object.arch.asm(asm_code) if asm_code else bin_code)
    main_bin.seek(0)
    return main_bin.read()
