diff --git a/requirements.txt b/requirements.txt index a8ab8f55..ebe280de 100644 --- a/requirements.txt +++ b/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 \ No newline at end of file diff --git a/securetea/lib/waf/Server/classifier.py b/securetea/lib/waf/Server/classifier.py index 8b814a31..19571896 100644 --- a/securetea/lib/waf/Server/classifier.py +++ b/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) diff --git a/securetea/lib/waf/Server/reqHandler.py b/securetea/lib/waf/Server/reqHandler.py index 553edc1c..9ea41070 100644 --- a/securetea/lib/waf/Server/reqHandler.py +++ b/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"]) diff --git a/securetea/lib/waf/data/model b/securetea/lib/waf/data/model new file mode 100644 index 00000000..9e67696e Binary files /dev/null and b/securetea/lib/waf/data/model differ diff --git a/train.py b/train.py new file mode 100644 index 00000000..ac82cf27 --- /dev/null +++ b/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() \ No newline at end of file