Create app.py

app.py

@@ -0,0 +1,312 @@
+import numpy as np
+import xgboost as xgb
+import gradio as gr
+from scapy.all import rdpcap
+from collections import defaultdict
+import os
+
+def transform_new_input(new_input):
+    #Scale input features based on predetermined min/max values
+    scaled_min = np.array([
+        1.0, 10.0, 856.0, 5775.0, 42.0, 26.0, 0.0, 278.0, 4.0, 1.0, 
+        -630355.0, 4.0, 50.0
+    ])
+    
+    scaled_max = np.array([
+        4.0, 352752.0, 271591638.0, 239241314.0, 421552.0, 3317.0, 
+        6302708.0, 6302708.0, 5.0, 5.0, 1746749.0, 608.0, 1012128.0
+    ])
+    
+    new_input = np.array(new_input)
+    scaled_input = (new_input - scaled_min) / (scaled_max - scaled_min)
+    return scaled_input
+
+class PcapProcessor:
+    def __init__(self, pcap_file):
+        #Initialize PCAP processor with file path
+        self.packets = rdpcap(pcap_file)
+        self.start_time = None
+        self.port_stats = defaultdict(lambda: {
+            'rx_packets': 0,
+            'rx_bytes': 0,
+            'tx_packets': 0,
+            'tx_bytes': 0,
+            'first_seen': None,
+            'last_seen': None,
+            'active_flows': set(),
+            'packets_matched': 0
+        })
+        
+    def process_packets(self, window_size=60):
+        #Process all packets and extract features
+        if not self.packets:
+            return []
+            
+        self.start_time = float(self.packets[0].time)
+        
+        # Process each packet
+        for packet in self.packets:
+            current_time = float(packet.time)
+            
+            if 'TCP' in packet or 'UDP' in packet:
+                try:
+                    src_port = packet.sport
+                    dst_port = packet.dport
+                    pkt_size = len(packet)
+                    
+                    # Track flow information
+                    flow_tuple = (packet['IP'].src, packet['IP'].dst, 
+                                src_port, dst_port)
+                    
+                    # Update port statistics
+                    self._update_port_stats(src_port, pkt_size, True, 
+                                          current_time, flow_tuple)
+                    self._update_port_stats(dst_port, pkt_size, False, 
+                                          current_time, flow_tuple)
+                except Exception as e:
+                    print(f"Error processing packet {packet}: {str(e)}")
+                    continue
+        
+        # Extract features for each port
+        features_list = []
+        for port, stats in self.port_stats.items():
+            if stats['first_seen'] is not None:
+                features = self._extract_port_features(port, stats, window_size)
+                features_list.append(features)
+        
+        return features_list
+    
+    def _update_port_stats(self, port, pkt_size, is_source, current_time, 
+                          flow_tuple):
+        #Update statistics for a given port
+        stats = self.port_stats[port]
+        
+        if stats['first_seen'] is None:
+            stats['first_seen'] = current_time
+        
+        stats['last_seen'] = current_time
+        
+        if is_source:
+            stats['tx_packets'] += 1
+            stats['tx_bytes'] += pkt_size
+        else:
+            stats['rx_packets'] += 1
+            stats['rx_bytes'] += pkt_size
+        
+        stats['active_flows'].add(flow_tuple)
+        stats['packets_matched'] += 1
+    
+    def _extract_port_features(self, port, stats, window_size):
+        #Extract the 13 features needed for the IDS model
+        port_alive_duration = stats['last_seen'] - stats['first_seen']
+        delta_alive_duration = min(port_alive_duration, window_size)
+        
+        # Calculate rates and loads
+        total_load = (stats['rx_bytes'] + stats['tx_bytes']) / \
+                    max(port_alive_duration, 1)
+        
+        features = [
+            min(port % 4 + 1, 4),  # Port Number (1-4)
+            stats['rx_packets'],    # Received Packets
+            stats['rx_bytes'],      # Received Bytes
+            stats['tx_bytes'],      # Sent Bytes
+            stats['tx_packets'],    # Sent Packets
+            port_alive_duration,    # Port alive Duration
+            stats['rx_bytes'],      # Delta Received Bytes
+            stats['tx_bytes'],      # Delta Sent Bytes
+            min(delta_alive_duration, 5),  # Delta Port alive Duration
+            min((port % 5) + 1, 5), # Connection Point
+            total_load,             # Total Load/Rate
+            len(stats['active_flows']), # Active Flow Entries
+            stats['packets_matched']    # Packets Matched
+        ]
+        
+        return features
+
+def process_pcap_for_ids(pcap_file):
+    """Process PCAP file and return features for IDS model"""
+    processor = PcapProcessor(pcap_file)
+    features = processor.process_packets()
+    return features
+
+def predict_from_features(features, model):
+    """Make prediction from extracted features"""
+    # Scale features
+    scaled_features = transform_new_input(features)
+    features_matrix = xgb.DMatrix(scaled_features.reshape(1, -1))
+    
+    # Make prediction and get probability distribution
+    raw_prediction = model.predict(features_matrix)
+    probabilities = raw_prediction[0]  # Get probability distribution
+    prediction = np.argmax(probabilities)
+    
+    # Add threshold for normal traffic
+    # If highest probability is for normal (class 0) and exceeds threshold
+    if prediction == 0 and probabilities[0] > 0.6:  # 60% confidence threshold
+        return get_prediction_message(0)
+    # If no class has high confidence, consider it normal
+    elif np.max(probabilities) < 0.4:  # Low confidence threshold
+        return get_prediction_message(0)
+    else:
+        return get_prediction_message(prediction)
+
+def get_prediction_message(prediction):
+    """Get formatted prediction message with confidence levels"""
+    messages = {
+        0: ("NORMAL TRAFFIC - No indication of attack.", 
+            "Traffic patterns appear to be within normal parameters."),
+        1: ("ALERT: Potential BLACKHOLE attack detected.", 
+            "Information: BLACKHOLE attacks occur when a router maliciously drops "
+            "packets it should forward. Investigate affected routes and traffic patterns."),
+        2: ("ALERT: Potential TCP-SYN flood attack detected.", 
+            "Information: TCP-SYN flood is a DDoS attack exhausting server resources "
+            "with half-open connections. Check connection states and implement SYN cookies."),
+        3: ("ALERT: PORTSCAN activity detected.", 
+            "Information: Port scanning detected - systematic probing of system ports. "
+            "Review firewall rules and implement connection rate limiting."),
+        4: ("ALERT: Potential DIVERSION attack detected.", 
+            "Information: Traffic diversion detected. Verify routing integrity and "
+            "check for signs of traffic manipulation or social engineering attempts.")
+    }
+    return messages.get(prediction, ("Unknown Traffic Pattern", "Additional analysis required."))
+
+def process_pcap_input(pcap_file):
+    """Process PCAP file input"""
+    try:
+        model = xgb.Booster()
+        model.load_model("m3_xg_boost.model")
+        features_list = process_pcap_for_ids(pcap_file.name)
+        if not features_list:
+            return "No valid network traffic found in PCAP file."
+        
+        results = []
+        for idx, features in enumerate(features_list):
+            result_msg, result_info = predict_from_features(features, model)
+            results.append(f"Traffic Pattern {idx + 1}:\n{result_msg}\n{result_info}\n")
+        
+        return "\n".join(results)
+    except Exception as e:
+        return f"Error processing PCAP file: {str(e)}"
+
+def process_manual_input(port_num, rx_packets, rx_bytes, tx_bytes, tx_packets, 
+                        port_duration, delta_rx_bytes, delta_tx_bytes, 
+                        delta_duration, conn_point, total_load, active_flows, 
+                        packets_matched):
+    #Process manual input values
+    try:
+        model = xgb.Booster()
+        model.load_model("m3_xg_boost.model")
+        features = [
+            port_num, rx_packets, rx_bytes, tx_bytes, tx_packets,
+            port_duration, delta_rx_bytes, delta_tx_bytes, delta_duration,
+            conn_point, total_load, active_flows, packets_matched
+        ]
+        
+        result_msg, result_info = predict_from_features(features, model)
+        return f"{result_msg}\n{result_info}"
+    except Exception as e:
+        return f"Error processing manual input: {str(e)}"
+
+# Main execution
+if __name__ == "__main__":
+    # Create the interface
+    with gr.Blocks(theme="default") as interface:
+        gr.Markdown("""
+        # Network Intrusion Detection System
+        Upload a PCAP file or use manual input to detect potential network attacks.
+        """)
+        
+        with gr.Tab("PCAP Analysis"):
+            pcap_input = gr.File(
+                label="Upload PCAP File", 
+                file_types=[".pcap", ".pcapng"]
+            )
+            pcap_output = gr.Textbox(label="Analysis Results")
+            pcap_button = gr.Button("Analyze PCAP")
+            pcap_button.click(
+                fn=process_pcap_input,
+                inputs=[pcap_input],
+                outputs=pcap_output
+            )
+        
+        with gr.Tab("Manual Input"):
+            # Manual input components
+            with gr.Row():
+                port_num = gr.Slider(1, 4, value=1, 
+                    label="Port Number - The switch port through which the flow passed")
+                rx_packets = gr.Slider(0, 352772, value=0, 
+                    label="Received Packets - Number of packets received by the port")
+            
+            with gr.Row():
+                rx_bytes = gr.Slider(0, 2.715916e08, value=0, 
+                    label="Received Bytes - Number of bytes received by the port")
+                tx_bytes = gr.Slider(0, 2.392430e08, value=0, 
+                    label="Sent Bytes - Number of bytes sent by the port")
+            
+            with gr.Row():
+                tx_packets = gr.Slider(0, 421598, value=0, 
+                    label="Sent Packets - Number of packets sent by the port")
+                port_duration = gr.Slider(0, 3317, value=0, 
+                    label="Port alive Duration (S) - The time port has been alive in seconds")
+            
+            with gr.Row():
+                delta_rx_bytes = gr.Slider(0, 6500000, value=0, 
+                    label="Delta Received Bytes")
+                delta_tx_bytes = gr.Slider(0, 6500000, value=0, 
+                    label="Delta Sent Bytes")
+            
+            with gr.Row():
+                delta_duration = gr.Slider(0, 5, value=0, 
+                    label="Delta Port alive Duration (S)")
+                conn_point = gr.Slider(1, 5, value=1, 
+                    label="Connection Point")
+            
+            with gr.Row():
+                total_load = gr.Slider(0, 1800000, value=0, 
+                    label="Total Load/Rate")
+                active_flows = gr.Slider(0, 610, value=0, 
+                    label="Active Flow Entries")
+            
+            with gr.Row():
+                packets_matched = gr.Slider(0, 1020000, value=0, 
+                    label="Packets Matched")
+            
+            manual_output = gr.Textbox(label="Analysis Results")
+            manual_button = gr.Button("Analyze Manual Input")
+            
+            # Connect manual input components
+            manual_button.click(
+                fn=process_manual_input,
+                inputs=[
+                    port_num, rx_packets, rx_bytes, tx_bytes, tx_packets,
+                    port_duration, delta_rx_bytes, delta_tx_bytes,
+                    delta_duration, conn_point, total_load, active_flows,
+                    packets_matched
+                ],
+                outputs=manual_output
+            )
+            
+            # Example inputs
+            gr.Examples(
+                examples=[
+                    [4, 350188, 14877116, 101354648, 159524, 2910, 278, 280, 
+                     5, 4, 0, 6, 667324],
+                    [2, 2326, 12856942, 31777516, 2998, 2497, 560, 560, 
+                     5, 2, 0, 4, 7259],
+                    [4, 150, 19774, 6475473, 3054, 166, 556, 6068, 
+                     5, 4, 502, 6, 7418],
+                    [2, 209, 20671, 6316631, 274, 96, 3527, 2757949, 
+                     5, 2, 183877, 8, 90494],
+                    [2, 1733, 37865130, 38063670, 3187, 2152, 0, 556, 
+                     5, 3, 0, 4, 14864]
+                ],
+                inputs=[
+                    port_num, rx_packets, rx_bytes, tx_bytes, tx_packets,
+                    port_duration, delta_rx_bytes, delta_tx_bytes,
+                    delta_duration, conn_point, total_load, active_flows,
+                    packets_matched
+                ]
+            )
+
+    # Launch the interface
+    interface.launch()