import os
from os.path import isfile, join
import random

path_to_benign = "./Offline/dump_file_benign/"
path_to_malware = "./Offline/dump_file_malware/"
path_to_train_result = "./train_results/"

def get_dir_files(path_to_dir):
	files = [path_to_dir + file for file in os.listdir(path_to_dir) if isfile(join(path_to_dir + file))]
	return files

def create():
    print("Creating Train Set and Test Set Files from the database..")
    files_benign = get_dir_files(path_to_benign)
    files_malware = get_dir_files(path_to_malware)
    
    num_benign = len(files_benign)
    num_malware = len(files_malware)
    
    num_test_benign = int(round(num_benign * 0.25))
    num_test_malware = int(round(num_malware * 0.25))
    
    train_benign_index = random.sample(range(num_benign), num_benign - num_test_benign)
    train_malware_index = random.sample(range(num_malware), num_malware - num_test_malware)
    
    test_benign_index = []
    for item in range(num_benign):
    	if item not in train_benign_index:
    		test_benign_index.append(item)
    
    test_malware_index = []
    for item in range(num_malware):
    	if item not in train_malware_index:
    		test_malware_index.append(item)

    print("Train and Test Files created..")
    return files_benign, train_benign_index, test_benign_index, files_malware, train_malware_index, test_malware_index