from shutil import copyfile
import os

BENIGNDIR = ''
TRAINDIR = ''
VALIDDIR = ''

def get_file_list(dir):
    files = []
    for _, _, filenames in os.walk(dir):
        files.extend(filenames)
        break

    return files


def main():
    train_list = get_file_list(TRAINDIR)
    benign_list = get_file_list(BENIGNDIR)
    for file in train_list:
        with open('train-label.csv', 'w') as f:
            if file in benign_list:
                f.write(file + ',' + '0\n')
            else:
                f.write(file + ',' + '1\n')

    del train_list

    valid_list = get_file_list(VALIDDIR)
    for file in valid_list:
        with open('valid-label.csv', 'w') as f:
            if file in benign_list:
                f.write(file + ',' + '0\n')
            else:
                f.write(file + ',' + '1\n')


if __name__ == '__main__':
    main()
