#!/usr/bin/env python3
#-*- encoding: Utf-8 -*-
from re import search, IGNORECASE
from argparse import Namespace
from io import BytesIO
import logging

"""
    The ElfSymbolizer class, defined in this file, gathers information from
    the other modules (such as kallsyms_finder, which extracts the kernel's
    runtime symbol table, or vmlinuz_decompressor, which processes possible
    kernel compressions), in order to generate the output ELF file.
"""

try:
    from kallsyms_finder import KallsymsFinder, KallsymsSymbolType

    from utils.elf import ElfFile, ElfSymtab, ElfRel, Elf32LittleEndianSymbolTableEntry, Elf32BigEndianSymbolTableEntry, Elf64LittleEndianSymbolTableEntry, Elf64BigEndianSymbolTableEntry, SPECIAL_SECTION_INDEX, ST_INFO_TYPE, ST_INFO_BINDING, ElfStrtab, ElfProgbits, ElfNullSection, ElfNoBits, SH_FLAGS

except ImportError:
    from vmlinux_to_elf.kallsyms_finder import KallsymsFinder, KallsymsSymbolType

    from vmlinux_to_elf.utils.elf import ElfFile, ElfSymtab, ElfRel, Elf32LittleEndianSymbolTableEntry, Elf32BigEndianSymbolTableEntry, Elf64LittleEndianSymbolTableEntry, Elf64BigEndianSymbolTableEntry, SPECIAL_SECTION_INDEX, ST_INFO_TYPE, ST_INFO_BINDING, ElfStrtab, ElfProgbits, ElfNullSection, ElfNoBits, SH_FLAGS


class ElfSymbolizer():
    
    def __init__(self, file_contents : bytes, output_file : str,
        elf_machine : int = None, bit_size : int = None,
        base_address : int = None, file_offset : int = None):
        
        if file_contents.startswith(b'\x27\x05\x19\x56'): # uImage header magic (always big-endian)
            
            if file_offset is None:
                file_offset = 64 # uImage header size (image_header_t from u-boot/image.h)
            
            if base_address is None:
                base_address = int.from_bytes(file_contents[4 * 4:4 * 5], 'big')
            
            
        if file_offset:
            file_contents = file_contents[file_offset:]
        
        kallsyms_finder = KallsymsFinder(file_contents, bit_size)
        
        
        if file_contents.startswith(b'\x7fELF'):
            
            kernel = ElfFile.from_bytes(BytesIO(file_contents))
        
        else:
            
            kernel = ElfFile(kallsyms_finder.is_big_endian, kallsyms_finder.is_64_bits)
            
            #  Previsouly the register size was based on the kernel version string:       bool(kallsyms_finder.offset_table_element_size >= 8 or search('itanium|(?:amd|aarch|ia|arm|x86_|\D-)64', kallsyms_finder.version_string, flags = IGNORECASE))
            
            if elf_machine is not None:
                kernel.file_header.e_machine = elf_machine
            else:
                kernel.file_header.e_machine = kallsyms_finder.elf_machine
            
            ET_EXEC = 2
            kernel.file_header.e_type = ET_EXEC
            
            null = ElfNullSection(kernel)
            null.section_name = ''
            
            progbits = ElfProgbits(kernel)
            progbits.section_name = '.kernel'
            progbits.section_header.sh_flags = SH_FLAGS.SHF_ALLOC | SH_FLAGS.SHF_EXECINSTR | SH_FLAGS.SHF_WRITE
            progbits.section_header.sh_size = len(file_contents)
            
            first_symbol_virtual_address = next((symbol.virtual_address for symbol in kallsyms_finder.symbols if symbol.symbol_type == KallsymsSymbolType.TEXT), None)
            
            if kallsyms_finder.has_base_relative:
                first_symbol_virtual_address = min(first_symbol_virtual_address, kallsyms_finder.relative_base_address)
            
            if base_address is not None:
                progbits.section_header.sh_addr = base_address
            elif kallsyms_finder.kernel_text_candidate:
                progbits.section_header.sh_addr = kallsyms_finder.kernel_text_candidate
            else:
                progbits.section_header.sh_addr = first_symbol_virtual_address & 0xfffffffffffff000
            
            progbits.section_contents = file_contents
            
            
            bss = ElfNoBits(kernel)
            bss.section_name = '.bss'
            bss.section_header.sh_flags = SH_FLAGS.SHF_ALLOC | SH_FLAGS.SHF_EXECINSTR | SH_FLAGS.SHF_WRITE
            bss.section_header.sh_size = 0x100000
            bss.section_header.sh_addr = progbits.section_header.sh_addr + len(file_contents)
            
            kernel.sections += [null, progbits, bss]
            

        
        """
            Find the entry point symbol. Based on executing this command
            on the Linux tree source:
            
            for i in $(find -iname 'vmlinux.lds.S' -o -iname 'dyn.lds.S' -o -iname 'vmlinux-std.lds');
                do echo "$i:"$(grep -P '^ENTRY\(' $i);
            done | grep -Po 'ENTRY\((.+?)\)' | sort -u

            You can find the possible symbols that are used as an entry
            point for the kernel, here sorted from the most specific to
            the less specific
        """
        
        POSSIBLE_ENTRY_POINT_SYMBOLS = [
            'kernel_entry', 'microblaze_start', 'parisc_kernel_start',
            'phys_startup_32', 'phys_startup_64', 'phys_start', '_stext_lma',
            'res_service', '_c_int00',
            'startup_32', 'startup_64', 'startup_continue', 'startup',
            '__start', '_start', 'start_kernel',
            'stext', '_stext', '_text'
        ]
        
        entry_point_address : int = None
        
        for symbol_name in POSSIBLE_ENTRY_POINT_SYMBOLS:
            
            symbol = kallsyms_finder.name_to_symbol.get(symbol_name)
            
            if symbol:
                entry_point_address = symbol.virtual_address
                
                break
        
        if entry_point_address is None:
            
            raise ValueError('No entry point symbol found in the kallsyms')
        
        kernel.file_header.e_entry = entry_point_address
    
        # Add symbols
        
        symtab = next((i for i in kernel.sections if i.section_name == '.symtab'), None)
        
        if not symtab:
            symtab = ElfSymtab(kernel)
            symtab.section_name = '.symtab'
            
            strtab = ElfStrtab(kernel)
            strtab.section_name = '.strtab'
            symtab.string_table = strtab
            
            shstrtab = ElfStrtab(kernel)
            shstrtab.section_name = '.shstrtab'
            
            kernel.section_string_table = shstrtab
            kernel.sections += [symtab, strtab, shstrtab]

        sections = sorted([i for i in kernel.sections if i.section_header.sh_addr > 0], key=lambda x: x.section_header.sh_addr)
        def _find_section(address):
            """
            Uses binary search to quickly find the section which the address belongs to
            """
            # Set baseline and roofline hypotheses, expressed in
            # section table indexes:
            lower_range, upper_range = 0, len(sections) - 1
            # Wait for the hypotheses to converge
            while lower_range < upper_range:
                # Mean operation to pick a new tentative hypothesis
                # (add one to ensure to ceil-round the upper
                # hypothesis in case of a difference of 1)
                middle = (lower_range + upper_range + 1) // 2
                if sections[middle].section_header.sh_addr <= address: # Test the hypothesis
                    lower_range = middle # Use the hypothesis as a baseline
                else:
                    upper_range = middle - 1 # Disqualify the hypothesis
            if (sections[lower_range].section_header.sh_addr <= address <=
                sections[lower_range].section_header.sh_addr +
                sections[lower_range].section_header.sh_size):
                return sections[lower_range] # Select the best hypothesis if it qualifies
            return None

        elf_symbol_class = {
            (False, False): Elf32LittleEndianSymbolTableEntry,
            (True, False): Elf32BigEndianSymbolTableEntry,
            (False, True): Elf64LittleEndianSymbolTableEntry,
            (True, True): Elf64BigEndianSymbolTableEntry,
        }[(kernel.is_big_endian, kernel.is_64_bits)]

        for symbol in kallsyms_finder.symbols:

            elf_symbol = elf_symbol_class(kernel.is_big_endian, kernel.is_64_bits)
            
            elf_symbol.symbol_name = symbol.name
            elf_symbol.st_value = symbol.virtual_address
            
            if symbol.symbol_type not in (KallsymsSymbolType.TEXT, KallsymsSymbolType.WEAK_SYMBOL_WITH_DEFAULT):
                elf_symbol.st_info_type = ST_INFO_TYPE.STT_OBJECT
            else:
                elf_symbol.st_info_type = ST_INFO_TYPE.STT_FUNC
            
            if symbol.symbol_type in (KallsymsSymbolType.WEAK_OBJECT_WITH_DEFAULT, KallsymsSymbolType.WEAK_SYMBOL_WITH_DEFAULT):
                elf_symbol.st_info_binding = ST_INFO_BINDING.STB_WEAK
            elif symbol.is_global:
                elf_symbol.st_info_binding = ST_INFO_BINDING.STB_GLOBAL
            else:
                elf_symbol.st_info_binding = ST_INFO_BINDING.STB_LOCAL

            if symbol.symbol_type == KallsymsSymbolType.ABSOLUTE:
                elf_symbol.st_shndx = SPECIAL_SECTION_INDEX.SHN_ABS
            else:
                elf_symbol.associated_section = _find_section(symbol.virtual_address)

            symtab.symbol_table.append(elf_symbol)
        
        # Save the modified ELF
        
        with open(output_file, 'wb') as fd:
            
            kernel.serialize(fd)
        
        logging.info('[+] Successfully wrote the new ELF kernel to %s' % output_file)
