From 28fff1c9244c738fb0251563b5827e0e3e7eaf5f Mon Sep 17 00:00:00 2001
From: yulonger's Desktop <yulonger@outlook.com>
Date: Fri, 7 Jul 2023 14:36:28 +0800
Subject: [PATCH] bump feat: stable version

---
 config.py |  2 +-
 main.py   | 31 ++++++++++++++++++++++++-------
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/config.py b/config.py
index 28ebe4a..5f764bf 100644
--- a/config.py
+++ b/config.py
@@ -1,6 +1,6 @@
 import datetime
 
-CSV_PATH = './_dataset/TrafficLabelling_/Friday-WorkingHours-DDoS.csv'
+CSV_PATH = './_dataset/TrafficLabelling_/Friday-WorkingHours-merged.csv'
 BYPASS_COLUMNS= ('Destination Port', 'Label')
 UNIQUE_COLUMNS = [' Bwd PSH Flags', ' Fwd URG Flags', ' Bwd URG Flags', ' CWE Flag Count',
        'Fwd Avg Bytes/Bulk', ' Fwd Avg Packets/Bulk', ' Fwd Avg Bulk Rate',
diff --git a/main.py b/main.py
index dab8fe5..6099545 100644
--- a/main.py
+++ b/main.py
@@ -16,6 +16,7 @@ def input_csv_to_df(file_path: str) -> pd.DataFrame:
 
 
 def averaging_df(df: pd.DataFrame):
+    logger.info(f"Total: {len(df)} averaging...")
     numeric_features = df.dtypes[df.dtypes != 'object'].index
     scaler = QuantileTransformer()
     df[numeric_features] = scaler.fit_transform(df[numeric_features])
@@ -27,21 +28,37 @@ def averaging_df(df: pd.DataFrame):
 def clean_data(df: pd.DataFrame) -> pd.DataFrame:
     df = df.replace([np.inf, -np.inf], np.nan)
     df = df.dropna(axis=0)  # 删除具有NaN值的行
-    df = get_ddos_df(df)
+    # df = get_ddos_df(df)
     df = drop_columns(df, UNIQUE_COLUMNS)
     # df = drop_unique_columns(df)
     df = df.iloc[:, 7:]
     return df
 
+def slice_df(df: pd.DataFrame):
+    logger.info(f"Total: {len(df)} slicing...")
+    ddos_df = select_label_rows(df, 'DDoS')
+    normal_df = select_label_rows(df, 'BENIGN')
+    return ddos_df, normal_df
+
 
 def process(df: pd.DataFrame, label: str = None):
-    df = clean_data(df)
-    df_clean_data = averaging_df(df)
-    create_dir(IMG_SAVE_PATH)
-    generate_and_save(df_clean_data)
+
+    ddos_df, normal_df = slice_df(df)
+
+    ddos_df = averaging_df(clean_data(ddos_df))
+    normal_df = averaging_df(clean_data(normal_df))
+
+    logger.info(f"DDoS: {len(ddos_df)}, Normal: {len(normal_df)}")
+
+    ddos_save_path = f"{IMG_SAVE_PATH}/ddos"
+    benign_save_path = f"{IMG_SAVE_PATH}/benign"
+    create_dir(ddos_save_path)
+    generate_and_save(ddos_df, ddos_save_path)
+    create_dir(benign_save_path)
+    generate_and_save(normal_df, benign_save_path)
 
 
-def generate_and_save(df_clean_data: pd.DataFrame):
+def generate_and_save(df_clean_data: pd.DataFrame, save_path: str = IMG_SAVE_PATH):
     row_length = len(df_clean_data.columns)
     col_length = len(df_clean_data)
     count = 0
@@ -61,7 +78,7 @@ def generate_and_save(df_clean_data: pd.DataFrame):
                 logger.info(f"Shape: {ims.shape}")
             array = np.array(ims, dtype=np.uint8)
             new_image = Image.fromarray(array)
-            new_image.save(f"{IMG_SAVE_PATH}/{saves_count}.png")
+            new_image.save(f"{save_path}/{saves_count}.png")
             count = 0
             ims = []