Network-Aware Hybrid GNN for Attack Prediction

Advanced Graph Neural Network with network-aware edges for predicting next attack techniques in ICS environments. Features federated learning support and comprehensive visualization.

🌟 Key Features

Network-Aware Graph Edges: 6D edge features capturing network traffic patterns
Hybrid Architecture: Combines graph structure + network features
Federated Learning: Horizontal FL for privacy-preserving collaborative training
Comprehensive Visualization: 7+ evaluation figures and analysis reports
High Performance: ~60% Top-1, ~75% Top-3, ~85% Top-5 accuracy

🚀 Quick Start

1. Installation

# Install dependencies
pip install -r requirements.txt

# Or using Poetry
poetry install
poetry shell

# Verify installation
python test_network_aware.py

2. Compare Approaches

# See difference between original and network-aware
python compare_approaches.py

3. Train Network-Aware Model

# Train with network-aware graph edges
python scripts/analyze_and_train_network_aware.py

# Or use notebook
jupyter notebook notebooks/analyze_and_train_network_aware.ipynb

4. View Results

# Check generated figures
ls figures/

# View training history
python -c "import pickle; h=pickle.load(open('models/training_history.pkl','rb')); print(h)"

📁 Project Structure

gnn_clean/
├── data/                                    # Training data
│   ├── sequences_with_network_features.json # 50K+ training pairs
│   ├── all_attack_sequences.json           # Attack chains
│   └── combined_statistics.json            # Data statistics
│
├── scripts/                                 # Training scripts
│   ├── analyze_and_train_network_aware.py  # Main training (network-aware)
│   └── generate_realistic_network_data.py  # Data generation
│
├── notebooks/                               # Jupyter notebooks
│   └── analyze_and_train_network_aware.ipynb
│
├── models/                                  # Saved models
│   ├── best_network_aware_model.pt         # Trained model weights
│   ├── network_aware_mappings.pkl          # Technique mappings
│   └── training_history.pkl                # Training curves
│
├── figures/                                 # Generated visualizations
│   ├── network_aware_training.png          # Training curves
│   ├── edge_features_distribution.png      # Edge analysis
│   ├── network_variance_analysis.png       # Variance patterns
│   ├── network_aware_confusion_matrix.png  # Confusion matrix
│   ├── network_aware_confidence.png        # Confidence analysis
│   ├── data_distribution.png               # Data characteristics
│   ├── graph_structure.png                 # Graph topology
│   └── high_variance_transitions.txt       # Analysis report
│
├── docs/                                    # Documentation
│   ├── GNN_MODEL_ARCHITECTURE.md           # Model architecture (Mermaid)
│   ├── NETWORK_AWARE_IMPLEMENTATION.md     # Implementation guide
│   ├── INTEGRATING_CURRENT_TRAFFIC.md      # Traffic integration guide
│   ├── NETWORK_INTEGRATION_OPTIONS.md      # Integration approaches
│   ├── FL_DESIGN_PLAN.md                   # Federated learning plan
│   └── VFL_DESIGN_PLAN.md                  # Vertical FL (archived)
│
├── compare_approaches.py                    # Compare original vs network-aware
├── test_network_aware.py                   # Validation script
├── requirements.txt                         # Dependencies
├── pyproject.toml                          # Poetry config
├── INSTALLATION.md                         # Installation guide
├── VISUALIZATION_GUIDE.md                  # Figure descriptions
├── NETWORK_AWARE_README.md                 # Detailed user guide
├── QUICK_REFERENCE.md                      # Quick reference card
└── README.md                               # This file

📊 Model Architecture

Network-Aware Hybrid GNN

Input:
├── Current Attack ID (e.g., T1566 - Phishing)
├── Network Features (18D: packets, bytes, protocols, etc.)
└── Graph Structure (edge_index + 6D edge features)

Model:
├── GNN Path: Node Embedding → 3 GAT Layers → Node Representation
│   └── GAT uses 6D edge features (network-aware attention)
├── Network Path: Network Features → 2-Layer MLP → Network Representation
└── Fusion: Concatenate → MLP → Classifier → Predictions

Edge Features (6D):
├── [0] Transition probability
├── [1] Average packet count
├── [2] Std packet count
├── [3] Average byte count
├── [4] Std byte count
└── [5] Log transition count

Key Innovation: Edge features include network statistics, enabling the model to distinguish attack variants by traffic patterns.

See docs/GNN_MODEL_ARCHITECTURE.md for detailed architecture with Mermaid diagrams.

📈 Performance

Network-Aware Model (Current)

Training Data: 50,000+ pairs with network features
Top-1 Accuracy: ~60%
Top-3 Accuracy: ~75%
Top-5 Accuracy: ~85%
Training Time: ~20 minutes (100 epochs, GPU)
Parameters: ~1.75M

Comparison

Model	Top-1	Top-3	Top-5	Edge Features
Original GNN	55%	70%	80%	1D (frequency)
Network-Aware	60%	75%	85%	6D (frequency + network)
Improvement	+5%	+5%	+5%	Better variant detection

💡 Key Concepts

Network-Aware Edges

Traditional GNN: Edges only have transition frequency

edge_weight = 0.45  # 45% probability

Network-Aware GNN: Edges have 6D features

edge_features = [
    0.45,      # Transition probability
    250.0,     # Usually 250 packets
    80.0,      # ±80 packet variance
    75000.0,   # Usually 75KB
    25000.0,   # ±25KB variance
    6.11       # log(450 occurrences)
]

Benefit: Model can distinguish "Phishing → Command Execution with 250 packets" vs "Phishing → Command Execution with 50 packets" (different attack variants).

Dual-Path Architecture

GNN Path: Learns attack sequence patterns from graph structure
Network Path: Learns network behavior from traffic features
Fusion: Combines both for robust predictions

See docs/INTEGRATING_CURRENT_TRAFFIC.md for integration options.

📖 Documentation

Getting Started

INSTALLATION.md - Installation guide
QUICK_REFERENCE.md - Quick reference card
NETWORK_AWARE_README.md - Detailed user guide

Architecture & Design

docs/GNN_MODEL_ARCHITECTURE.md - Model architecture (Mermaid diagrams)
docs/NETWORK_AWARE_IMPLEMENTATION.md - Implementation details
docs/INTEGRATING_CURRENT_TRAFFIC.md - Traffic integration guide
docs/NETWORK_INTEGRATION_OPTIONS.md - Integration approaches

Federated Learning

docs/FL_DESIGN_PLAN.md - Horizontal FL design (3 facilities)
docs/VFL_DESIGN_PLAN.md - Vertical FL design (archived)

Visualization

VISUALIZATION_GUIDE.md - Guide to generated figures

🔬 Usage Examples

Basic Prediction

import torch
import pickle
from scripts.analyze_and_train_network_aware import NetworkAwareHybridGNN

# Load model and mappings
model = NetworkAwareHybridGNN(num_techniques, num_network_features, 6, 256, 8, 0.4)
model.load_state_dict(torch.load('models/best_network_aware_model.pt'))
model.eval()

with open('models/network_aware_mappings.pkl', 'rb') as f:
    mappings = pickle.load(f)

# Prepare input
current_attack = "T1566"  # Phishing
network_features = [150.0, 50000.0, ...]  # 18D vector

current_idx = torch.tensor([mappings['technique_to_idx'][current_attack]])
network_tensor = torch.tensor([network_features])

# Predict
with torch.no_grad():
    logits = model(current_idx, network_tensor, edge_index, edge_attr)
    probs = torch.softmax(logits, dim=1)
    top5_probs, top5_indices = torch.topk(probs, 5)

# Display results
print("Top 5 predictions:")
for prob, idx in zip(top5_probs[0], top5_indices[0]):
    technique = mappings['idx_to_technique'][idx.item()]
    print(f"  {technique}: {prob.item():.3f}")

Output:

Top 5 predictions:
  T1059: 0.350  (Command Execution)
  T1071: 0.250  (Application Layer Protocol)
  T1003: 0.150  (Credential Dumping)
  T1082: 0.100  (System Information Discovery)
  T1083: 0.080  (File and Directory Discovery)

🎯 Training & Evaluation

Training Configuration

# Default hyperparameters
HIDDEN_DIM = 256
NUM_HEADS = 8
DROPOUT = 0.4
LEARNING_RATE = 0.0005
BATCH_SIZE = 128
NUM_EPOCHS = 100

Hyperparameter Tuning

# Modify in script or pass as arguments
HIDDEN_DIM = 512      # Increase model capacity
NUM_HEADS = 16        # More attention heads
DROPOUT = 0.5         # More regularization

Generated Outputs

After training, you'll get:

Models:

models/best_network_aware_model.pt - Trained weights
models/network_aware_mappings.pkl - Technique mappings
models/training_history.pkl - Training curves

Figures (7 visualizations):

figures/network_aware_training.png - Training/validation curves
figures/edge_features_distribution.png - Edge feature analysis
figures/network_variance_analysis.png - High-variance transitions
figures/network_aware_confusion_matrix.png - Confusion matrix
figures/network_aware_confidence.png - Confidence & calibration
figures/data_distribution.png - Data characteristics
figures/graph_structure.png - Graph topology

Reports:

figures/high_variance_transitions.txt - Detailed analysis

Evaluation Metrics

# Typical results on test set
Top-1 Accuracy: 60%  # Exact match
Top-3 Accuracy: 75%  # True attack in top 3
Top-5 Accuracy: 85%  # True attack in top 5

# Per-class accuracy for top 10 classes
# High-variance transitions show biggest improvement

🔮 Federated Learning (Planned)

Horizontal Federated Learning

Train collaboratively across 3 facilities without sharing raw data:

Facility 1 (10K samples) ──┐
                           ├──> Central Server (Aggregates)
Facility 2 (15K samples) ──┤         ↓
                           │    Global Model
Facility 3 (12K samples) ──┘         ↓
                                Broadcast back

Features:

Privacy-preserving (no raw data sharing)
FedAvg aggregation
Supports IID and Non-IID data splits
Expected: 95-98% of centralized accuracy

See docs/FL_DESIGN_PLAN.md for implementation plan.

🧪 Testing & Validation

Run Tests

# Validate setup
python test_network_aware.py

# Compare approaches
python compare_approaches.py

# Check data quality
python -c "
import json
data = json.load(open('data/sequences_with_network_features.json'))
print(f'Loaded {len(data)} training pairs')
print(f'Network features: {len(data[0][\"network_features\"])}D')
"

Verify Model

# Check model parameters
import torch
model = torch.load('models/best_network_aware_model.pt')
total_params = sum(p.numel() for p in model.values())
print(f"Total parameters: {total_params:,}")

🛠️ Advanced Features

Traffic Integration Options

Three approaches to integrate current traffic with graph:

Dual-Path (Current): Separate processing, late fusion
Node Integration (Recommended): Add traffic to node features
Attention Modulation (Advanced): Traffic modulates attention weights

See docs/INTEGRATING_CURRENT_TRAFFIC.md for implementation.

High-Variance Transitions

Transitions with high network variance benefit most from network-aware edges:

# View high-variance transitions
cat figures/high_variance_transitions.txt

# These show biggest accuracy improvements
# Example: T1566→T1059 with CV=1.2 (high variance)

🐛 Troubleshooting

Common Issues

Issue: CUDA out of memory

# Solution: Reduce batch size or use CPU
BATCH_SIZE = 64  # or 32
device = torch.device('cpu')

Issue: Import errors

# Solution: Reinstall dependencies
pip install -r requirements.txt --force-reinstall

Issue: No visualization

# Solution: Install matplotlib
pip install matplotlib seaborn scikit-learn

See INSTALLATION.md for detailed troubleshooting.

📚 Additional Resources

Papers

Graph Attention Networks (Veličković et al., 2018)
Federated Learning (McMahan et al., 2017)
MITRE ATT&CK for ICS

Tools

PyTorch: https://pytorch.org/
PyTorch Geometric: https://pytorch-geometric.readthedocs.io/
MITRE ATT&CK: https://attack.mitre.org/matrices/ics/

🤝 Contributing

Contributions welcome! Areas for improvement:

Model Enhancements
- Implement fully integrated traffic model
- Add temporal features
- Experiment with different GNN architectures
Federated Learning
- Implement horizontal FL
- Add differential privacy
- Optimize communication efficiency
Visualization
- Interactive dashboards
- Attention weight visualization
- Real-time prediction monitoring
Documentation
- More examples
- Tutorial notebooks
- API documentation

📄 License

MIT License - see LICENSE file for details

📧 Contact

For questions or issues:

Open an issue on GitHub
Check documentation in docs/
Review troubleshooting in INSTALLATION.md

🎓 Citation

If you use this code in your research, please cite:

@software{network_aware_gnn_attack_prediction,
  title={Network-Aware Hybrid GNN for Attack Technique Prediction},
  author={Your Name},
  year={2025},
  url={https://github.com/yourusername/gnn-attack-prediction},
  note={Features network-aware graph edges and federated learning support}
}

🌟 Acknowledgments

MITRE ATT&CK for ICS framework
PyTorch Geometric library
Graph Attention Networks paper

Version: 0.2.0
Last Updated: 2025
Status: Active Development

FilesExpand file tree

README.md

Latest commit

History