import argparse
import os
import sys

from pwn import *

from elftools.elf.elffile import ELFFile
from elftools.elf.dynamic import DynamicSection
from pathlib import Path
from rgbprint import rgbprint, Color
from typing import Callable, Dict, Iterable, List, Type

#
# constants.py
#

DEBUG_FILE = Path(__file__).resolve().parent / 'debug.txt'

#
# helpers.py
#

def write_data_to_a_debug_file(file: str, data: str) -> None:
    try:
        with open(file, 'a+') as debug_file:
            debug_file.write(data)
    except Exception as error:
        raise Exception(f'FUNCTION: write_data_to_a_debug_file() - {error}')

#
# elf_local_data.py
#

class ELFLocalData():
    _el_file: Type[ELFFile]
    _el_segments_dict: Dict
    _el_segments_iterable: Iterable
    _el_sections_dict: Dict
    _el_sections_iterable: Iterable
    _el_tags_dict: Dict
    _el_elf_base_address: int
    _el_elf_init_section_offset: int

    #

    def __init__(self) -> None:
        file = open(self._program_path, 'rb')
        self._el_file = ELFFile(file)

        self._el_set_all_elf_segments()
        self._el_set_all_elf_sections()
        self._el_set_all_elf_tags()
        self._el_set_elf_base_address()
        self._el_set_elf_init_section_offset()

    #

    def el_get_segment_by_number(self, number) -> Dict:
        return self._el_segments_dict[number]


    def el_get_section_by_name(self, name) -> Dict:
        return self._el_sections_dict[name]


    def el_get_tag_by_name(self, name) -> Dict:
        return self._el_tags_dict[name]


    def el_get_elf_base_address(self) -> int:
        return self._el_elf_base_address


    def el_get_elf_init_section_offset(self) -> int:
        return self._el_elf_init_section_offset

    #

    def _el_set_all_elf_segments(self) -> Dict:
        self._el_segments_dict = {}
        counter = 0
        self._el_segments_iterable = self._el_file.iter_segments()
        for segment in self._el_segments_iterable:
            current_segment = {}
            for key, value in segment.header.items():
                current_segment[key] = value
            self._el_segments_dict.update({counter : current_segment})
            counter += 1

        return self._el_segments_dict


    def _el_set_all_elf_sections(self) -> Dict:
        self._el_sections_dict = {}
        self._el_sections_iterable = self._el_file.iter_sections()
        for section in self._el_sections_iterable:
            current_section = {}
            for key, value in section.header.items():
                current_section[key] = value
            self._el_sections_dict.update({section.name : current_section})

        return self._el_sections_dict


    def _el_set_all_elf_tags(self) -> Dict:
        self._el_tags_dict = {}
        for section in self._el_sections_iterable:
            if section.is_null() == False:
                if not isinstance(section, DynamicSection):
                    for tag in section.iter_tags():
                        self._el_tags_dict.update({tag['d_tag'] : tag['d_val']})

        return self._el_tags_dict


    def _el_set_elf_base_address(self) -> int:
        for i in range(len(self._el_segments_dict)):
            if self._el_segments_dict[i]['p_type'] == 'PT_LOAD':
                if self._virtual_memory_mapping == True:
                    address = self._el_segments_dict[i]['p_vaddr']
                else:
                    address = self._el_segments_dict[i]['p_paddr']
                offset = self._el_segments_dict[i]['p_offset']
                break
        self._el_elf_base_address = address - offset

        return self._el_elf_base_address


    def _el_set_elf_init_section_offset(self) -> None:
        self._el_elf_init_section_offset = self._el_sections_dict['.init']['sh_addr'] - self._el_elf_base_address

        return self._el_elf_init_section_offset