run method

Author: ollycassidy13

.gitignore

@@ -1,2 +1,8 @@
 nids/__pycache__
-dataset/MachineLearningCVE
+dataset/MachineLearningCVE
+nids_logs.log
+packet_sniffer.log
+nids/model_metadata.pkl
+nids/model.pth
+nids/scaler.pkl
+nids/training_test_accuracy.png

nids/flow.py

@@ -0,0 +1,59 @@
+# nids/flow.py
+
+import time
+import pandas as pd
+from scapy.all import IP, IPv6, TCP, UDP
+
+class Flow:
+    def __init__(self, packet):
+        self.packets = [packet]
+        self.start_time = time.time()
+        self.end_time = time.time()
+        if IP in packet:
+            self.src_ip = packet[IP].src
+            self.dst_ip = packet[IP].dst
+            self.src_port = packet[IP].sport if TCP in packet or UDP in packet else 0
+            self.dst_port = packet[IP].dport if TCP in packet or UDP in packet else 0
+        elif IPv6 in packet:
+            self.src_ip = packet[IPv6].src
+            self.dst_ip = packet[IPv6].dst
+            self.src_port = packet[IPv6].sport if TCP in packet or UDP in packet else 0
+            self.dst_port = packet[IPv6].dport if TCP in packet or UDP in packet else 0
+        self.total_fwd_packets = 1 if self.src_ip == (packet[IP].src if IP in packet else packet[IPv6].src) else 0
+        self.total_bwd_packets = 1 if self.src_ip == (packet[IP].dst if IP in packet else packet[IPv6].dst) else 0
+        self.total_fwd_bytes = len(packet) if self.src_ip == (packet[IP].src if IP in packet else packet[IPv6].src) else 0
+        self.total_bwd_bytes = len(packet) if self.src_ip == (packet[IP].dst if IP in packet else packet[IPv6].dst) else 0
+
+    def update(self, packet):
+        self.packets.append(packet)
+        self.end_time = time.time()
+        if self.src_ip == (packet[IP].src if IP in packet else packet[IPv6].src):
+            self.total_fwd_packets += 1
+            self.total_fwd_bytes += len(packet)
+        else:
+            self.total_bwd_packets += 1
+            self.total_bwd_bytes += len(packet)
+
+    def get_duration(self):
+        return (self.end_time - self.start_time) * 1e6  # duration in microseconds
+
+    def get_features(self):
+        features = {
+            'Destination Port': self.dst_port,
+            'Flow Duration': self.get_duration(),
+            'Total Fwd Packets': self.total_fwd_packets,
+            'Total Backward Packets': self.total_bwd_packets,
+            'Total Length of Fwd Packets': self.total_fwd_bytes,
+            'Total Length of Bwd Packets': self.total_bwd_bytes,
+            'Fwd Packet Length Max': max([len(p) for p in self.packets if self.src_ip == (p[IP].src if IP in p else p[IPv6].src)], default=0),
+            'Fwd Packet Length Min': min([len(p) for p in self.packets if self.src_ip == (p[IP].src if IP in p else p[IPv6].src)], default=0),
+            'Fwd Packet Length Mean': sum([len(p) for p in self.packets if self.src_ip == (p[IP].src if IP in p else p[IPv6].src)]) / self.total_fwd_packets if self.total_fwd_packets > 0 else 0,
+            'Fwd Packet Length Std': pd.Series([len(p) for p in self.packets if self.src_ip == (p[IP].src if IP in p else p[IPv6].src)]).std(),
+            'Bwd Packet Length Max': max([len(p) for p in self.packets if self.src_ip == (p[IP].dst if IP in p else p[IPv6].dst)], default=0),
+            'Bwd Packet Length Min': min([len(p) for p in self.packets if self.src_ip == (p[IP].dst if IP in p else p[IPv6].dst)], default=0),
+            'Bwd Packet Length Mean': sum([len(p) for p in self.packets if self.src_ip == (p[IP].dst if IP in p else p[IPv6].dst)]) / self.total_bwd_packets if self.total_bwd_packets > 0 else 0,
+            'Bwd Packet Length Std': pd.Series([len(p) for p in self.packets if self.src_ip == (p[IP].dst if IP in p else p[IPv6].dst)]).std(),
+            'Flow Bytes/s': (self.total_fwd_bytes + self.total_bwd_bytes) / self.get_duration() * 1e6 if self.get_duration() > 0 else 0,
+            'Flow Packets/s': (self.total_fwd_packets + self.total_bwd_packets) / self.get_duration() * 1e6 if self.get_duration() > 0 else 0,
+        }
+        return features

nids/logging.py

@@ -1,11 +1,22 @@
 # nids/logging.py
 
 import logging
+from scapy.all import IP, IPv6
+import os
 
 def setup_logging():
-    # Configure logging
-    logging.basicConfig(filename='nids_logs.log', level=logging.INFO,
+    # Clear the log file if it exists
+    log_file = 'nids_logs.log'
+    if os.path.exists(log_file):
+        with open(log_file, 'w'):
+            pass
+    
+    # Setup logging configuration
+    logging.basicConfig(filename=log_file, level=logging.INFO,
                         format='%(asctime)s:%(levelname)s:%(message)s')
 
-def log_prediction(data, prediction):
-    logging.info(f'Data: {data}, Prediction: {prediction.item()}')
+def log_prediction(packet, prediction, original_data, traffic_type, src_ip, dst_ip):
+    summary = packet.summary() if IP in packet or IPv6 in packet else "Non-IP packet"
+    features_str = ', '.join([f'{k}: {v}' for k, v in original_data.items()])
+    log_message = f'Packet: {summary}, Prediction: {prediction.item()} ({traffic_type}), Source IP: {src_ip}, Destination IP: {dst_ip}, Features: [{features_str}]'
+    logging.info(log_message)

nids/prediction.py

@@ -1,34 +1,63 @@
 # 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 nids.flow import Flow
 from sklearn.preprocessing import StandardScaler
+from scapy.all import IP, TCP, UDP, IPv6
+import pickle
 
-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')))
+# Load the metadata including class mapping and feature names
+with open('nids/model_metadata.pkl', 'rb') as f:
+    metadata = pickle.load(f)
+    LABEL_TO_TRAFFIC_TYPE = metadata['class_mapping']
+    FEATURE_NAMES = metadata['feature_names']
+
+flows = {}
+
+def preprocess_packet(packet, scaler):
+    if IP in packet:
+        flow_key = (packet[IP].src, packet[IP].dst, packet[IP].sport, packet[IP].dport)
+    elif IPv6 in packet:
+        flow_key = (packet[IPv6].src, packet[IPv6].dst, packet[IPv6].sport, packet[IPv6].dport)
+    else:
+        return None, None, None, None
+
+    if flow_key not in flows:
+        flows[flow_key] = Flow(packet)
+    else:
+        flows[flow_key].update(packet)
 
+    flow = flows[flow_key]
+    features = flow.get_features()
+
+    # Create DataFrame with consistent feature names
+    df = pd.DataFrame([features])
+    df = pd.get_dummies(df)
+    df = df.reindex(columns=scaler.feature_names_in_, fill_value=0)  # Ensure consistent feature order
+
+    scaled_data = scaler.transform(df)
+    src_ip = flow.src_ip
+    dst_ip = flow.dst_ip
+    return torch.tensor(scaled_data, dtype=torch.float32), features, src_ip, dst_ip
+
+def run_prediction(packet, model, scaler):
     model.eval()
-    # Real-time prediction loop
-    for message in consumer:
-        data = message.value
-        data_tensor = preprocess_data(data, scaler)
+    try:
+        data_tensor, original_data, src_ip, dst_ip = preprocess_packet(packet, scaler)
+        if data_tensor is None:
+            print(f"Packet ignored: {packet.summary()}")
+            return
 
         # Make prediction
         with torch.no_grad():
             output = model(data_tensor)
             _, prediction = torch.max(output, 1)
-            log_prediction(data, prediction)
-            print(f'Prediction: {prediction.item()}')
+            traffic_type = LABEL_TO_TRAFFIC_TYPE.get(prediction.item(), "Unknown")
+        
+        # Log detailed information
+        log_prediction(packet, prediction, original_data, traffic_type, src_ip, dst_ip)
+        print(f'Prediction: {prediction.item()} ({traffic_type}), Source IP: {src_ip}, Destination IP: {dst_ip}')
+    except Exception as e:
+        print(f'Error processing packet: {e}')

nids/retraining.py

@@ -8,7 +8,7 @@
 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
+from nids import Net
 
 def load_and_preprocess_data(file_path):
     data = pd.read_csv(file_path)

run.py

@@ -2,24 +2,31 @@
 
 import torch
 import pickle
-from nids import Net, setup_logging, run_prediction
+from scapy.all import sniff, IP, IPv6
+from nids import setup_logging
+from nids.model import Net
+from nids.prediction import run_prediction
 
-if __name__ == '__main__':
-    # Setup logging
-    setup_logging()
+# Setup logging
+setup_logging()
+
+# Load the model and scaler
+with open('nids/model_metadata.pkl', 'rb') as f:
+    metadata = pickle.load(f)
+
+model = Net(input_size=metadata['num_features'], num_classes=metadata['num_classes'])
+model.load_state_dict(torch.load('nids/model.pth'))
+model.eval()
 
-    # 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)
+with open('nids/scaler.pkl', 'rb') as f:
+    scaler = pickle.load(f)
+
+def packet_handler(packet):
+    if IP or IPv6 in packet:
+        run_prediction(packet, model, scaler)
+    else:
+        print("Non-IP/IPv6 packet ignored")
+
+if __name__ == '__main__':
+    print("Starting packet capture...")
+    sniff(prn=packet_handler, store=0)  # prn specifies the function to apply to each packet