Add patches and scripts needed to run on Octeon

README.md

@@ -72,7 +72,7 @@ cd prepro &&　sh multi_gen.sh ../data/JSON Table.csv 0 && python3 toTrain.py Ta
 To train a model, do:
 
 ```shell
-# To train the model, you should specify the mode(DNN/XGB):
+# To train the model, you should specify the mode:
 cd ./main && sh TRAIN.sh DNN
 ```
 

main/train_updated.py

@@ -0,0 +1,169 @@
+import os
+import numpy as np
+import argparse
+import matplotlib.pyplot as plt
+from sklearn.metrics import (
+    confusion_matrix, f1_score, precision_score, recall_score
+)
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import normalize
+from tensorflow.keras.utils import to_categorical
+from tensorflow.keras.models import Model
+from tensorflow.keras.layers import Input, Dense, Dropout, BatchNormalization, ReLU
+from tensorflow.keras.optimizers import SGD
+from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau
+
+# Label mappings
+LABEL2DIG = {'chat': 0, 'voip': 1, 'trap2p': 2, 'stream': 3, 'file_trans': 4, 'email': 5,
+             'vpn_chat': 6, 'vpn_voip': 7, 'vpn_trap2p': 8, 'vpn_stream': 9, 'vpn_file_trans': 10, 'vpn_email': 11}
+DIG2LABEL = {v: k for k, v in LABEL2DIG.items()}
+nclass = len(LABEL2DIG)
+
+
+def build_dnn(input_dim, dropout_rates):
+    inputs = Input(shape=(input_dim,))
+    x = Dense(1280)(inputs)
+    x = BatchNormalization()(x)
+    x = ReLU()(x)
+    x = Dropout(dropout_rates[0])(x)
+
+    x = Dense(960)(x)
+    x = BatchNormalization()(x)
+    x = ReLU()(x)
+    x = Dropout(dropout_rates[1])(x)
+
+    x = Dense(640)(x)
+    x = BatchNormalization()(x)
+    x = ReLU()(x)
+    x = Dropout(dropout_rates[2])(x)
+
+    x = Dense(640)(x)
+    x = BatchNormalization()(x)
+    x = ReLU()(x)
+    x = Dropout(dropout_rates[3])(x)
+
+    x = Dense(480)(x)
+    x = BatchNormalization()(x)
+    x = ReLU()(x)
+    x = Dropout(dropout_rates[4])(x)
+
+    x = Dense(320)(x)
+    x = BatchNormalization()(x)
+    x = ReLU()(x)
+
+    x = Dense(320)(x)
+    x = BatchNormalization()(x)
+    x = ReLU()(x)
+
+    outputs = Dense(nclass, activation='softmax')(x)
+    return Model(inputs=inputs, outputs=outputs)
+
+
+def plot_confusion(cm, class_names, normalize=False, title="Confusion Matrix", path="cm.png"):
+    if normalize:
+        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
+
+    plt.figure(figsize=(10, 10))
+    plt.imshow(cm, interpolation='nearest', cmap=plt.cm.Blues)
+    plt.title(title)
+    plt.colorbar()
+    tick_marks = np.arange(len(class_names))
+    plt.xticks(tick_marks, class_names, rotation=90)
+    plt.yticks(tick_marks, class_names)
+
+    fmt = '.2f' if normalize else 'd'
+    thresh = cm.max() / 2.
+    for i in range(len(cm)):
+        for j in range(len(cm[i])):
+            plt.text(j, i, format(cm[i, j], fmt),
+                     ha="center", va="center",
+                     color="white" if cm[i, j] > thresh else "black")
+
+    plt.tight_layout()
+    plt.ylabel('True label')
+    plt.xlabel('Predicted label')
+    plt.savefig(path)
+    plt.close()
+
+
+def DNN(opts):
+    # Load data
+    X = np.load(opts.source_data_folder + '/X_train.npy', allow_pickle=True)
+    y = np.load(opts.source_data_folder + '/y_train.npy', allow_pickle=True)
+
+    print('Loaded data:', X.shape, y.shape)
+
+    # Normalize
+    X = normalize(X, norm='max', axis=0)
+
+    # One-hot encode labels
+    y_categorical = to_categorical(y, num_classes=nclass)
+
+    # Split
+    X_train, X_test, y_train, y_test = train_test_split(
+        X, y_categorical, test_size=0.15, random_state=42)
+
+    # Build model
+    dropout_rates = [0.5, 0.5, 0.5, 0.3, 0.2]
+    model = build_dnn(X.shape[1], dropout_rates)
+
+    model.compile(
+        optimizer=SGD(learning_rate=0.01, momentum=0.9),
+        loss='categorical_crossentropy',
+        metrics=['accuracy']
+    )
+
+    os.makedirs(opts.output_folder, exist_ok=True)
+
+    # Callbacks
+    callbacks = [
+        ModelCheckpoint(os.path.join(opts.output_folder,
+                        "model.keras"), save_best_only=True),
+        EarlyStopping(patience=opts.patience, restore_best_weights=True),
+        ReduceLROnPlateau(patience=5, factor=0.5)
+    ]
+
+    model.fit(
+        X_train, y_train,
+        validation_data=(X_test, y_test),
+        epochs=100,
+        batch_size=opts.batch_size,
+        callbacks=callbacks,
+        verbose=1
+    )  # to onnx
+
+    model.save(os.path.join(opts.output_folder, "final_model.keras"))
+
+    # Evaluation
+    y_pred = model.predict(X_test).argmax(axis=1)
+    y_true = y_test.argmax(axis=1)
+
+    print("F1 score:", f1_score(y_true, y_pred, average='macro'))
+    print("Precision:", precision_score(y_true, y_pred, average=None))
+    print("Recall:", recall_score(y_true, y_pred, average=None))
+
+    class_names = [DIG2LABEL[i] for i in range(nclass)]
+    cm = confusion_matrix(y_true, y_pred)
+    plot_confusion(cm, class_names, False, "Confusion Matrix",
+                   os.path.join(opts.output_folder, 'confusion_matrix.png'))
+    plot_confusion(cm, class_names, True, "Normalized Confusion Matrix",
+                   os.path.join(opts.output_folder, 'confusion_matrix_norm.png'))
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(description='Modern DNN Training')
+    parser.add_argument('--mode', type=str, default='DNN',
+                        choices=['DNN', 'XGB'])
+    parser.add_argument('--source_data_folder', type=str, default='../data')
+    parser.add_argument('--output_folder', type=str, default='./output')
+    parser.add_argument('--batch_size', type=int, default=1024)
+    parser.add_argument('--patience', type=int, default=10)
+    return parser.parse_args()
+
+
+if __name__ == '__main__':
+    opts = parse_arguments()
+    if opts.mode == 'DNN':
+        DNN(opts)
+    else:
+        print("XGB mode not modernized yet.")

patches/joy.patch

@@ -0,0 +1,30 @@
+diff --git a/sleuth b/sleuth
+index 2fc96ed..30f0e41 100755
+--- a/sleuth
++++ b/sleuth
+@@ -1,4 +1,4 @@
+-#!/usr/bin/python
++#!/usr/bin/python2.7
+
+ """
+ sleuth performs query operations on JSON-formatted flow objects; see sleuth --help for more details
+@@ -40,15 +40,16 @@ sleuth performs query operations on JSON-formatted flow objects; see sleuth --he
+
+ import sys
+ import argparse
+-sys.path.append("/usr/local/lib/python")
++sys.path.append("/usr/local/lib/python2.7")
+
+ try:
+     from sleuth import *
++    from sleuth_pkg.flow import FlowIteratorFromFile
+ except ImportError:
+     # If in source code directory
+     from sleuth_pkg import *
+-    
+-    
++
++
+ def attach_iterators(source, args):
+     source = DNSLinkedFlowEnrichIterator(source)
+

patches/traffic.patch

@@ -0,0 +1,79 @@
+diff --git a/main/TRAIN.sh b/main/TRAIN.sh
+old mode 100644
+new mode 100755
+index 612ba0d..930230a
+--- a/main/TRAIN.sh
++++ b/main/TRAIN.sh
+@@ -1 +1 @@
+-python3 train.py --mode $1 --source_data_folder ../data --output_folder TEST1212 --patience 5
++python3 train_updated.py --mode $1 --source_data_folder ../data --output_folder TEST1212 --patience 5
+diff --git a/prepro/Table_Generator.py b/prepro/Table_Generator.py
+index 35016a2..1085c5c 100644
+--- a/prepro/Table_Generator.py
++++ b/prepro/Table_Generator.py
+@@ -4,7 +4,8 @@
+ # In[48]:
+
+
+-from pandas.io.json import json_normalize
++#from pandas.io.json import json_normalize
++from pandas import json_normalize
+ import pandas as pd
+ import numpy as np
+ import json
+diff --git a/prepro/toTrain.py b/prepro/toTrain.py
+index 4f155d8..fb3091d 100644
+--- a/prepro/toTrain.py
++++ b/prepro/toTrain.py
+@@ -18,24 +18,33 @@ if __name__=="__main__":
+     #create column vpn
+     df['vpn'] = df.file_name.str.contains('vpn').astype(int)
+
+-    df = df[(df.type!='browsing')&(df.file_name!='skype_audio1a_test')]
+-    # Broadcast and Multicast
+-    df = df[~(df.da.str.contains('224.0.'))]
+-    df = df[~(df.da.str.contains('239.255.'))]
+-    df = df[~(df.da.str.contains('255.255.'))]
+-    # Zzro packet size flow
+-    df = df[(df.formean!=0)|(df.backmean!=0)]
+-    # file transfer should not use udp
+-    df = df[~((df.type=='file_trans')&(df.pr==17))]
+-    # clean icmp
+-    df = df[df.pr!=1]
+-    # clean NetBIOS
+-    df = df[(df.sp>139)|(df.sp<137)]
+-    # one packet flow
+-    df = df[(df.tot_forpkts + df.tot_backpkts) > 1]
+-
+-    #get y
+-    type_ = df.type.values
++    # Filter out specific types and files
++    df = df[(df['type'] != 'browsing') & (
++    df['file_name'] != 'skype_audio1a_test')]
++
++    # Remove broadcast and multicast traffic
++    df = df[~df['da'].astype(str).str.contains('224.0.', na=False)]
++    df = df[~df['da'].astype(str).str.contains('239.255.', na=False)]
++    df = df[~df['da'].astype(str).str.contains('255.255.', na=False)]
++
++    # Remove zero-packet-size flows
++    df = df[(df['formean'] != 0) | (df['backmean'] != 0)]
++
++    # Remove file transfers using UDP (pr == 17)
++    df = df[~((df['type'] == 'file_trans') & (df['pr'] == 17))]
++
++    # Remove ICMP traffic (pr == 1)
++    df = df[df['pr'] != 1]
++
++    # Remove NetBIOS traffic (sp between 137 and 139 inclusive)
++    df = df[(df['sp'] > 139) | (df['sp'] < 137)]
++
++    # Remove one-packet flows
++    df = df[(df['tot_forpkts'] + df['tot_backpkts']) > 1]
++
++    # Map labels to digits
++    type_ = df['type'].values
++
+     type_ = [LABEL2DIG[label] for label in type_]
+
+     y_train = df.vpn.values * 6 + type_

prepro/generate_inference_files.py

@@ -0,0 +1,75 @@
+import pandas as pd
+import numpy as np
+import os
+import sys
+
+CSV_INPUT   = sys.argv[1] 
+COL_SAMPLE  = "ColSample.csv"     # from TRAINING
+X_TRAIN_NPY = "../data/X_train.npy"         # from TRAINING
+OUT_DIR     = "inference_bins"           # output folder
+
+os.makedirs(OUT_DIR, exist_ok=True)
+
+print("[*] Loading CSV:", CSV_INPUT)
+df = pd.read_csv(CSV_INPUT)
+
+df["vpn"] = df.file_name.astype(str).str.contains("vpn").astype(int)
+
+df = df[(df["type"] != "browsing") &
+        (df["file_name"] != "skype_audio1a_test")]
+df = df[~((df["type"] == "file_trans") & (df["pr"] == 17))]
+
+df = df[~df["da"].astype(str).str.contains("224.0.", na=False)]
+df = df[~df["da"].astype(str).str.contains("239.255.", na=False)]
+df = df[~df["da"].astype(str).str.contains("255.255.", na=False)]
+
+df = df[(df["formean"] != 0) | (df["backmean"] != 0)]
+df = df[df["pr"] != 1]
+df = df[(df["sp"] > 139) | (df["sp"] < 137)]
+df = df[(df["tot_forpkts"] + df["tot_backpkts"]) > 1]
+
+if len(df) == 0:
+    raise RuntimeError("No valid flows after filtering")
+
+print("[+] Flows after filtering:", len(df))
+
+df_feat = df.drop(
+    ["sp", "dp", "sa", "da", "type", "file_name", "label", "vpn"],
+    axis=1,
+    errors="ignore"
+)
+
+table = pd.get_dummies(df_feat)
+
+col_sample = pd.read_csv(COL_SAMPLE)
+table = table.reindex(columns=col_sample.columns, fill_value=0)
+
+X = table.values.astype(np.float32)
+
+X_train = np.load(X_TRAIN_NPY, allow_pickle=True).astype(np.float32)
+max_vec = X_train.max(axis=0)
+max_vec[max_vec == 0] = 1.0   # avoid divide-by-zero
+
+X_norm = X / max_vec
+
+expected_features = X_train.shape[1]
+
+
+if X.shape[1] != expected_features:
+    raise RuntimeError(
+        f"Feature mismatch: {X.shape[1]} != {expected_features}"
+    )
+
+# sanity check
+print("[*] Normalization stats:")
+print("    max =", X_norm.max())
+print("    mean =", X_norm.mean())
+print("    non-zero =", (X_norm != 0).sum())
+
+for i in range(X_norm.shape[0]):
+    out = os.path.join(OUT_DIR, f"input_{i:04d}.bin")
+    X_norm[i].astype(np.float32).tofile(out)
+
+print(f"[+] {X_norm.shape[0]} BIN files written to {OUT_DIR}/")
+
+print("DONE")

requirements.txt

@@ -0,0 +1,8 @@
+keras
+numpy
+matplotlib
+scikit-learn
+xgboost
+argparse
+pandas
+tensorflow