#! /usr/bin/env python3
############################################################################
# tools/pynuttx/nxelf/elf.py
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.  The
# ASF licenses this file to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance with the
# License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
# License for the specific language governing permissions and limitations
# under the License.
#
############################################################################


import bisect
import functools
import hashlib
import importlib
import logging
import os
import pickle
import sys
import time
from enum import IntEnum
from typing import Tuple

from .macros import Macro

try:
    import cxxfilt
    import lief
    from construct import (
        Array,
        Construct,
        Float16l,
        Float32l,
        Float64l,
        Int8sl,
        Int8ul,
        Int16sl,
        Int16ul,
        Int32sb,
        Int32sl,
        Int32ub,
        Int32ul,
        Int64sb,
        Int64sl,
        Int64ub,
        Int64ul,
        Struct,
    )
    from elftools.elf.elffile import ELFFile
    from elftools.elf.sections import SymbolTableSection
except ModuleNotFoundError as e:
    print(f"Error:{e}.\nPlease execute the following command to install dependencies:")
    print("pip install construct pyelftools cxxfilt lief ")

log = logging.getLogger(__name__)


class TypeConflictError(Exception):
    """
    Symbols have different definitions
    """

    pass


class Types:
    def __init__(self, tag):
        self.types = {}
        self.tag = tag
        self.result = dict()

    def set_type(self, die):
        if "DW_AT_name" not in die.attributes:
            return

        name = die.attributes["DW_AT_name"].value.decode("utf-8")
        if name not in self.types:
            self.types[name] = set()

        self.types[name].add(die)

    def get_types(self, type_name):
        if type_name in self.types:
            sets = self.types[type_name]
            return sets
        else:
            return None

    def set_result(self, type_name, result):
        if len(result) != 1:
            raise TypeConflictError(
                f"Multiple different definitions or values ​​exist for a symbol: {type_name} {result}"
            )

        result = result.pop()
        if type_name in self.types:
            self.result[type_name] = result
        return result

    def get_result(self, type_name):
        if type_name in self.result:
            return self.result[type_name]
        return None


class ELFParser:
    """
    ELF file parser class for extracting the following information from ELF files:
    - Symbol addresses
    - Structure definitions
    - Enumeration type definitions
    - Enumeration values

    Main functionality:
    1. Get structure definitions
        elf_parser = ELFParser("nuttx")
        struct = elf_parser.get_type("file_operations")  # Returns construct.Struct object
        print(struct._subcons)

        # result:
        Container:
            open = <Renamed open <FormatField>>
            close = <Renamed close <FormatField>>
            read = <Renamed read <FormatField>>
            write = <Renamed write <FormatField>>
            seek = <Renamed seek <FormatField>>
            ioctl = <Renamed ioctl <FormatField>>
            mmap = <Renamed mmap <FormatField>>
            truncate = <Renamed truncate <FormatField>>
            poll = <Renamed poll <FormatField>>
            unlink = <Renamed unlink <FormatField>>

    2. Get symbol addresses
       addr = elf_parser.symbol_addr("_SeggerRTT")  # Returns symbol address

    3. Get enumeration type definitions
       enum = elf_parser.get_type("tstate_e")  # Returns construct.IntEnum object

    4. Get enumeration values
        value = elf_parser.enum_value("TSTATE_TASK_RUNNING")  # Returns integer value

    """

    def __init__(self, elf_path):
        self.elf = ELFFile(open(elf_path, "rb"))
        self.addr = list()
        self.types = dict()
        self.info = dict()
        self.symbol = dict()
        self.result = dict()
        self.dwarf = self.elf.get_dwarf_info()

        t = time.time()
        print("Parsing ELF file...")
        self.parse_header()
        self.parse_symbol()
        self.macro = Macro(elf_path)
        print(f"ELF file parsed in {time.time() - t:.1f} seconds")

    def parse_header(self):
        header = self.elf.header
        self.info["bitwides"] = (
            32 if header["e_ident"]["EI_CLASS"] == "ELFCLASS32" else 64
        )
        self.info["byteorder"] = (
            "little" if header["e_ident"]["EI_DATA"] == "ELFDATA2LSB" else "big"
        )
        self.info["arch"] = header["e_machine"]
        self.info["size_t"] = "uint%d" % self.info["bitwides"]

    def symbol_filter(self, symbol):
        if symbol["st_info"]["type"] != "STT_FUNC":
            return None
        if symbol["st_info"]["bind"] == "STB_WEAK":
            return None
        if symbol["st_shndx"] == "SHN_UNDEF":
            return None
        return symbol

    def parse_symbol(self):
        tables = [
            s
            for _, s in enumerate(self.elf.iter_sections())
            if isinstance(s, SymbolTableSection)
            and s.name == ".symtab"
            and s["sh_entsize"]
        ]

        for section in tables:
            for nsym, symbol in enumerate(section.iter_symbols()):
                try:
                    name = cxxfilt.demangle(symbol.name)
                except Exception:
                    name = symbol.name
                self.symbol[name] = symbol["st_value"]
                if symbol_filter := self.symbol_filter(symbol):
                    self.addr.append(
                        (symbol_filter["st_value"], symbol_filter["st_size"], name)
                    )

        self.addr.sort(key=lambda x: x[0])

    @functools.lru_cache(maxsize=None)
    def symbol_addr(self, name):
        if len(self.symbol.keys()) == 0:
            self.parse_symbol()

        if name not in self.symbol:
            return None

        return self.symbol[name]

    @functools.lru_cache(maxsize=None)
    def addr2symbol(self, address):
        """Get function name from address"""
        if len(self.addr) == 0:
            self.parse_symbol()

        if i := bisect.bisect_right(self.addr, address, key=lambda x: x[0]):
            addr, size, name = self.addr[i - 1]
            if addr <= address < addr + size:
                return name

        return None

    @functools.lru_cache(maxsize=None)
    def readstring(self, addr):
        for segment in self.elf.iter_segments():
            seg_addr = segment["p_paddr"]
            seg_size = min(segment["p_memsz"], segment["p_filesz"])
            if addr >= seg_addr and addr <= seg_addr + seg_size:
                data = segment.data()[addr - seg_addr :]
                data = data.split(b"\x00")[0]
                return data.decode("utf-8")
        return None

    @functools.lru_cache(maxsize=None)
    def read(self, addr, size):
        for segment in self.elf.iter_segments():
            seg_addr = segment["p_paddr"]
            seg_size = min(segment["p_memsz"], segment["p_filesz"])
            if addr >= seg_addr and addr + size <= seg_addr + seg_size:
                data = segment.data()
                start = addr - seg_addr
                return data[start : start + size]

    @functools.lru_cache(maxsize=None)
    def parse_base_type(self, die):
        name = die.attributes["DW_AT_name"].value.decode("utf-8")
        size = die.attributes["DW_AT_byte_size"].value

        basetypes = {
            "unsigned": {1: Int8ul, 2: Int16ul, 4: Int32ul, 8: Int64ul},
            "float": {2: Float16l, 4: Float32l, 8: Float64l},
            "signed": {1: Int8sl, 2: Int16sl, 4: Int32sl, 8: Int64sl},
        }

        if "unsigned" in name:
            return basetypes["unsigned"].get(size)
        elif "double" in name or "float" in name:
            return basetypes["float"].get(size)
        elif "_Bool" in name:
            return Int8ul
        elif "char" in name or "short" in name or "int" in name:
            return basetypes["signed"].get(size)

        raise ValueError(f"Unsupported base type: {name}")

    @functools.lru_cache(maxsize=None)
    def parse_array(self, die):
        nums = 0
        for child in die.iter_children():
            if "DW_AT_upper_bound" in child.attributes:
                nums = child.attributes["DW_AT_upper_bound"].value + 1
            elif "DW_AT_count" in child.attributes:
                nums = child.attributes["DW_AT_count"].value

        type_die = self.dwarf.get_DIE_from_refaddr(
            die.attributes["DW_AT_type"].value + die.cu.cu_offset
        )

        item_type = self.parse_die(type_die)

        def dynamic_array(ctx):
            # If array_field is specified during parse, array parsing uses the specified field as the length
            if hasattr(ctx, "_params") and "array_field" in ctx._params:
                field_name = ctx._params["array_field"]
                if hasattr(ctx, field_name):
                    return getattr(ctx, field_name)
            elif hasattr(ctx, "_params") and "array_length" in ctx._params:
                # If array_length is specified during parse, the array is parsed using the specified length.
                return ctx._params["array_length"]

            return nums

        """
            Usage example:
            struct example_s {
                int length;
                int buffer[0];
            };
            struct = elf.get_type("example")
            struct.parse(
                b"\x00\x00\x00\x04"  # length = 4
                b"\x01\x02\x03\x04",  # buffer = [1, 2, 3, 4]
                array_field="length"  # Use the length field to determine the size of the array
            )
        """

        return Array(dynamic_array, item_type)

    @functools.lru_cache(maxsize=None)
    def parse_typedef(self, die):
        type_attr = die.attributes["DW_AT_type"]
        die = self.dwarf.get_DIE_from_refaddr(type_attr.value + die.cu.cu_offset)
        return self.parse_die(die)

    @functools.lru_cache(maxsize=None)
    def parse_struct(self, die):
        members = dict()
        for child in die.iter_children():
            member_name = child.attributes["DW_AT_name"].value.decode("utf-8")
            member_type = child.attributes["DW_AT_type"].value
            type_die = self.dwarf.get_DIE_from_refaddr(member_type + die.cu.cu_offset)
            member_type = self.parse_die(type_die)
            members[member_name] = member_type

        struct = Struct(**members)
        return struct

    @functools.lru_cache(maxsize=None)
    def parse_enum(self, die) -> IntEnum:
        if die.tag != "DW_TAG_enumeration_type":
            raise ValueError(f"type is not enum: {die.tag}")

        enum = dict()
        name = die.attributes["DW_AT_name"].value.decode("utf-8")
        for child in die.iter_children():
            name = child.attributes["DW_AT_name"].value.decode("utf-8")
            value = child.attributes["DW_AT_const_value"].value
            enum[name] = value

        return IntEnum(name, enum)

    @functools.lru_cache(maxsize=None)
    def parse_enum_value(self, die):
        if die.tag != "DW_TAG_enumerator":
            raise ValueError(f"type is not enum: {die.tag}")

        name = die.attributes["DW_AT_name"].value.decode("utf-8")
        value = die.attributes["DW_AT_const_value"].value
        return name, value

    @functools.lru_cache(maxsize=None)
    def find_die_by_name(self, name):
        if name in self.types:
            return self.types[name]

        for CU in self.dwarf.iter_CUs():
            result = None
            for DIE in CU.iter_DIEs():
                if "DW_AT_name" not in DIE.attributes:
                    continue

                AT_name = DIE.attributes["DW_AT_name"].value.decode("utf-8")
                self.types[AT_name] = DIE
                if name == AT_name:
                    result = DIE

            if result:
                return result

        return None

    @functools.lru_cache(maxsize=None)
    def parse_die(self, die):
        if (
            "DW_AT_name" not in die.attributes
            and die.tag != "DW_TAG_pointer_type"
            and die.tag != "DW_TAG_array_type"
        ):
            return None

        tag_handlers = {
            "DW_TAG_structure_type": self.parse_struct,
            "DW_TAG_enumeration_type": self.parse_enum,
            "DW_TAG_base_type": self.parse_base_type,
            "DW_TAG_typedef": self.parse_typedef,
            "DW_TAG_pointer_type": lambda _: (
                Int32ul if self.info["bitwides"] == 32 else Int64ul
            ),
            "DW_TAG_array_type": self.parse_array,
        }

        if die.tag not in tag_handlers:
            raise ValueError(f"Unsupported type: {die.tag}")

        return tag_handlers[die.tag](die)

    @functools.lru_cache(maxsize=None)
    def get_type(self, type_name):
        if type_name in self.result:
            return self.result[type_name]

        die = self.find_die_by_name(type_name)
        if die is None:
            return None

        return self.parse_die(die)


class LiefELF:
    def __init__(self, filename):
        self.elf = lief.parse(filename)
        if not self.elf:
            raise BaseException(f"Failed to parse ELF file: {filename}")

        self.endian = (
            "l"
            if self.elf.abstract.header.endianness == lief.Header.ENDIANNESS.LITTLE
            else "b"
        )
        self.architecture = self.elf.abstract.header.architecture
        self.bits = 64 if self.elf.abstract.header.is_64 else 32

    def get_symbol(self, symbol):
        # Try to get LTO private symbol if not found.
        # This may still not work if LTO renames the symbol, or
        # it has different suffix.
        # Try to get symbol with .0 suffix too for some static variable defined in function scope, like g_statenames
        return (
            self.elf.get_symbol(symbol)
            or self.elf.get_symbol(f"{symbol}.lto_priv.0")
            or self.elf.get_symbol(f"{symbol}.0")
        )

    def read_symbol(
        self, symbol, struct: Construct = None
    ) -> Tuple[lief.Symbol, memoryview]:
        sym = self.get_symbol(symbol)
        if sym is None:
            return None

        data = self.read_from(sym.value, sym.size)
        if struct:
            data = struct.parse(data)
        return sym, data

    def read_from(self, addr, len=1) -> memoryview:
        for section in self.elf.sections:
            if section.type == lief.ELF.Section.TYPE.PROGBITS:
                off = addr - section.virtual_address
                if (
                    section.virtual_address
                    <= addr
                    < section.virtual_address + section.size
                    and section.size - off >= len
                ):
                    return section.content[off : off + len]

        for segment in self.elf.segments:
            if segment.type == lief.ELF.Segment.TYPE.LOAD:
                off = addr - segment.virtual_address
                if (
                    segment.virtual_address
                    <= addr
                    < segment.virtual_address + segment.virtual_size
                    and segment.virtual_size - off >= len
                ):
                    return segment.content[off : off + len]

        return None

    def read_string(self, addr) -> str:
        """Read const string from ELF file"""
        output = b""
        while True:
            c = self.read_from(addr, 1)
            if c == b"\0":
                break

            output += c.tobytes()
            addr += 1

        return output.decode("utf-8", errors="replace")

    def get_inttype(self) -> Construct:
        return {
            "32l": Int32sl,
            "32b": Int32sb,
            "64l": Int64sl,
            "64b": Int64sb,
        }.get(f"{self.bits}{self.endian}", Int32sl)

    def get_pointer_type(self) -> Construct:
        return {
            "32l": Int32ul,
            "32b": Int32ub,
            "64l": Int64ul,
            "64b": Int64ub,
        }.get(f"{self.bits}{self.endian}", Int32ul)

    def get_pointer_size(self):
        return 8 if self.elf.abstract.header.is_64 else 4


class AngrElf:
    def __init__(self, elf: str):
        try:
            self.angr = importlib.import_module("angr")
            self.capstone = importlib.import_module("capstone")
        except Exception as e:
            print(
                f"Error:{e}.\nPlease execute the following command to install dependencies:"
            )
            print("pip install angr capstone")
            sys.exit(1)

        with open(elf, "rb") as f:
            hash = hashlib.md5(f.read()).hexdigest()

        project_cache_path = f"{hash}.project"
        cfg_cache_path = f"{hash}.cfg"

        self.project = self.load_save_cache(
            project_cache_path, lambda: self.angr.Project(elf, auto_load_libs=False)
        )
        self.cfg = self.load_save_cache(
            cfg_cache_path, lambda: self.project.analyses.CFGFast()
        )

        for func in self.cfg.kb.functions.values():
            func._project = self.project

    def load_save_cache(self, file_path, func):
        cache_path = f"{file_path}.pkl"
        if os.path.exists(cache_path):
            with open(cache_path, "rb") as f:
                print(f"load cache {cache_path}")
                return pickle.load(f)
        else:
            print(f"not found cache {cache_path}")
            obj = func()
            with open(cache_path, "wb") as f:
                print(f"create cache {cache_path}")
                pickle.dump(obj, f)
            return obj

    def sym2addr(self, sym: str):
        sym = self.cfg.kb.functions.floor_func(sym)
        if sym:
            return sym.addr
        else:
            return None

    def addr2sym(self, addr: int):
        sym = self.cfg.kb.functions.floor_func(addr)
        if sym:
            return sym.name
        else:
            return f"0x{addr:08x}"

    def addr2func(self, addr: int):
        addr = addr
        sym = self.cfg.kb.functions.floor_func(addr)
        if sym:
            return self.cfg.kb.functions[sym.addr]
        else:
            return None

    def sym2func(self, sym: str):
        addr = self.sym2addr(sym)
        if addr:
            return self.cfg.kb.functions[addr]
        else:
            return None

    def addr2block(self, addr: int):
        addr = addr
        sym = self.cfg.kb.functions.floor_func(addr)
        if sym:
            func = self.cfg.kb.functions[sym.addr]
            for block in func._addr_to_block_node.values():
                if addr >= block.addr and addr < block.addr + block.size:
                    return block
        else:
            return None

    def print_block(self, block):
        print(f"basic block: {hex(block.addr)}, size: {block.size}")
        for insn in block.capstone.insns:
            print(f"    {hex(insn.address)}: {insn.mnemonic} {insn.op_str}")
            if insn.group(self.capstone.CS_GRP_CALL):
                print(f" call {hex(insn.operands[0].imm)}")
            elif insn.group(self.capstone.CS_GRP_RET):
                print(f" return {hex(insn.operands[0].imm)}")
            elif insn.group(self.capstone.CS_GRP_INT):
                print(f" interrupt {hex(insn.operands[0].imm)}")
            elif insn.group(self.capstone.CS_GRP_BRANCH_RELATIVE):
                target_addr = insn.operands[0].imm
                target_function = self.cfg.kb.functions.floor_func(target_addr)
                print(
                    f" conditional jump {hex(target_addr)} (target function: {target_function.name})"
                )
            elif insn.group(self.capstone.CS_GRP_JUMP):
                sub_function = self.cfg.kb.functions.floor_func(insn.operands[0].imm)
                print(
                    f" unconditional jump {hex(insn.operands[0].imm)} (target function: {sub_function.name})"
                )
        sys.stdout.flush()

    def print_function(self, function):
        print(f"function: {function.name} {function.addr} {function.size}")
        for block in function.blocks:
            self.print_block(block)

    def get_block_next_address(self, block):
        for insn in block.capstone.insns:
            if insn.group(self.capstone.CS_GRP_BRANCH_RELATIVE):
                try:
                    addr = insn.operands[0].imm
                    func = self.addr2func(addr)
                    if func:
                        next_block = func.get_block(addr)
                        return next_block.addr
                except Exception as e:
                    log.error(
                        f"get_block_next_address failed: addr: {block.addr:#08x} {e}"
                    )
                    return None
        return None
