AI Network Traffic Analysis: Detecting Anomalies with ML

mkdir -p ai-traffic-analysis/{src,data,models,logs}
cd ai-traffic-analysis
python3 -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate
pip install --upgrade pip

pip install scapy==2.5.0 pandas==2.1.4 numpy==1.26.2 scikit-learn==1.3.2 tensorflow==2.15.0 matplotlib==3.8.2 seaborn==0.13.0 jupyter==1.0.0

cat > src/__init__.py << 'EOF'
# AI Network Traffic Analysis
EOF

cat > src/config.py << 'EOF'
"""Configuration for AI traffic analysis system."""
import os
from pathlib import Path

BASE_DIR = Path(__file__).parent.parent
DATA_DIR = BASE_DIR / "data"
MODELS_DIR = BASE_DIR / "models"
LOGS_DIR = BASE_DIR / "logs"

# Network interface to monitor
INTERFACE = os.getenv("NETWORK_INTERFACE", "eth0")  # Change to your interface

# Capture settings
CAPTURE_DURATION = 60  # seconds
MAX_PACKETS = 10000

# Feature extraction settings
FEATURE_WINDOW_SIZE = 100  # packets per window
FLOW_TIMEOUT = 300  # seconds

# ML model settings
MODEL_TYPE = "isolation_forest"  # isolation_forest, autoencoder, lstm
ANOMALY_THRESHOLD = 0.7  # threshold for anomaly detection
TRAINING_SAMPLES = 10000  # samples needed for training

# Alert settings
ALERT_COOLDOWN = 60  # seconds between alerts for same source
MIN_ANOMALY_SCORE = 0.8  # minimum score to trigger alert
EOF

# src/traffic_capture.py
"""Network traffic capture using Scapy."""
import logging
from collections import deque
from scapy.all import sniff, IP, TCP, UDP, ICMP
from typing import List, Dict, Optional
import time

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class TrafficCaptureError(Exception):
    """Custom error for traffic capture failures."""
    pass


class PacketCapture:
    """Captures and stores network packets."""
    
    def __init__(self, interface: str, max_packets: int = 10000):
        """
        Initialize packet capture.
        
        Args:
            interface: Network interface to capture on
            max_packets: Maximum packets to store in buffer
        """
        self.interface = interface
        self.max_packets = max_packets
        self.packets = deque(maxlen=max_packets)
        self.capturing = False
        self.stats = {
            "total_packets": 0,
            "ip_packets": 0,
            "tcp_packets": 0,
            "udp_packets": 0,
            "icmp_packets": 0,
            "errors": 0
        }
    
    def _packet_handler(self, packet) -> None:
        """Handle captured packet."""
        try:
            self.stats["total_packets"] += 1
            
            if packet.haslayer(IP):
                self.stats["ip_packets"] += 1
                
                if packet.haslayer(TCP):
                    self.stats["tcp_packets"] += 1
                elif packet.haslayer(UDP):
                    self.stats["udp_packets"] += 1
                elif packet.haslayer(ICMP):
                    self.stats["icmp_packets"] += 1
                
                # Store packet metadata
                packet_data = {
                    "timestamp": time.time(),
                    "src_ip": packet[IP].src,
                    "dst_ip": packet[IP].dst,
                    "protocol": packet[IP].proto,
                    "length": len(packet),
                    "ttl": packet[IP].ttl
                }
                
                # Add protocol-specific data
                if packet.haslayer(TCP):
                    packet_data["src_port"] = packet[TCP].sport
                    packet_data["dst_port"] = packet[TCP].dport
                    packet_data["flags"] = packet[TCP].flags
                elif packet.haslayer(UDP):
                    packet_data["src_port"] = packet[UDP].sport
                    packet_data["dst_port"] = packet[UDP].dport
                
                self.packets.append(packet_data)
                
        except Exception as e:
            self.stats["errors"] += 1
            logger.warning(f"Error processing packet: {e}")
    
    def start_capture(self, duration: int = 60) -> None:
        """
        Start capturing packets.
        
        Args:
            duration: Capture duration in seconds
        """
        if self.capturing:
            raise TrafficCaptureError("Capture already in progress")
        
        logger.info(f"Starting capture on {self.interface} for {duration}s")
        self.capturing = True
        self.packets.clear()
        self.stats = {k: 0 for k in self.stats}
        
        try:
            sniff(
                iface=self.interface,
                prn=self._packet_handler,
                timeout=duration,
                store=False
            )
        except Exception as e:
            logger.error(f"Capture error: {e}")
            raise TrafficCaptureError(f"Failed to capture: {e}")
        finally:
            self.capturing = False
            logger.info(f"Capture complete: {self.stats}")
    
    def get_packets(self) -> List[Dict]:
        """Get captured packets."""
        return list(self.packets)
    
    def get_stats(self) -> Dict:
        """Get capture statistics."""
        return self.stats.copy()

# src/feature_extraction.py
"""Feature extraction from network traffic."""
import numpy as np
import pandas as pd
from typing import List, Dict
from collections import defaultdict
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class FeatureExtractionError(Exception):
    """Custom error for feature extraction failures."""
    pass


class TrafficFeatureExtractor:
    """Extracts features from network traffic."""
    
    def __init__(self, window_size: int = 100):
        """
        Initialize feature extractor.
        
        Args:
            window_size: Number of packets per analysis window
        """
        self.window_size = window_size
    
    def extract_packet_features(self, packets: List[Dict]) -> pd.DataFrame:
        """
        Extract features from packet list.
        
        Args:
            packets: List of packet dictionaries
            
        Returns:
            DataFrame with extracted features
        """
        if not packets:
            raise FeatureExtractionError("No packets provided")
        
        try:
            df = pd.DataFrame(packets)
            
            # Basic packet features
            features = {
                "packet_count": len(packets),
                "total_bytes": df["length"].sum() if "length" in df.columns else 0,
                "avg_packet_size": df["length"].mean() if "length" in df.columns else 0,
                "std_packet_size": df["length"].std() if "length" in df.columns else 0,
            }
            
            # Protocol distribution
            if "protocol" in df.columns:
                protocol_counts = df["protocol"].value_counts()
                features["tcp_ratio"] = (df["protocol"] == 6).sum() / len(df)
                features["udp_ratio"] = (df["protocol"] == 17).sum() / len(df)
                features["icmp_ratio"] = (df["protocol"] == 1).sum() / len(df)
            
            # Port statistics
            if "dst_port" in df.columns:
                features["unique_dst_ports"] = df["dst_port"].nunique()
                features["port_entropy"] = self._calculate_entropy(df["dst_port"])
                features["top_port_ratio"] = df["dst_port"].value_counts().iloc[0] / len(df) if len(df) > 0 else 0
            
            # IP statistics
            if "src_ip" in df.columns:
                features["unique_src_ips"] = df["src_ip"].nunique()
                features["ip_entropy"] = self._calculate_entropy(df["src_ip"])
            
            # Timing features
            if "timestamp" in df.columns:
                timestamps = df["timestamp"].values
                if len(timestamps) > 1:
                    intervals = np.diff(np.sort(timestamps))
                    features["avg_interval"] = np.mean(intervals)
                    features["std_interval"] = np.std(intervals)
                    features["min_interval"] = np.min(intervals)
                    features["max_interval"] = np.max(intervals)
                else:
                    features["avg_interval"] = 0
                    features["std_interval"] = 0
                    features["min_interval"] = 0
                    features["max_interval"] = 0
            
            return pd.DataFrame([features])
            
        except Exception as e:
            logger.error(f"Feature extraction error: {e}")
            raise FeatureExtractionError(f"Failed to extract features: {e}")
    
    def extract_flow_features(self, packets: List[Dict]) -> pd.DataFrame:
        """
        Extract flow-based features.
        
        Args:
            packets: List of packet dictionaries
            
        Returns:
            DataFrame with flow features
        """
        if not packets:
            raise FeatureExtractionError("No packets provided")
        
        try:
            df = pd.DataFrame(packets)
            
            # Group by flow (src_ip, dst_ip, dst_port)
            flows = defaultdict(lambda: {"packets": 0, "bytes": 0, "duration": 0})
            
            for _, row in df.iterrows():
                if "src_ip" in row and "dst_ip" in row and "dst_port" in row:
                    flow_key = (row["src_ip"], row["dst_ip"], row["dst_port"])
                    flows[flow_key]["packets"] += 1
                    flows[flow_key]["bytes"] += row.get("length", 0)
                    
                    if "timestamp" in row:
                        if flows[flow_key]["duration"] == 0:
                            flows[flow_key]["start_time"] = row["timestamp"]
                        flows[flow_key]["end_time"] = row["timestamp"]
            
            # Calculate flow features
            flow_features = []
            for flow_key, flow_data in flows.items():
                duration = flow_data.get("end_time", 0) - flow_data.get("start_time", 0)
                flow_features.append({
                    "src_ip": flow_key[0],
                    "dst_ip": flow_key[1],
                    "dst_port": flow_key[2],
                    "packet_count": flow_data["packets"],
                    "byte_count": flow_data["bytes"],
                    "duration": duration,
                    "bytes_per_packet": flow_data["bytes"] / flow_data["packets"] if flow_data["packets"] > 0 else 0,
                    "packets_per_second": flow_data["packets"] / duration if duration > 0 else 0
                })
            
            return pd.DataFrame(flow_features)
            
        except Exception as e:
            logger.error(f"Flow feature extraction error: {e}")
            raise FeatureExtractionError(f"Failed to extract flow features: {e}")
    
    def _calculate_entropy(self, series: pd.Series) -> float:
        """Calculate Shannon entropy."""
        value_counts = series.value_counts()
        probabilities = value_counts / len(series)
        entropy = -np.sum(probabilities * np.log2(probabilities + 1e-10))
        return entropy
    
    def extract_window_features(self, packets: List[Dict]) -> pd.DataFrame:
        """
        Extract features for sliding windows.
        
        Args:
            packets: List of packet dictionaries
            
        Returns:
            DataFrame with window features
        """
        if len(packets) < self.window_size:
            return self.extract_packet_features(packets)
        
        windows = []
        for i in range(0, len(packets), self.window_size):
            window_packets = packets[i:i + self.window_size]
            window_features = self.extract_packet_features(window_packets)
            windows.append(window_features)
        
        return pd.concat(windows, ignore_index=True)

# src/ml_models.py
"""ML models for anomaly detection."""
import numpy as np
import pandas as pd
from sklearn.ensemble import IsolationForest
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import pickle
import logging
from pathlib import Path
from typing import Tuple, Optional

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class MLModelError(Exception):
    """Custom error for ML model failures."""
    pass


class AnomalyDetector:
    """ML-based anomaly detector for network traffic."""
    
    def __init__(self, model_type: str = "isolation_forest"):
        """
        Initialize anomaly detector.
        
        Args:
            model_type: Type of model (isolation_forest, autoencoder)
        """
        self.model_type = model_type
        self.model = None
        self.scaler = StandardScaler()
        self.feature_names = None
        self.is_trained = False
    
    def train(self, features: pd.DataFrame, contamination: float = 0.1) -> None:
        """
        Train the anomaly detection model.
        
        Args:
            features: DataFrame with extracted features
            contamination: Expected proportion of anomalies
        """
        if features.empty:
            raise MLModelError("No features provided for training")
        
        try:
            # Store feature names
            self.feature_names = features.columns.tolist()
            
            # Scale features
            X_scaled = self.scaler.fit_transform(features)
            
            # Train model
            if self.model_type == "isolation_forest":
                self.model = IsolationForest(
                    contamination=contamination,
                    random_state=42,
                    n_estimators=100
                )
                self.model.fit(X_scaled)
            else:
                raise MLModelError(f"Unsupported model type: {self.model_type}")
            
            self.is_trained = True
            logger.info(f"Model trained on {len(features)} samples")
            
        except Exception as e:
            logger.error(f"Training error: {e}")
            raise MLModelError(f"Failed to train model: {e}")
    
    def predict(self, features: pd.DataFrame) -> Tuple[np.ndarray, np.ndarray]:
        """
        Predict anomalies in features.
        
        Args:
            features: DataFrame with features to predict
            
        Returns:
            Tuple of (predictions, scores)
            - predictions: 1 for normal, -1 for anomaly
            - scores: Anomaly scores (lower = more anomalous)
        """
        if not self.is_trained:
            raise MLModelError("Model not trained")
        
        if features.empty:
            raise MLModelError("No features provided")
        
        try:
            # Ensure same features as training
            if set(features.columns) != set(self.feature_names):
                raise MLModelError("Feature mismatch with training data")
            
            # Reorder columns to match training
            features = features[self.feature_names]
            
            # Scale features
            X_scaled = self.scaler.transform(features)
            
            # Predict
            predictions = self.model.predict(X_scaled)
            scores = self.model.score_samples(X_scaled)
            
            # Convert scores to anomaly scores (lower = more anomalous)
            anomaly_scores = -scores  # Invert so higher = more anomalous
            
            return predictions, anomaly_scores
            
        except Exception as e:
            logger.error(f"Prediction error: {e}")
            raise MLModelError(f"Failed to predict: {e}")
    
    def save(self, filepath: Path) -> None:
        """Save model to file."""
        if not self.is_trained:
            raise MLModelError("Model not trained")
        
        try:
            model_data = {
                "model": self.model,
                "scaler": self.scaler,
                "feature_names": self.feature_names,
                "model_type": self.model_type
            }
            
            with open(filepath, "wb") as f:
                pickle.dump(model_data, f)
            
            logger.info(f"Model saved to {filepath}")
            
        except Exception as e:
            logger.error(f"Save error: {e}")
            raise MLModelError(f"Failed to save model: {e}")
    
    def load(self, filepath: Path) -> None:
        """Load model from file."""
        try:
            with open(filepath, "rb") as f:
                model_data = pickle.load(f)
            
            self.model = model_data["model"]
            self.scaler = model_data["scaler"]
            self.feature_names = model_data["feature_names"]
            self.model_type = model_data["model_type"]
            self.is_trained = True
            
            logger.info(f"Model loaded from {filepath}")
            
        except Exception as e:
            logger.error(f"Load error: {e}")
            raise MLModelError(f"Failed to load model: {e}")

# src/detection_pipeline.py
"""Real-time anomaly detection pipeline."""
import time
import logging
from typing import Dict, List, Optional
from collections import defaultdict
from src.traffic_capture import PacketCapture
from src.feature_extraction import TrafficFeatureExtractor
from src.ml_models import AnomalyDetector
from src.config import ALERT_COOLDOWN, MIN_ANOMALY_SCORE

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class DetectionPipelineError(Exception):
    """Custom error for detection pipeline failures."""
    pass


class RealTimeDetector:
    """Real-time anomaly detection pipeline."""
    
    def __init__(
        self,
        detector: AnomalyDetector,
        extractor: TrafficFeatureExtractor,
        interface: str = "eth0"
    ):
        """
        Initialize real-time detector.
        
        Args:
            detector: Trained anomaly detector
            extractor: Feature extractor
            interface: Network interface to monitor
        """
        self.detector = detector
        self.extractor = extractor
        self.interface = interface
        self.capture = PacketCapture(interface)
        self.alert_cooldown = defaultdict(float)
        self.alerts = []
    
    def process_window(self, window_packets: List[Dict]) -> Optional[Dict]:
        """
        Process a window of packets and detect anomalies.
        
        Args:
            window_packets: List of packets in current window
            
        Returns:
            Alert dictionary if anomaly detected, None otherwise
        """
        if len(window_packets) < 10:  # Need minimum packets
            return None
        
        try:
            # Extract features
            features = self.extractor.extract_packet_features(window_packets)
            
            # Predict anomalies
            predictions, scores = self.detector.predict(features)
            
            # Check for anomalies
            max_score = np.max(scores)
            if max_score >= MIN_ANOMALY_SCORE:
                # Get source IPs from window
                src_ips = [p.get("src_ip") for p in window_packets if "src_ip" in p]
                top_src_ip = max(set(src_ips), key=src_ips.count) if src_ips else "unknown"
                
                # Check cooldown
                current_time = time.time()
                if current_time - self.alert_cooldown.get(top_src_ip, 0) < ALERT_COOLDOWN:
                    return None
                
                # Generate alert
                alert = {
                    "timestamp": current_time,
                    "src_ip": top_src_ip,
                    "anomaly_score": float(max_score),
                    "packet_count": len(window_packets),
                    "features": features.iloc[0].to_dict()
                }
                
                self.alert_cooldown[top_src_ip] = current_time
                self.alerts.append(alert)
                
                logger.warning(f"ANOMALY DETECTED: {top_src_ip} (score: {max_score:.2f})")
                return alert
            
            return None
            
        except Exception as e:
            logger.error(f"Processing error: {e}")
            return None
    
    def run_continuous(self, window_size: int = 100, duration: int = 300) -> List[Dict]:
        """
        Run continuous detection.
        
        Args:
            window_size: Packets per analysis window
            duration: Total monitoring duration in seconds
            
        Returns:
            List of alerts
        """
        logger.info(f"Starting continuous detection for {duration}s")
        self.alerts.clear()
        
        start_time = time.time()
        window_packets = []
        
        def packet_handler(packet):
            nonlocal window_packets
            window_packets.append(packet)
            
            if len(window_packets) >= window_size:
                self.process_window(window_packets)
                window_packets = []
        
        try:
            self.capture.start_capture(duration=duration)
            packets = self.capture.get_packets()
            
            # Process in windows
            for i in range(0, len(packets), window_size):
                window = packets[i:i + window_size]
                self.process_window(window)
            
            logger.info(f"Detection complete: {len(self.alerts)} alerts")
            return self.alerts
            
        except Exception as e:
            logger.error(f"Continuous detection error: {e}")
            raise DetectionPipelineError(f"Failed to run detection: {e}")