from os import listdir

__author__ = "Wei-Chung Huang"
__copyright__ = "Copyright 2018, The SJSU MSCS Master project"
__license__ = "MIT"
__version__ = "1.0.0"


def prefetch_images(file_path):
    """
    Get all image names in file_path
    :param file_path: the file path containing malware images with its family name as the parent folder
    :return: a list of tuples with [0]: family name [1]: list of image names
    """
    families = listdir(file_path)
    [families.remove(family) if family[0] == '.' else None for family in families]

    image_list = []
    for family in families:
        files = listdir(file_path + '/' + family)
        image_list.append((family, files))

    return image_list


def test():
    dataset_path = "../dataset/"
    image_list = prefetch_images(dataset_path)

    # test for getting the first image path
    if len(image_list) > 0:
        family = image_list[0][0]
        images = image_list[0][1]
        image_path = dataset_path + family + '/' + images[0]
        print(image_path)


if __name__ == "__main__":
    test()
