Initial scripts, directory structure and implementation

Author: ollycassidy13

.gitignore

@@ -0,0 +1,2 @@
+nids/__pycache__
+dataset/MachineLearningCVE

csvs.py

@@ -0,0 +1,10 @@
+import glob
+import pandas as pd
+
+csv_files = glob.glob('dataset/MachineLearningCVE/*.csv')
+print(f"Found CSV files: {csv_files}")
+
+# Load a single CSV file to inspect its columns
+df = pd.read_csv('dataset/MachineLearningCVE/Friday-WorkingHours-Afternoon-DDos.pcap_ISCX.csv')
+print(df.columns)
+

nids/__init__.py

@@ -0,0 +1,7 @@
+# nids/__init__.py
+
+from .data_preprocessing import load_and_preprocess_data
+from .model import Net, train_model
+from .logging import setup_logging, log_prediction
+from .prediction import run_prediction
+from .retraining import retrain

nids/data_preprocessing.py

@@ -0,0 +1,64 @@
+# nids/data_preprocessing.py
+
+import pandas as pd
+import glob
+from sklearn.preprocessing import StandardScaler
+from sklearn.model_selection import train_test_split
+import numpy as np
+
+def load_and_preprocess_data(csv_files_path):
+    # Load all CSV files
+    csv_files = glob.glob(csv_files_path)
+    print(f"Found CSV files: {csv_files}")  # Debug print
+    
+    if not csv_files:
+        raise FileNotFoundError(f"No CSV files found in the path: {csv_files_path}")
+    
+    dataframes = [pd.read_csv(file) for file in csv_files]
+    
+    # Concatenate all dataframes
+    data = pd.concat(dataframes, ignore_index=True)
+    print(f"Concatenated Data Shape: {data.shape}")  # Debug print
+
+    # Handle missing values
+    data = data.dropna()
+
+    # Check for the target label column
+    possible_label_columns = ['label', 'Label', 'class', 'Class', ' Label']
+    label_column = None
+    for col in possible_label_columns:
+        if col in data.columns:
+            label_column = col
+            break
+
+    if label_column is None:
+        raise ValueError("The target label column is not found in the dataset.")
+    
+    # Encode categorical variables
+    data[label_column] = data[label_column].astype('category').cat.codes
+
+    # Print unique values of the target labels
+    print(f"Unique target labels: {data[label_column].unique()}")
+
+    # Replace infinite values with NaN
+    data.replace([np.inf, -np.inf], np.nan, inplace=True)
+    
+    # Convert all columns to numeric, forcing non-numeric to NaN
+    data = data.apply(pd.to_numeric, errors='coerce')
+
+    # Drop rows with NaN values
+    data.dropna(inplace=True)
+
+    # Normalize numerical features
+    numerical_features = data.select_dtypes(include=['float64', 'int64']).columns
+    scaler = StandardScaler()
+    data[numerical_features] = scaler.fit_transform(data[numerical_features])
+    
+    # Split features and labels
+    X = data.drop(columns=[label_column])
+    y = data[label_column]
+    
+    # Split data into training and testing sets
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+
+    return X_train, X_test, y_train, y_test, scaler

nids/logging.py

@@ -0,0 +1,11 @@
+# nids/logging.py
+
+import logging
+
+def setup_logging():
+    # Configure logging
+    logging.basicConfig(filename='nids_logs.log', level=logging.INFO,
+                        format='%(asctime)s:%(levelname)s:%(message)s')
+
+def log_prediction(data, prediction):
+    logging.info(f'Data: {data}, Prediction: {prediction.item()}')

nids/model.py

@@ -0,0 +1,52 @@
+# nids/model.py
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import pandas as pd
+
+class Net(nn.Module):
+    def __init__(self, input_size, num_classes):
+        super(Net, self).__init__()
+        self.fc1 = nn.Linear(input_size, 128)
+        self.fc2 = nn.Linear(128, 64)
+        self.fc3 = nn.Linear(64, num_classes)
+
+    def forward(self, x):
+        x = torch.relu(self.fc1(x))
+        x = torch.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x
+
+def train_model(X_train, y_train):
+    # Convert data to numpy arrays before creating tensors
+    X_train_array = X_train.to_numpy()
+    y_train_array = y_train.to_numpy()
+    
+    # Convert data to PyTorch tensors
+    X_train_tensor = torch.tensor(X_train_array, dtype=torch.float32)
+    y_train_tensor = torch.tensor(y_train_array, dtype=torch.long)
+    
+    # Create DataLoader
+    train_dataset = torch.utils.data.TensorDataset(X_train_tensor, y_train_tensor)
+    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)
+    
+    # Get the number of unique classes
+    num_classes = len(pd.unique(y_train))
+    
+    # Initialize the model, loss function, and optimizer
+    model = Net(X_train.shape[1], num_classes)
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+    
+    # Training loop
+    for epoch in range(20):  # Number of epochs
+        for inputs, labels in train_loader:
+            optimizer.zero_grad()
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+        print(f'Epoch {epoch+1}/{20}, Loss: {loss.item()}')
+    
+    return model, num_classes

nids/prediction.py

@@ -0,0 +1,34 @@
+# nids/prediction.py
+
+from kafka import KafkaConsumer
+import torch
+import pandas as pd
+import json
+from nids.model import Net
+from nids.logging import log_prediction
+from sklearn.preprocessing import StandardScaler
+
+def preprocess_data(data, scaler):
+    data = pd.DataFrame([data])
+    data = pd.get_dummies(data)
+    data = scaler.transform(data)
+    return torch.tensor(data, dtype=torch.float32)
+
+def run_prediction(model, scaler):
+    # Initialize Kafka consumer
+    consumer = KafkaConsumer('network_traffic',
+                             bootstrap_servers='localhost:9092',
+                             value_deserializer=lambda v: json.loads(v.decode('utf-8')))
+    
+    model.eval()
+    # Real-time prediction loop
+    for message in consumer:
+        data = message.value
+        data_tensor = preprocess_data(data, scaler)
+        
+        # Make prediction
+        with torch.no_grad():
+            output = model(data_tensor)
+            _, prediction = torch.max(output, 1)
+            log_prediction(data, prediction)
+            print(f'Prediction: {prediction.item()}')

nids/retraining.py

@@ -0,0 +1,65 @@
+# nids/retraining.py
+
+import pandas as pd
+import pickle
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, TensorDataset
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import StandardScaler
+from nids.model import Net
+
+def load_and_preprocess_data(file_path):
+    data = pd.read_csv(file_path)
+    data = data.drop(columns=['timestamp'])
+    X = data.drop(columns=['label'])
+    y = data['label']
+    X = pd.get_dummies(X)
+    scaler = StandardScaler()
+    X = scaler.fit_transform(X)
+    return train_test_split(X, y, test_size=0.2, random_state=42), scaler
+
+def retrain_model(model, train_loader):
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+    for epoch in range(10):  # Fewer epochs for incremental training
+        for inputs, labels in train_loader:
+            optimizer.zero_grad()
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+    return model
+
+def retrain(csv_files_path):
+    (X_train, X_test, y_train, y_test), scaler = load_and_preprocess_data(csv_files_path)
+    
+    # Convert data to PyTorch tensors
+    X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
+    y_train_tensor = torch.tensor(y_train.values, dtype=torch.long)
+    train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
+    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
+    
+    # Load the existing model
+    with open('nids/model_metadata.pkl', 'rb') as f:
+        metadata = pickle.load(f)
+        num_features = metadata['num_features']
+        num_classes = metadata['num_classes']
+    
+    model = Net(num_features, num_classes)
+    model.load_state_dict(torch.load('nids/model.pth'))
+    
+    # Retrain the model
+    model = retrain_model(model, train_loader)
+    
+    # Save the updated model
+    torch.save(model.state_dict(), 'nids/updated_model.pth')
+    
+    # Save the scaler and metadata
+    with open('nids/scaler.pkl', 'wb') as f:
+        pickle.dump(scaler, f)
+    with open('nids/model_metadata.pkl', 'wb') as f:
+        pickle.dump(metadata, f)
+
+    print("Model, scaler, and metadata (number of features and classes) updated successfully.")

requirements.txt

@@ -0,0 +1,4 @@
+pandas
+scikit-learn
+torch
+kafka-python

retrain_and_run.py

@@ -0,0 +1,28 @@
+# retrain_and_run.py
+
+import torch
+import pickle
+from nids import Net, setup_logging, run_prediction, retrain
+
+if __name__ == '__main__':
+    # Setup logging
+    setup_logging()
+    
+    # Retrain the model
+    retrain('nids/new_data_directory/*.csv')
+    
+    # Load the number of features and classes
+    with open('nids/model_metadata.pkl', 'rb') as f:
+        metadata = pickle.load(f)
+        num_features = metadata['num_features']
+        num_classes = metadata['num_classes']
+    
+    # Load the updated model and scaler
+    model = Net(input_size=num_features, num_classes=num_classes)
+    model.load_state_dict(torch.load('nids/updated_model.pth'))
+    
+    with open('nids/scaler.pkl', 'rb') as f:
+        scaler = pickle.load(f)
+    
+    # Run real-time prediction
+    run_prediction(model, scaler)

run.py

@@ -0,0 +1,25 @@
+# run.py
+
+import torch
+import pickle
+from nids import Net, setup_logging, run_prediction
+
+if __name__ == '__main__':
+    # Setup logging
+    setup_logging()
+
+    # Load the number of features and classes
+    with open('nids/model_metadata.pkl', 'rb') as f:
+        metadata = pickle.load(f)
+        num_features = metadata['num_features']
+        num_classes = metadata['num_classes']
+    
+    # Load the model and scaler
+    model = Net(input_size=num_features, num_classes=num_classes)
+    model.load_state_dict(torch.load('nids/model.pth'))
+    
+    with open('nids/scaler.pkl', 'rb') as f:
+        scaler = pickle.load(f)
+    
+    # Run real-time prediction
+    run_prediction(model, scaler)

train.py

@@ -0,0 +1,27 @@
+# train.py
+
+import torch
+import pickle
+from nids import load_and_preprocess_data, train_model
+
+if __name__ == '__main__':
+    # Load and preprocess data
+    X_train, X_test, y_train, y_test, scaler = load_and_preprocess_data('dataset/MachineLearningCVE/*.csv')
+    
+    # Train the model
+    model, num_classes = train_model(X_train, y_train)
+    
+    # Save the model and scaler
+    torch.save(model.state_dict(), 'nids/model.pth')
+    with open('nids/scaler.pkl', 'wb') as f:
+        pickle.dump(scaler, f)
+    
+    # Save the number of features and classes
+    with open('nids/model_metadata.pkl', 'wb') as f:
+        metadata = {
+            'num_features': X_train.shape[1],
+            'num_classes': num_classes
+        }
+        pickle.dump(metadata, f)
+
+    print("Model, scaler, and metadata (number of features and classes) saved successfully.")