Merge pull request #405 from Sarthak-Biswas/Fixes

Fixed unfitted gaussianDB problem

requirements.txt

@@ -97,3 +97,4 @@ Werkzeug>=2.1.2
 wget>=3.2
 XlsxWriter>=3.0.3
 yara-python>=4.2.0
+tensorflow>=2.12.0

securetea/lib/waf/Server/classifier.py

@@ -42,14 +42,13 @@ def __init__(self,live_data):
         A class that initialise the required variables
 
         """
-
-
-        self.live_data=[live_data]
+        # Takes 2 live data
+        self.live_data=[[live_data[0], live_data[2]]]
 
 
 
         datapath = Path(os.path.dirname(__file__)).parent / "data/data_updated.csv"
-        modelpath = Path(os.path.dirname(__file__)).parent / "data/modeltestgram.sav"
+        modelpath = Path(os.path.dirname(__file__)).parent / "data/model"
 
 
 
@@ -63,12 +62,14 @@ def __init__(self,live_data):
 
         self.target = self.data["label"]
         self.path_vectorizer = TfidfVectorizer(tokenizer=get3Grams,encoding="cp1252")
+
+        # TODO: body currently contains no useful values, for which predictions are False. Will be updated soon.
         self.body_vectorizer = TfidfVectorizer(tokenizer=get3Grams,encoding="cp1252")
         self.model=GaussianNB()
 
         # Feature selection
 
-        self.X=self.data[['path','body','path_len']]
+        self.X=self.data[['path','path_len']]
 
         self.X_train,self.X_test,self.Y_train,self.Y_test=train_test_split(self.X,self.target,test_size=0.2)
 
@@ -90,7 +91,7 @@ def train_model(self):
 
         # Column Transformer
 
-        self.column_transformer = ColumnTransformer([('tf-1', self.path_vectorizer, 'path'),('tf-2',self.body_vectorizer,'body'),], remainder='passthrough', sparse_threshold=0)
+        self.column_transformer = ColumnTransformer([('tf-1', self.path_vectorizer, 'path')], remainder='passthrough', sparse_threshold=0)
 
         # Creating Pipeline
 
@@ -104,12 +105,12 @@ def train_model(self):
 
         self.pipe.fit(self.X_train,self.Y_train)
 
-        with open("model", "wb") as f:
-            pickle.dump(self.pipe, f)
-
+        joblib.dump(self.pipe, self.MODEL_PATH)
 
-        self.ddos = TrainDDoS()
-        self.ddos.train()
+
+        # TODO: will be implemented soon   
+        """ self.ddos = TrainDDoS()
+         self.ddos.train() """
 
 
 
@@ -139,7 +140,7 @@ def predict_model(self):
         except Exception as E:
             print(E)
 
-        self.live_df = pd.DataFrame(self.live_data,columns=['path','body','path_len','special_char','whitespaces'])
+        self.live_df = pd.DataFrame(self.live_data,columns=['path','path_len'])
         return self.model.predict(self.live_df)
 
 

securetea/lib/waf/Server/reqHandler.py

@@ -138,34 +138,34 @@ def data_received(self, data):
 
                 # Model to detect DDoS attacks
 
-                self.ddos = DetectDDoS(self.feature_value)
-                self.predicted_ddos = self.ddos.predict()
+                # self.ddos = DetectDDoS(self.feature_value)
+                # self.predicted_ddos = self.ddos.predict()
 
                 # Blocks if detects ddos
 
-                if self.predicted_ddos == 1:
+                # if self.predicted_ddos == 1:
 
-                    message="DoS from {}".format(headers["X-Real-IP"])
+                #     message="DoS from {}".format(headers["X-Real-IP"])
 
-                    self.logger.log(
-                        message,
-                        logtype="warning"
-                    )
+                #     self.logger.log(
+                #         message,
+                #         logtype="warning"
+                #     )
 
-                self.transport.close()
-                self.waflogger.write_log(message)
+                # # self.transport.close()
+                # self.waflogger.write_log(message)
 
                 # Based on mode Block or Log Request
 
                 if self.mode==0 and predicted_value[0]==1:
 
-                    # Inserts block status, attacker IP, payload to the database for frontend
-                    db_log = {'blocked': 0,
+                    # TODO: Inserts block status, attacker IP, payload to the database for frontend
+                    """ db_log = {'blocked': 0,
                               'From': headers["X-Real-IP"],
                               'payload': path}
                     
                     dataObj = DatabaseLogs(db_log)
-                    dataObj.insert_log()
+                    dataObj.insert_log() """
 
                     # Log the file and send the Request
                     message="Attack Detected from :{} Payload:{}".format(headers["X-Real-IP"],path)
@@ -179,13 +179,13 @@ def data_received(self, data):
 
                 if self.mode==1 and predicted_value[0]==1:
 
-                    # Inserts block status, attacker IP, payload to the database for frontend
-                    db_log = {'blocked': 1,
+                    # TODO: Inserts block status, attacker IP, payload to the database for frontend
+                    """ db_log = {'blocked': 1,
                               'From': headers["X-Real-IP"],
                               'payload': path}
                     
                     dataObj = DatabaseLogs(db_log)
-                    dataObj.insert_log()
+                    dataObj.insert_log() """
 
                     # Reset the Request
                     message="Attack Detected ! Request Blocked from :{}".format(headers["X-Real-IP"])

train.py

@@ -0,0 +1,6 @@
+from securetea.lib.waf.Server.classifier import WAF
+
+# Train the WAF classifier in case of any error occurs.
+
+train = WAF([None, None, None])
+train.train_model()