#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import hashlib
import math
import os
import re
import subprocess
from binascii import *

# Utility method for converting hex strings to decimal numbers
def hex_to_dec(hex_string):
    try:
        return int(hex_string, base=16)
    except:
        return -1

# Class: PE32
# Encapsulates the properties of the the PE file it is initialised with
class PE32:
    # Starting constant offset
    LOC_PE_SIG_OFFSET = "0x3c"
    # Constants used for calculating offsets
    PE_SIG_LENGTH = 4
    COFF_LENGTH = 20
    SECTION_HEADER_SIZE = 40

    # Constructor
    # Param: None
    def __init__(self):
        self.file = None
        self.fileName = None
        self.directory = None

    # Used to initialise the properties the PE32 object
    def File(self, file, filename, directory):
        self.file = file
        if(directory != ""):
            self.fileName = directory + "/" + filename
        else:
            self.fileName = filename
        self.directory = directory

    @property
    def get_file_name(self):
        return self.fileName

    # Read specified number of bytes (length) from an offset (parent_offset + offset)
    def read_bytes(self, parent_offset, offset, length):
        absolute_offset = parent_offset + offset
        _bytes = list()
        # Read single byte and append to list of bytes read so far
        for i in range(length):
            _bytes.append(self.read_byte(absolute_offset + i).decode())
        return _bytes

    # Construct hex string from list of bytes and reverse using slicing syntax
    def read_big_endian_bytes(self, parent_offset, offset, length):
        _bytes = self.read_bytes(parent_offset, offset, length)
        return ''.join(_bytes[::-1])

    # Construct hex string from list of bytes
    def read_little_endian_bytes(self, parent_offset, offset, length):
        _bytes = self.read_bytes(parent_offset, offset, length)
        return ''.join(_bytes)

    # Read 1 byte from specified offset
    def read_byte(self, offset):
        self.file.seek(offset)
        return hexlify(self.file.read(1))

    # calculate the ratio of code to data, return 0 if size of code or data is 0
    def get_code_data_ratio(self):
        if hex_to_dec(self.get_size_of_code()) == 0 or hex_to_dec(self.get_size_of_initialized_data()) == 0:
            return 0
        else:
            return hex_to_dec(self.get_size_of_code()) / hex_to_dec(self.get_size_of_initialized_data())

    # Calcualte an entropy score between 0 and 8 between start offset and start offset + length
    def calculate_entropy(self, start_offset, length):
        # Initialise array to store counts of each byte value
        # Index in the array corresponds to the byte value
        byte_occurences = [0] * 256

        # Read bytes in range start_offset, start_offset + length individually
        # Increment the number of occurances at the index corresponding with the value of the byte
        for offset in range(start_offset, start_offset + length):
            hex_value = self.read_byte(offset)
            decimal_value = hex_to_dec(hex_value)
            byte_occurences[decimal_value] += 1

        entropy = 0

        for occurences in byte_occurences:
            # For each byte count, continue if count is zero
            if occurences == 0:
                continue
            # Calculate frequency as count / length (number of bytes)
            frequency = float(occurences) / float(length)
            # Calculate 8frequency log(base256)(frequency) and subtract from the entropy score
            entropy -= frequency * math.log(frequency, 256) * 8
        return entropy

    # Call strings command line tool and store output in file_strings variable
    def get_imports(self):
        file_strings = subprocess.check_output(['strings', self.get_file_name])
        # Find all strings which end in .dll/.exe/.DLL/.EXE and return
        imports = re.findall('\w+\.dll|\w+\.DLL|\w+\.EXE|\w+\.exe', file_strings.decode())
        return imports

    # Methods used to fetch features from PE FILE
    # Values used as parameters taken from the PE File Specification

    def get_pe_sig_offset(self):
        return self.read_byte(hex_to_dec(self.LOC_PE_SIG_OFFSET))

    def get_optional_header_offset(self):
        return self.get_coff_header_offset() + self.COFF_LENGTH

    def get_section_header_table_offset(self):
        return self.get_optional_header_offset() + hex_to_dec(self.get_size_of_optional_header())

    def get_coff_header_offset(self):
        return hex_to_dec(self.get_pe_sig_offset()) + self.PE_SIG_LENGTH

    def get_file_length(self):
        return os.path.getsize(self.get_file_name)

    def calc_section_offset(self, index):
        return self.SECTION_HEADER_SIZE * index

    def get_file_entropy(self):
        return self.calculate_entropy(0, self.get_file_length())

    def get_number_of_imports(self):
        return int((hex_to_dec(self.get_import_table_size()) / 20) - 1)

    def get_machine(self):
        return self.read_big_endian_bytes(self.get_coff_header_offset(), offset=0, length=2)

    def get_number_of_sections(self):
        return hex_to_dec(self.read_big_endian_bytes(self.get_coff_header_offset(), offset=2, length=2))

    def get_time_date_stamp(self):
        return self.read_big_endian_bytes(self.get_coff_header_offset(), offset=4, length=4)

    def get_pointer_to_symbol_table(self):
        return self.read_big_endian_bytes(self.get_coff_header_offset(), offset=8, length=4)

    def get_number_of_symbols(self):
        return self.read_big_endian_bytes(self.get_coff_header_offset(), offset=12, length=4)

    def get_size_of_optional_header(self):
        return self.read_big_endian_bytes(self.get_coff_header_offset(), offset=16, length=2)

    def get_characteristics(self):
        return self.read_big_endian_bytes(self.get_coff_header_offset(), offset=18, length=2)

    def get_magic(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=0, length=2)

    def get_major_linker_version(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=2, length=1)

    def get_minor_linker_version(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=3, length=1)

    def get_size_of_code(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=4, length=4)

    def get_size_of_initialized_data(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=8, length=4)

    def get_size_of_uninitialized_data(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=12, length=4)

    def get_address_of_entry_point(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=16, length=4)

    def get_base_of_code(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=20, length=4)

    def get_base_of_data(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=24, length=4)

    def get_image_base(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=28, length=4)

    def get_section_alignment(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=32, length=4)

    def get_file_alignment(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=36, length=4)

    def get_major_operating_system_version(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=40, length=2)

    def get_minor_operating_system_version(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=42, length=2)

    def get_major_image_version(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=44, length=2)

    def get_minor_image_version(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=46, length=2)

    def get_major_subsystem_version(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=48, length=2)

    def get_minor_subsystem_version(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=50, length=2)

    def get_win32_version_value(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=52, length=4)

    def get_size_of_image(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=56, length=4)

    def get_size_of_headers(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=60, length=4)

    def get_check_sum(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=64, length=4)

    def get_subsystem(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=68, length=2)

    def get_dll_characteristics(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=70, length=2)

    def get_size_of_stack_reserve(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=72, length=4)

    def get_size_of_stack_commit(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=76, length=4)

    def get_size_of_heap_reserve(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=80, length=4)

    def get_size_of_heap_commit(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=84, length=4)

    def get_loader_flags(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=88, length=4)

    def get_number_of_rva_and_sizes(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=92, length=4)

    def get_export_table_rva(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=96, length=4)

    def get_export_table_size(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=100, length=4)

    def get_import_table_rva(self):
        return hex_to_dec(self.read_big_endian_bytes(self.get_optional_header_offset(), offset=104, length=4))

    def get_import_table_size(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=108, length=4)

    def get_resource_table_rva(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=112, length=4)

    def get_resource_table_size(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=116, length=4)

    def get_exception_table_rva(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=120, length=4)

    def get_exception_table_size(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=124, length=4)

    def get_certificate_table_rva(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=128, length=4)

    def get_certificate_table_size(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=132, length=4)

    def get_base_relocation_table_rva(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=136, length=4)

    def get_base_relocation_table_size(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=140, length=4)

    def get_debug_rva(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=144, length=4)

    def get_debug_size(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=148, length=4)

    def get_architecture_rva(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=152, length=4)

    def get_architecture_size(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=156, length=4)

    def get_global_ptr_rva(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=160, length=4)

    def get_global_ptr_size(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=164, length=4)

    def get_tls_table_rva(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=168, length=4)

    def get_tls_table_size(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=172, length=4)

    def get_load_config_table_rva(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=176, length=4)

    def get_load_config_table_size(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=180, length=4)

    def get_bound_import_rva(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=184, length=4)

    def get_bound_import_size(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=188, length=4)

    def get_iat_rva(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=192, length=4)

    def get_iat_size(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=196, length=4)

    def get_delay_import_descriptor_rva(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=200, length=4)

    def get_delay_import_descriptor_size(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=204, length=4)

    def get_clr_runtime_header_rva(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=208, length=4)

    def get_clr_runtime_header_size(self):
        return self.read_big_endian_bytes(self.get_optional_header_offset(), offset=212, length=4)

    def get_section_name(self, index):
        _bytesString = self.read_little_endian_bytes(self.get_section_header_table_offset() + self.calc_section_offset(index),
                                                     offset=0, length=8)
        return bytearray.fromhex(_bytesString).decode()

    def get_section_virtual_size(self, index):
        return self.read_big_endian_bytes(self.get_section_header_table_offset() + self.calc_section_offset(index), offset=8,
                                          length=4)

    def get_section_md5(self, index):
        h = hashlib.new('md5')
        pointer = hex_to_dec(self.get_section_pointer_to_raw_data(index))
        length = hex_to_dec(self.get_section_size_of_raw_data(index))
        content = self.read_big_endian_bytes(pointer, 0, length).encode('UTF-8')
        h.update(content)
        return h.hexdigest()

    def get_section_virtual_address(self, index):
        return self.read_big_endian_bytes(self.get_section_header_table_offset() + self.calc_section_offset(index), offset=12,
                                          length=4)

    def get_section_size_of_raw_data(self, index):
        return self.read_big_endian_bytes(self.get_section_header_table_offset() + self.calc_section_offset(index), offset=16,
                                          length=4)

    def get_section_pointer_to_raw_data(self, index):
        return self.read_big_endian_bytes(self.get_section_header_table_offset() + self.calc_section_offset(index), offset=20,
                                          length=4)

    def get_section_pointer_to_relocations(self, index):
        return self.read_big_endian_bytes(self.get_section_header_table_offset() + self.calc_section_offset(index), offset=24,
                                          length=4)

    def get_section_pointer_to_linenumbers(self, index):
        return self.read_big_endian_bytes(self.get_section_header_table_offset() + self.calc_section_offset(index), offset=28,
                                          length=4)

    def get_section_number_of_relocations(self, index):
        return self.read_big_endian_bytes(self.get_section_header_table_offset() + self.calc_section_offset(index), offset=32,
                                          length=2)

    def get_section_number_of_linenumbers(self, index):
        return self.read_big_endian_bytes(self.get_section_header_table_offset() + self.calc_section_offset(index), offset=34,
                                          length=2)

    def get_section_characteristics(self, index):
        return self.read_big_endian_bytes(self.get_section_header_table_offset() + self.calc_section_offset(index), offset=36,
                                          length=4)
