#This script takes pathname of directory containing malware and benign file as input and output the static.csv file
#Taking pathname of directory containing malware and benign file as input
print("Enter Path : ")
pathname = input()
import pickle
import os
from ast import literal_eval
import numpy as np
import csv
import warnings
warnings.filterwarnings("ignore")

#Loading the trained model
model_filename = 'model.sav'                          
clf = pickle.load(open(model_filename, 'rb'))



if(pathname[-1]!='/'):                  #Make sure that pathname contains '/' at the end
    pathname+='/'

#Fetching the name of folders inside the directory 
folders = os.listdir(pathname)

#Initialising the feature array fot test datset
X_test = np.zeros((len(folders),25))

#Fetures to be selected from file_header of Structure_Info.txt
file_header = ['NumberOfSections:','NumberOfSymbols:','SizeOfOptionalHeader:']
#Fetures to be selected from image_header of Structure_Info.txt
image_header = ['MajorLinkerVersion:','MinorLinkerVersion:','SizeOfCode:','SizeOfInitializedData:','SizeOfUninitializedData:','SizeOfImage:','SizeOfHeaders:','SizeOfStackReserve:','SizeOfStackCommit:','SizeOfHeapReserve:','SizeOfHeapCommit:','NumberOfRvaAndSizes:']
#Fetures to be selected from sections of Structure_Info.txt
sections = ['.text','.rdata','.data','.rsrc','.idata','.reloc']
#Fetures to be selected from exported_symbols of Structure_Info.txt
exported_symbols = ['NumberOfFunctions:','NumberOfNames:']

filenames = []
k = 0

#Selecting features from Structure_Info.txt'

for x in folders:
    filenames.append(x)             #apending the name of files, that will come handy in producing output.csv
    f = open(('{}{}{}{}'.format(pathname, x,'/','Structure_Info.txt')),"r",errors = "ignore")
    lines = (f.readlines())
    a=[]
    b=[]
    count = 0
    flag1 = 1
    flag2 = 1
    flag3 = 1
    flag4 = 1
    
    for j in range(len(lines)):
        b = lines[j].split()
        a += b
        
    for j in range(len(a)):
        if(flag1 and a[j] in file_header):
            X_test[k][file_header.index(a[j])] = literal_eval(a[j+1])       #To convert hex value to integer
            
        if(flag2 and a[j] in image_header):
            X_test[k][image_header.index(a[j])+3] = literal_eval(a[j+1])    #To convert hex value to integer
            flag1 = 0
            
        if(a[j] in sections):
            X_test[k][sections.index(a[j])+15] = 1                          #If the section (.text,idata, etc) is present then 1 else 0, as malware try to hide their purpose by changing the name of sections
            
        if(flag3 and a[j] in exported_symbols):
            X_test[k][exported_symbols.index(a[j])+21] = literal_eval(a[j+1])   #To convert hex value to unteger
            flag2 = 0
            
        if(a[j] == '----------Imported'):
            flag4 = 1
            flag3 = 0
            
        if(a[j] == '----------Resource'):
            flag4 = 0
            
        if(flag4 == 1):             #Checks for dll only if flag4 is set to save computation, as imported dlls occurs in imports sections of Structutre.txt 
            x = str(a[j])
            z = x.find('.dll')
            if(z!=-1):
                count = count+1
    X_test[k][23] = count
    k = k+1 
    f.close()


k = 0
#Selecting features from String.txt(basically number of strings as packed alwares usually have less number of strings)
for x in folders:
    f = open(('{}{}{}{}'.format(pathname, x,'/','String.txt')),"r",errors = "ignore")
    lines = (f.readlines())
    n = len(lines)
    X_test[k][24] = n
    k = k+1
    f.close()


#Selecting relevant features (Since others play no role in classification as they had almost same values for both malware and benign found during training and validation)
X_new = X_test[:,3:11]
X_new = np.append(X_new,X_test[:,12:13],axis = 1)
X_new = np.append(X_new,X_test[:,15:25],axis = 1)
X_new = np.append(X_new,X_test[:,0:1],axis =1)
X_test = X_new
    
#Predecting the labels of test dataset
y_pred = clf.predict(X_test)

out_pred = []

#Setting 'B' for Benign and 'M' for Malware 
for i in range(y_pred.size):
    if(y_pred[i]==0):
        out_pred.append('B')
    else:
        out_pred.append('M')

#To output the csv file
with open('static.csv', 'w', newline='') as file:       
    to_write = csv.writer(file)
    to_write.writerow(["File_Hash","Predicted Label"])
    for i in range(y_pred.size):
        to_write.writerow([filenames[i],out_pred[i]])
