import sys

import getopt

import os

import multiprocessing

def list_all_file(path):

    """

    list all files of a directory

    :param path:

    :return:

    """

    file_list = []

    for (path, _, files) in os.walk(path):

        for filename in files:

            file_list.append(os.path.join(path, filename))

    return file_list

def process_file(filename, is_black):

    suffix = os.path.splitext(filename)[-1][1:]

    if suffix != "pcap":

        return

    if is_black:

        cmd = "python extract_tls_flow4.py -vr {} -o black/{}.txt  >logs/black/{}.log".format(filename, os.path.basename(filename), os.path.basename(filename))

    else:

        cmd = "python extract_tls_flow4.py -vr {} -o white/{}.txt  >logs/white/{}.log".format(filename, os.path.basename(filename), os.path.basename(filename))

    os.system(cmd)

def process_black_file(filename):

    process_file(filename, 1)

def process_white_file(filename):

    process_file(filename, 0)

def process_dir(sample_dir, is_black):

    file_list = list_all_file(sample_dir)

    process_num = 30

    pool = multiprocessing.Pool(processes=process_num)

    if is_black:

        pool.map(process_black_file, file_list)

    else:

        pool.map(process_white_file, file_list)

    pool.close()

    pool.join()

    print("End...........")

black_sample_dir = "/opt/data/samples/black_pcap"

white_sample_dir = "/opt/data/samples/white_pcap"

process_dir(black_sample_dir, 1)

process_dir(white_sample_dir, 0)