from enum import Enum, Flag
import struct

class ELFPType(Enum):
    NULL = 0
    LOAD = 1
    DYNAMIC = 2
    INTERP = 3
    NOTE = 4
    SHLIB = 5
    PHDR = 6
    TLS = 7


class ELFPFlag(Flag):
    NONE = 0
    EXEC = 1
    WRITE = 2
    READ = 4


class ELFEndian(Enum):
    LITTLE = 1
    BIG = 2


class ELF64Phdr:
    def __init__(self):
        self.type = ELFPType.NULL
        self.flags = 0
        self.offset = 0
        self.vaddr = 0
        self.paddr = 0
        self.filesz = 0
        self.memsz = 0
        self.align = 0

        return

    def parse(self, buf, endian):
        endian_marker = '>' if endian == ELFEndian.BIG else '<'
        entry = struct.unpack('{}IIQQQQQQ'.format(endian_marker), buf)

        self.type = ELFPType(entry[0])
        self.flags = ELFPFlag.NONE
        self.offset = entry[2]
        self.vaddr = entry[3]
        self.paddr = entry[4]
        self.filesz = entry[5]
        self.memsz = entry[6]
        self.align = entry[7]

        if entry[1] & ELFPFlag.READ.value: self.flags |= ELFPFlag.READ
        if entry[1] & ELFPFlag.WRITE.value: self.flags |= ELFPFlag.WRITE
        if entry[1] & ELFPFlag.EXEC.value: self.flags |= ELFPFlag.EXEC


class ELF64Image:
    def __init__(self):
        self.ph_list = []
        self.text = None
        self.data = None
        self.entry = None

    def load(self, path):
        f = open(path, mode='rb')
        ident_bytes = f.read(16)

        ident = struct.unpack('BBBBBBBBBBBBBBBB', ident_bytes)

        if ident[0] != 0x7F or ident[1] != ord('E') or ident[2] != ord('L') or ident[3] != ord('F'):
            f.close()
            return -1

        endian = ELFEndian(ident[5])
        endian_marker = '>' if endian == ELFEndian.BIG else '<'

        hdr_bytes = f.read(48)
        hdr = struct.unpack('{}HHIQQQIHHHHHH'.format(endian_marker), hdr_bytes)

        self.entry = hdr[3]
        phoff = hdr[4]
        phnum = hdr[9]
        phentsize = hdr[8]

        for i in range(0, phnum):
            f.seek(phoff + (i * phentsize))
            phdr_bytes = f.read(phentsize)

            phdr = ELF64Phdr()
            phdr.parse(phdr_bytes, endian)

            if phdr.flags == ELFPFlag.READ | ELFPFlag.EXEC:
                self.text = phdr
            if phdr.flags == ELFPFlag.READ | ELFPFlag.WRITE:
                self.data = phdr

            self.ph_list.append(phdr)

        f.close()
        return 0