import os
import pefile
import time
import re
import click
import subprocess

data_directory_list = ['DIRECTORY_ENTRY_DEBUG', 'DIRECTORY_ENTRY_EXPORT', 'DIRECTORY_ENTRY_LOAD_CONFIG',
                         'DIRECTORY_ENTRY_RESOURCE', 'DIRECTORY_ENTRY_BASERELOC', 'DIRECTORY_ENTRY_TLS']
normal_section_names = ['.text', '.rdata', '.data', '.pdata', '.rsrc', '.idata', '.bss', '.code', '.edata']

def entropy(name, path):
    entropy_list = []
    entropy = subprocess.check_output("ent '{}' | head -n 1 | cut -d' ' -f 3".format((path + name)),
                                      shell=True).decode('utf8')
    entropy_list.append(entropy[0:-1])

    pe = pefile.PE(path + name)
    text_flag = False
    data_flag = False
    for section in pe.sections:
        try:
            section_name = (section.Name).decode('utf-8')
            section_name = section_name.replace('\x00','')
            if section_name =='.text':
                text_entropy = section.get_entropy()
                text_flag = True
            elif section_name =='.data':
                data_entropy = section.get_entropy()
                data_flag = True
        except:
            continue
    entropy_list.append(text_entropy if text_flag else -1)
    entropy_list.append(data_entropy if data_flag else -1)

    return entropy_list

def section_name_checker(section_names):
    """
    :param section_names:
    an array of section names of a program
    :return:
    a 1*2d array that indicate number of nonsuspicious sections and number of suspicious sections,respectively
    """
    number_of_suspicious_names = 0
    number_of_nonsuspicious_names = 0
    for name in section_names:
        if name in normal_section_names:
            number_of_nonsuspicious_names += 1
        else:
            number_of_suspicious_names += 1

    return number_of_suspicious_names,number_of_nonsuspicious_names

def empty_section_name_checker(section_names):

    #---- normalize names --------
    for i in range(len(section_names)):
        section_names[i] = re.sub(' +', ' ',section_names[i])

    if '' in section_names or ' ' in section_names:
        # print(file_name)
        return 0
    else:
        return 1

def data_directory_checker(pe,data_directory_name):
    try:
        if getattr(pe,data_directory_name):
            return 1
        else:
            return 0
    except:
        return 0


@click.command()
@click.option('--path', required=True, help='path of samples')
@click.option('--outputfile', default='features.txt', help='output file name for storing extracted features')

def feature_extractor(path,outputfile):
    start_time = time.time()

    samples = os.listdir(path)
    features_outputfile = open(outputfile,'w')
    for sample in samples:
        try:
            pe = pefile.PE(path + sample)

            # ----------------- Data Directories --------------------
            temp = ''
            for data_directory in data_directory_list:
                temp += str(data_directory_checker(pe, data_directory))
            features_outputfile.write('{},'.format(int(temp,2)))
            print(int(temp,2))
            # ---------------------- file_info -----------------------
            count = 0
            try:
                for entry in pe.FileInfo:
                    if entry[0].Key == b'StringFileInfo':
                        entry = entry[0]
                        for st in entry.StringTable:
                            for entry in (st.entries.items()):
                                count += 1
                    if entry[1].Key == b'StringFileInfo':
                        entry = entry[1]
                        for st in entry.StringTable:
                            for entry in (st.entries.items()):
                                count += 1
                features_outputfile.write('{},'.format(count))
            except:
                features_outputfile.write('{},'.format(count))
            print(count)
            # ---------------------- checksum ------------------------
            try:
                checksum = pe.OPTIONAL_HEADER.CheckSum
                features_outputfile.write('0,'.format(sample)) if checksum == 0 else features_outputfile.write(
                    '1,'.format(sample))
            except:
                features_outputfile.write('0,'.format(sample))

            # ------------------------- entropy ---------------------------
            entropies = entropy(sample, path)
            for entro in entropies:
                features_outputfile.write('{},'.format(entro))
            print(entropies)
            # ----------------------- section names -----------------------
            section_names = []
            try:
                sections = pe.sections
                for section in sections:
                    name = (section.Name).decode('utf-8')
                    name = name.replace('\x00', '')
                    section_names.append(name)
            except:
                continue

            section_name_features = section_name_checker(section_names)
            features_outputfile.write('{},{},'.format(section_name_features[0], section_name_features[1]))
            empty_section_names = empty_section_name_checker(section_names)
            features_outputfile.write('{},{}\n'.format(empty_section_names, sample))
            print(section_name_features)
            print(empty_section_names)

        except:
            print('{} is not a pe file'.format(sample))

    end_time = time.time()
    print('feature extraction time: {}s'.format(end_time - start_time))


if __name__ == '__main__':
    feature_extractor()

import os
import pefile
import time
import re
import click
import subprocess

data_directory_list = ['DIRECTORY_ENTRY_DEBUG', 'DIRECTORY_ENTRY_EXPORT', 'DIRECTORY_ENTRY_LOAD_CONFIG',
                         'DIRECTORY_ENTRY_RESOURCE', 'DIRECTORY_ENTRY_BASERELOC', 'DIRECTORY_ENTRY_TLS']
normal_section_names = ['.text', '.rdata', '.data', '.pdata', '.rsrc', '.idata', '.bss', '.code', '.edata']

def entropy(name, path):
    entropy_list = []
    entropy = subprocess.check_output("ent '{}' | head -n 1 | cut -d' ' -f 3".format((path + name)),
                                      shell=True).decode('utf8')
    entropy_list.append(entropy[0:-1])

    pe = pefile.PE(path + name)
    text_flag = False
    data_flag = False
    for section in pe.sections:
        try:
            section_name = (section.Name).decode('utf-8')
            section_name = section_name.replace('\x00','')
            if section_name =='.text':
                text_entropy = section.get_entropy()
                text_flag = True
            elif section_name =='.data':
                data_entropy = section.get_entropy()
                data_flag = True
        except:
            continue
    entropy_list.append(text_entropy if text_flag else -1)
    entropy_list.append(data_entropy if data_flag else -1)

    return entropy_list

def section_name_checker(section_names):
    """
    :param section_names:
    an array of section names of a program
    :return:
    a 1*2d array that indicate number of nonsuspicious sections and number of suspicious sections,respectively
    """
    number_of_suspicious_names = 0
    number_of_nonsuspicious_names = 0
    for name in section_names:
        if name in normal_section_names:
            number_of_nonsuspicious_names += 1
        else:
            number_of_suspicious_names += 1

    return number_of_suspicious_names,number_of_nonsuspicious_names

def empty_section_name_checker(section_names):

    #---- normalize names --------
    for i in range(len(section_names)):
        section_names[i] = re.sub(' +', ' ',section_names[i])

    if '' in section_names or ' ' in section_names:
        # print(file_name)
        return 0
    else:
        return 1

def data_directory_checker(pe,data_directory_name):
    try:
        if getattr(pe,data_directory_name):
            return 1
        else:
            return 0
    except:
        return 0


@click.command()
@click.option('--path', required=True, help='path of samples')
@click.option('--outputfile', default='features.txt', help='output file name for storing extracted features')

def feature_extractor(path,outputfile):
    start_time = time.time()

    samples = os.listdir(path)
    features_outputfile = open(outputfile,'w')
    for sample in samples:
        try:
            pe = pefile.PE(path + sample)

            # ----------------- Data Directories --------------------
            temp = ''
            for data_directory in data_directory_list:
                temp += str(data_directory_checker(pe, data_directory))
            features_outputfile.write('{},'.format(int(temp,2)))
            print(int(temp,2))
            # ---------------------- file_info -----------------------
            count = 0
            try:
                for entry in pe.FileInfo:
                    if entry[0].Key == b'StringFileInfo':
                        entry = entry[0]
                        for st in entry.StringTable:
                            for entry in (st.entries.items()):
                                count += 1
                    if entry[1].Key == b'StringFileInfo':
                        entry = entry[1]
                        for st in entry.StringTable:
                            for entry in (st.entries.items()):
                                count += 1
                features_outputfile.write('{},'.format(count))
            except:
                features_outputfile.write('{},'.format(count))
            print(count)
            # ---------------------- checksum ------------------------
            try:
                checksum = pe.OPTIONAL_HEADER.CheckSum
                features_outputfile.write('0,'.format(sample)) if checksum == 0 else features_outputfile.write(
                    '1,'.format(sample))
            except:
                features_outputfile.write('0,'.format(sample))

            # ------------------------- entropy ---------------------------
            entropies = entropy(sample, path)
            for entro in entropies:
                features_outputfile.write('{},'.format(entro))
            print(entropies)
            # ----------------------- section names -----------------------
            section_names = []
            try:
                sections = pe.sections
                for section in sections:
                    name = (section.Name).decode('utf-8')
                    name = name.replace('\x00', '')
                    section_names.append(name)
            except:
                continue

            section_name_features = section_name_checker(section_names)
            features_outputfile.write('{},{},'.format(section_name_features[0], section_name_features[1]))
            empty_section_names = empty_section_name_checker(section_names)
            features_outputfile.write('{},{}\n'.format(empty_section_names, sample))
            print(section_name_features)
            print(empty_section_names)

        except:
            print('{} is not a pe file'.format(sample))

    end_time = time.time()
    print('feature extraction time: {}s'.format(end_time - start_time))


if __name__ == '__main__':
    feature_extractor()

:param section_names:     an array of section names of a program     :return:     a 1*2d array that indicate number of nonsuspicious sections and number of suspicious sections,respectively #---- normalize names -------- # print(file_name) # ----------------- Data Directories -------------------- # ---------------------- file_info ----------------------- # ---------------------- checksum ------------------------ # ------------------------- entropy --------------------------- # ----------------------- section names -----------------------