NetGenRL/driver.py at master · DeepShield-AI/NetGenRL · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
import json
import os
from pre_process import pcap_process, ChiMerge,baseline
from stateFormer import model_train
from post_process import check_bestmodel, data_save

def check_and_make_forder(folder_path):
    if not os.path.exists(folder_path):
        os.makedirs(folder_path)

def construct_label_dict(pcap_folder, dataset):
    label_dict = {}
    count = 0
    for filename in os.listdir(f"./{pcap_folder}/{dataset}/"):
        # if filename.endswith('.pcap') or filename.endswith('.pcapng'):
        #     label = filename.split('.')[0]
        #     label_dict[label] = count
        #     count += 1
        if filename.endswith('.json'):
            label = filename.split('.')[0]
            label_dict[label] = count
            count += 1
    return label_dict


def driver(config):
    # with open("./config.json","r") as file:
    #     config = json.load(file)
    # print(config)

    # label_dict = construct_label_dict(config['path']['pcap_folder'],config['path']['dataset'])
    label_dict = construct_label_dict(config['path']['json_folder'],config['path']['dataset'])
    sery_attr_names = list(config['attributes']['sery_attribute'].keys())
    port_attr_names = list(config['attributes']['port_attribute'].keys())
    ip_attr_names = list(config['attributes']['ip_attribute'].keys())
    param_dicts = dict(config['attributes']['sery_attribute'])
    param_dicts.update(dict(config['attributes']['port_attribute']))
    param_dicts.update(dict(config['attributes']['ip_attribute']))

    check_and_make_forder(config['path']['json_folder'])
    check_and_make_forder(f"./{config['path']['json_folder']}/{config['path']['dataset']}")
    check_and_make_forder(config['path']['bins_folder'])
    check_and_make_forder(config['path']['model_folder'])
    check_and_make_forder(f"./{config['path']['model_folder']}/{config['path']['dataset']}")
    check_and_make_forder(config['path']['result_folder'])
    check_and_make_forder(f"./{config['path']['result_folder']}/{config['path']['dataset']}")

    # print("Processing pcap to json ...")

    # pcap_process.process_pcap(f"./{config['path']['pcap_folder']}/{config['path']['dataset']}",
    #                           label_dict,
    #                           f"./{config['path']['json_folder']}/{config['path']['dataset']}")

    print("Binning data ...")

    ChiMerge.chimerge(config['path']['dataset'],
                      config['path']['json_folder'],
                      config['path']['bins_folder'],
                      ip_attr_names,
                      port_attr_names,
                      sery_attr_names,
                      config['model_paras']['max_seq_len'],
                      param_dicts)
    # baseline.basemerge(config['path']['dataset'],
    #                     config['path']['json_folder'],
    #                     config['path']['bins_folder'],
    #                     ip_attr_names,
    #                     port_attr_names,
    #                     sery_attr_names,
    #                     config['model_paras']['max_seq_len'],
    #                     param_dicts)

    print("Training model ...")

    model_train.model_train(label_dict,
                      config['path']['dataset'],
                      config['path']['json_folder'],
                      config['path']['bins_folder'],
                      config['path']['model_folder'],
                      port_attr_names + ip_attr_names,
                      sery_attr_names,
                      config['model_paras'])

    print("Choosing best model ...")
    model_id = check_bestmodel.check_models(label_dict,
                                            config['path']['dataset'],
                                            config['path']['json_folder'],
                                            config['path']['bins_folder'],
                                            config['path']['model_folder'],
                                            port_attr_names + ip_attr_names,
                                            sery_attr_names,
                                            config['model_paras']['batch_size'],
                                            config['model_paras']['max_seq_len'],
                                            config['model_paras']['checkpoint'],
                                            config['model_paras']['epoch'])


    print("Generating data ...")
    data_save.generate_data(label_dict,
                            config['path']['dataset'],
                            config['path']['json_folder'],
                            config['path']['bins_folder'],
                            config['path']['model_folder'],
                            config['path']['result_folder'],
                            port_attr_names + ip_attr_names,
                            sery_attr_names,
                            config['model_paras']['batch_size'],
                            config['model_paras']['max_seq_len'],
                            config['model_paras']['checkpoint'],
                            model_id,
                            config['model_paras']['expand_times'])

def read_config():
    with open("./config.json","r") as file:
        config = json.load(file)
    return config

if __name__ == "__main__":
    config = read_config()
    # DATASETS = [
    #     "iscx_3", "ssrc_3"
    # ]
    # for dataset in DATASETS:
    #     print(f"Begin processing dataset: {dataset} ----------------------\n\n")
    #     config['path']['dataset'] = dataset
    #     driver(config)
    #     print(f"Finished processing dataset: {dataset} ----------------------\n\n")
    driver(config)