import os


def get_clean_and_infected_filepaths():
    base_path = "../data/elf_files"
    clean_path = os.path.join(base_path, "clean")
    infected_path = os.path.join(base_path, "infected")

    return {"clean": list(map(lambda filename: os.path.join(clean_path, filename), os.listdir(clean_path))),
            "infected": list(map(lambda filename: os.path.join(infected_path, filename), os.listdir(infected_path)))}


def read_byte_windows(path, label, output_filename="dataset.csv", window_size=32):
    base_output_path = "../data/datasets/elf"
    output_path = os.path.join(base_output_path, output_filename)
    data_f_string = "%s,"*window_size + "%s\n"

    with open(output_path, "w") as of:
        column_f_string = "%s,"*window_size + "class\n"
        of.write(column_f_string % tuple(map(lambda x: "byte" + str(x), range(window_size))))

        with open(path, "rb") as f:
            data = str(f.read(), "latin-1")

        for i in range(len(data) - window_size):
            bytes_window = list(map(lambda byte: ord(byte), data[i:i+window_size]))
            bytes_window.append(label)
            of.write(data_f_string % tuple(map(lambda x: str(x), bytes_window)))


if __name__ == "__main__":
    filepaths = get_clean_and_infected_filepaths()

    for key in filepaths.keys():
        if key == "clean":
            label = "0"
        elif key == "infected":
            label = "1"
        
        for path in filepaths[key]:
            print("Reading bytes_window from path: {}".format(path))
            output_path = key + path.split(os.path.sep)[-1] + ".csv"
            read_byte_windows(path, label, output_filename=output_path)

