# Copyright 2025 NXP

# SPDX-License-Identifier: BSD-3-Clause

import re
import struct

########################################################################
# Dwarf Location Parsers
#
# This module stores a collection of entries extracted from the .debug_loc
# section of DWARF debugging information. It also includes an extension
# for the 'DW_OP_regx' parameter, which can be reused in other scripts. Additional
# extensions can be added to enhance functionality and support more use cases.
########################################################################

DW_OPCODES = {
    0x03: "DW_OP_addr",
    0x06: "DW_OP_deref",
    0x08: "DW_OP_const1u",
    0x09: "DW_OP_const1s",
    0x0a: "DW_OP_const2u",
    0x0b: "DW_OP_const2s",
    0x0c: "DW_OP_const4u",
    0x0d: "DW_OP_const4s",
    0x0e: "DW_OP_const8u",
    0x0f: "DW_OP_const8s",
    0x10: "DW_OP_constu",
    0x11: "DW_OP_consts",
    0x12: "DW_OP_dup",
    0x13: "DW_OP_drop",
    0x22: "DW_OP_plus",
    0x90: "DW_OP_regx",
    0x91: "DW_OP_fbreg",
    0x92: "DW_OP_bregx",
    0x93: "DW_OP_piece",
    0x9f: "DW_OP_stack_value",
    # We can add more entries if needed.
}

class DwarfLocSection:
    """
    Parses and interprets the .debug_loc section of a DWARF-enabled ELF file.

    This section contains location lists that describe where variables are stored
    (e.g., in registers or memory) during different ranges of program execution.

    Attributes:
        elf_file (ElfFile): The ELF file object.
        section (SectionHeader): The .debug_loc section header.
        entries (list): Parsed location list entries.
        base_address (int): Current base address for relative location entries.
    """
    def __init__(self, elf_file):
        self.elf_file = elf_file
        self.section = elf_file.get_section_header_by_name('.debug_loc')
        self.entries = []
        self.regx_entries = []
        self.base_address = 0

        self._parse()

    def _parse(self):
        """
        Parses the .debug_loc section and extracts location list entries.

        Each entry includes a start and end address, and a DWARF expression
        describing the location of a variable. Handles base address selection
        and end-of-list markers.
        """
        if not self.section:
            print("No .debug_loc section found.")
            return

        data = self.section.read_data()
        offset = 0

        while offset < len(data):
            entry_offset = offset
            # Py3 version, can be changed later when we switch to a newer version
            # begin = int.from_bytes(data[offset:offset+4], 'little')
            # end = int.from_bytes(data[offset+4:offset+8], 'little')
            begin = struct.unpack('<I', data[offset:offset+4])[0]
            end = struct.unpack('<I', data[offset+4:offset+8])[0]
            offset += 8

            if begin == 0 and end == 0:
                self.entries.append((entry_offset, 'End of list'))
                continue

            if begin == 0xffffffff:
                self.base_address = end
                self.entries.append((entry_offset, 'Base address: 0x{0:08x}'.format(end)))
                continue

            expr_len = struct.unpack('<I', data[offset:offset+2])[0]
            offset += 2
            expr = data[offset:offset+expr_len]
            offset += expr_len

            expr_str = self._decode_dwarf_expression(expr)
            range_str = '0x{0:08x} - 0x{1:08x}: {2}'.format(begin + self.base_address, end + self.base_address, expr_str)
            self.entries.append((entry_offset, range_str))
            if "DW_OP_regx" in expr_str:   
                self._fill_regx_entries(expr_str, (begin+self.base_address), (end + self.base_address))


    def _fill_regx_entries(self, expr_with_regx, begin, end):
        
        match = re.search(r'DW_OP_regx:\s*(\d+)', expr_with_regx)
        val = -1
        if match:
            val = int(match.group(1))

        self.regx_entries.append({"begin": begin, "end": end, "regx_index": val})
        return
    
    def get_regx_entries(self):
        """
        Get DW_OP_regx entries with begin and end addresses
        """        
        return self.regx_entries

    def print_entries(self):
        """
        Prints all parsed location list entries in a human-readable format.
        """
        for offset, desc in self.entries:
            print("{0:08x} {1}".format(offset, desc))


    def _decode_dwarf_expression(self, expr_bytes):
        """
        Decodes a DWARF expression byte sequence into a human-readable string.

        Args:
            expr_bytes (bytes): The raw bytes of the DWARF expression.

        Returns:
            str: A string representation of the decoded expression.
        """
        offset = 0
        result = []
        while offset < len(expr_bytes):
            opcode = expr_bytes[offset]
            offset += 1
            opname = DW_OPCODES.get(opcode, "DW_OP_unknown_0x{0:02x}".format(opcode))

            if opcode in (0x90, 0x92):  # DW_OP_regx, DW_OP_bregx
                value, size = self._read_uleb128(expr_bytes, offset)
                offset += size
                result.append("{0}: {1}".format(opname, value))
            elif opcode == 0x93:  # DW_OP_piece
                value, size = self._read_uleb128(expr_bytes, offset)
                offset += size
                result.append("{0}: {1} bytes".format(opname, value))
            elif opcode == 0x9f:  # DW_OP_stack_value
                result.append(opname)
            elif opcode == 0x11:  # DW_OP_consts
                value, size = self._read_sleb128(expr_bytes, offset)
                offset += size
                result.append("{0}: {1}".format(opname, value))
            else:
                result.append(opname)

        return "; ".join(result)


    def _read_uleb128(self, data, offset):
        """
        Reads an unsigned LEB128-encoded integer from a byte array.

        Args:
            data (bytes): The byte array.
            offset (int): The starting offset.

        Returns:
            tuple: (decoded integer, number of bytes read)
        """
        result = 0
        shift = 0
        size = 0
        while True:
            byte = data[offset]
            offset += 1
            size += 1
            result |= (byte & 0x7f) << shift
            if (byte & 0x80) == 0:
                break
            shift += 7
        return result, size


    def _read_sleb128(self, data, offset):
        """
        Reads a signed LEB128-encoded integer from a byte array.

        Args:
            data (bytes): The byte array.
            offset (int): The starting offset.

        Returns:
            tuple: (decoded integer, number of bytes read)
        """
        result = 0
        shift = 0
        size = 0
        byte = 0
        while True:
            byte = data[offset]
            offset += 1
            size += 1
            result |= (byte & 0x7f) << shift
            shift += 7
            if (byte & 0x80) == 0:
                break
        if (shift < 32) and (byte & 0x40):
            result |= - (1 << shift)
        return result, size
