Add training script

train_model.py

@@ -0,0 +1,399 @@
+#!/usr/bin/env python3
+"""
+C2Sentinel Training Script v2 - Improved training with proper normalization
+"""
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+import random
+from tqdm import tqdm
+import json
+
+from c2sentinel import (
+    C2Sentinel, C2SentinelConfig, LogBERTC2Sentinel,
+    FeatureExtractor
+)
+from safetensors.torch import save_file
+
+
+class C2TrafficDataset(Dataset):
+    """Dataset with normalized features."""
+
+    def __init__(self, num_samples=10000, normalize=True):
+        self.samples = []
+        self.labels = []
+        self.c2_types = []
+        self.feature_extractor = FeatureExtractor()
+
+        print(f"Generating {num_samples} training samples...")
+
+        num_c2 = num_samples // 2
+        num_benign = num_samples - num_c2
+
+        # Generate C2 samples
+        for _ in tqdm(range(num_c2), desc="C2 samples"):
+            connections, c2_type = self._generate_c2_traffic()
+            features = self.feature_extractor.extract_features(connections)
+            self.samples.append(features)
+            self.labels.append(1)
+            self.c2_types.append(c2_type)
+
+        # Generate benign samples
+        for _ in tqdm(range(num_benign), desc="Benign samples"):
+            connections = self._generate_benign_traffic()
+            features = self.feature_extractor.extract_features(connections)
+            self.samples.append(features)
+            self.labels.append(0)
+            self.c2_types.append(0)
+
+        self.samples = np.array(self.samples, dtype=np.float32)
+        self.labels = np.array(self.labels, dtype=np.float32)
+        self.c2_types = np.array(self.c2_types, dtype=np.int64)
+
+        # Normalize features (critical for training stability)
+        if normalize:
+            self.mean = np.mean(self.samples, axis=0)
+            self.std = np.std(self.samples, axis=0) + 1e-8
+            self.samples = (self.samples - self.mean) / self.std
+
+            # Save normalization params
+            np.savez('normalization_params.npz', mean=self.mean, std=self.std)
+            print(f"Feature stats - mean range: [{self.mean.min():.2f}, {self.mean.max():.2f}], "
+                  f"std range: [{self.std.min():.4f}, {self.std.max():.2f}]")
+
+        # Shuffle
+        indices = np.random.permutation(len(self.samples))
+        self.samples = self.samples[indices]
+        self.labels = self.labels[indices]
+        self.c2_types = self.c2_types[indices]
+
+        print(f"C2 samples: {np.sum(self.labels)}, Benign: {len(self.labels) - np.sum(self.labels)}")
+
+    def _generate_c2_traffic(self):
+        """Generate C2 beacon traffic with clear patterns."""
+        c2_type = random.randint(1, 10)
+
+        # Strong C2 characteristics
+        if c2_type <= 3:  # Fast beacon (Metasploit-style)
+            interval = random.uniform(2, 15)
+            jitter = random.uniform(0, 0.15)  # Low jitter
+            port = random.choice([4444, 4445, 5555, 443])
+            bytes_sent = random.randint(80, 200)
+            bytes_recv = random.randint(40, 150)
+        elif c2_type <= 6:  # Medium beacon (Cobalt Strike-style)
+            interval = random.uniform(30, 90)
+            jitter = random.uniform(0, 0.2)
+            port = 443
+            bytes_sent = random.randint(60, 150)
+            bytes_recv = random.randint(40, 100)
+        else:  # Slow beacon (APT-style)
+            interval = random.uniform(120, 300)
+            jitter = random.uniform(0, 0.1)  # Very low jitter for APT
+            port = 443
+            bytes_sent = random.randint(50, 120)
+            bytes_recv = random.randint(40, 80)
+
+        # Single destination (key C2 indicator)
+        dst_ip = f"{random.randint(1,223)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}"
+        num_connections = random.randint(10, 40)
+
+        connections = []
+        timestamp = 1705600000
+
+        for _ in range(num_connections):
+            actual_interval = interval * (1 + random.uniform(-jitter, jitter))
+            timestamp += actual_interval
+
+            # Very consistent sizes (key C2 indicator)
+            size_var = random.uniform(0.95, 1.05)
+
+            connections.append({
+                'timestamp': timestamp,
+                'dst_ip': dst_ip,
+                'dst_port': port,
+                'bytes_sent': int(bytes_sent * size_var),
+                'bytes_recv': int(bytes_recv * size_var),
+                'protocol': 'tcp'
+            })
+
+        return connections, c2_type
+
+    def _generate_benign_traffic(self):
+        """Generate clearly benign traffic."""
+        pattern = random.choice(['browsing', 'api', 'streaming', 'interactive'])
+
+        connections = []
+        timestamp = 1705600000
+
+        if pattern == 'browsing':
+            # Multiple destinations, highly variable sizes
+            for _ in range(random.randint(10, 40)):
+                timestamp += random.uniform(0.5, 45)
+                connections.append({
+                    'timestamp': timestamp,
+                    'dst_ip': f"{random.randint(1,223)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}",
+                    'dst_port': random.choice([80, 443]),
+                    'bytes_sent': random.randint(200, 5000),
+                    'bytes_recv': random.randint(5000, 500000),
+                    'protocol': 'tcp'
+                })
+
+        elif pattern == 'api':
+            # Single dest but HIGHLY variable response sizes
+            dst_ip = f"{random.randint(1,223)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}"
+            for _ in range(random.randint(15, 40)):
+                timestamp += random.uniform(0.1, 20)
+                connections.append({
+                    'timestamp': timestamp,
+                    'dst_ip': dst_ip,
+                    'dst_port': 443,
+                    'bytes_sent': random.randint(100, 3000),
+                    'bytes_recv': random.randint(200, 100000),  # Highly variable
+                    'protocol': 'tcp'
+                })
+
+        elif pattern == 'streaming':
+            # Large downloads, irregular timing
+            dst_ip = f"{random.randint(1,223)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}"
+            for _ in range(random.randint(20, 60)):
+                timestamp += random.uniform(0.05, 3)
+                connections.append({
+                    'timestamp': timestamp,
+                    'dst_ip': dst_ip,
+                    'dst_port': 443,
+                    'bytes_sent': random.randint(30, 200),
+                    'bytes_recv': random.randint(5000, 150000),
+                    'protocol': 'tcp'
+                })
+
+        else:  # interactive (ssh-like)
+            dst_ip = f"192.168.{random.randint(0,255)}.{random.randint(1,254)}"
+            for _ in range(random.randint(15, 50)):
+                if random.random() < 0.3:
+                    timestamp += random.uniform(3, 45)  # Thinking
+                else:
+                    timestamp += random.uniform(0.1, 2)  # Typing
+                connections.append({
+                    'timestamp': timestamp,
+                    'dst_ip': dst_ip,
+                    'dst_port': 22,
+                    'bytes_sent': random.randint(20, 800),
+                    'bytes_recv': random.randint(50, 20000),
+                    'protocol': 'tcp'
+                })
+
+        return connections
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        return {
+            'features': torch.tensor(self.samples[idx]),
+            'label': torch.tensor(self.labels[idx]),
+            'c2_type': torch.tensor(self.c2_types[idx])
+        }
+
+
+def train_model(num_epochs=100, batch_size=32, learning_rate=0.0001, num_samples=20000):
+    """Train with improved stability."""
+
+    print("=" * 70)
+    print("C2Sentinel Model Training v2")
+    print("=" * 70)
+
+    config = C2SentinelConfig()
+    model = LogBERTC2Sentinel(config)
+
+    # Initialize weights properly
+    def init_weights(m):
+        if isinstance(m, nn.Linear):
+            nn.init.xavier_uniform_(m.weight, gain=0.5)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+    model.apply(init_weights)
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Device: {device}")
+    model.to(device)
+
+    # Count parameters
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Model parameters: {total_params:,}")
+
+    dataset = C2TrafficDataset(num_samples=num_samples, normalize=True)
+
+    train_size = int(0.9 * len(dataset))
+    val_size = len(dataset) - train_size
+    train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])
+
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)
+    val_loader = DataLoader(val_dataset, batch_size=batch_size)
+
+    print(f"Train: {train_size}, Val: {val_size}")
+
+    # Simple BCE loss - focus on main task only
+    criterion = nn.BCEWithLogitsLoss()
+
+    # Lower LR with warmup
+    optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.001)
+
+    # Warmup + cosine decay
+    warmup_epochs = 5
+    def lr_lambda(epoch):
+        if epoch < warmup_epochs:
+            return (epoch + 1) / warmup_epochs
+        return 0.5 * (1 + np.cos(np.pi * (epoch - warmup_epochs) / (num_epochs - warmup_epochs)))
+
+    scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+    best_val_acc = 0
+    patience = 15
+    patience_counter = 0
+
+    for epoch in range(num_epochs):
+        model.train()
+        train_loss = 0
+        train_correct = 0
+        train_total = 0
+
+        for batch in tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}", leave=False):
+            features = batch['features'].to(device)
+            labels = batch['label'].to(device)
+
+            optimizer.zero_grad()
+            outputs = model(features)
+
+            # Only C2 detection loss
+            loss = criterion(outputs['c2_logits'].squeeze(), labels)
+
+            loss.backward()
+
+            # Gradient clipping
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)
+
+            optimizer.step()
+
+            train_loss += loss.item()
+            predictions = (torch.sigmoid(outputs['c2_logits'].squeeze()) > 0.5).float()
+            train_correct += (predictions == labels).sum().item()
+            train_total += labels.size(0)
+
+        scheduler.step()
+
+        # Validation
+        model.eval()
+        val_correct = 0
+        val_total = 0
+        val_loss = 0
+
+        with torch.no_grad():
+            for batch in val_loader:
+                features = batch['features'].to(device)
+                labels = batch['label'].to(device)
+
+                outputs = model(features)
+                loss = criterion(outputs['c2_logits'].squeeze(), labels)
+                val_loss += loss.item()
+
+                predictions = (torch.sigmoid(outputs['c2_logits'].squeeze()) > 0.5).float()
+                val_correct += (predictions == labels).sum().item()
+                val_total += labels.size(0)
+
+        train_acc = 100 * train_correct / train_total
+        val_acc = 100 * val_correct / val_total
+        lr = optimizer.param_groups[0]['lr']
+
+        print(f"Epoch {epoch+1}: Loss={train_loss/len(train_loader):.4f}, "
+              f"Train={train_acc:.1f}%, Val={val_acc:.1f}%, LR={lr:.6f}")
+
+        if val_acc > best_val_acc:
+            best_val_acc = val_acc
+            patience_counter = 0
+            save_file(model.state_dict(), 'c2_sentinel.safetensors')
+            print(f"  -> Saved (Val: {val_acc:.1f}%)")
+        else:
+            patience_counter += 1
+            if patience_counter >= patience:
+                print(f"Early stopping at epoch {epoch+1}")
+                break
+
+    print(f"\nBest validation accuracy: {best_val_acc:.1f}%")
+    return model, config
+
+
+def test_model():
+    """Test the trained model."""
+    print("\n" + "=" * 70)
+    print("Testing Model")
+    print("=" * 70)
+
+    sentinel = C2Sentinel.load('c2_sentinel')
+
+    # Test 1: Cobalt Strike
+    print("\n[1] Cobalt Strike Beacon (60s interval)...")
+    cs = [{'timestamp': 1705600000 + i*60, 'dst_ip': '185.234.72.19', 'dst_port': 443,
+           'bytes_sent': 92, 'bytes_recv': 48} for i in range(16)]
+    r = sentinel.analyze(cs)
+    print(f"  {'✓ C2 DETECTED' if r.is_c2 else '✗ No C2'} (prob={r.c2_probability:.2%})")
+
+    # Test 2: Metasploit
+    print("\n[2] Metasploit Beacon (5s interval, port 4444)...")
+    msf = [{'timestamp': 1705600000 + i*5, 'dst_ip': '10.10.10.10', 'dst_port': 4444,
+            'bytes_sent': 150, 'bytes_recv': 400} for i in range(20)]
+    r = sentinel.analyze(msf)
+    print(f"  {'✓ C2 DETECTED' if r.is_c2 else '✗ No C2'} (prob={r.c2_probability:.2%})")
+
+    # Test 3: Slow APT beacon
+    print("\n[3] APT Slow Beacon (120s interval)...")
+    apt = [{'timestamp': 1705600000 + i*120, 'dst_ip': '45.33.32.156', 'dst_port': 443,
+            'bytes_sent': 80, 'bytes_recv': 60} for i in range(12)]
+    r = sentinel.analyze(apt)
+    print(f"  {'✓ C2 DETECTED' if r.is_c2 else '✗ No C2'} (prob={r.c2_probability:.2%})")
+
+    # Test 4: Web browsing (should be benign)
+    print("\n[4] Web Browsing (should be clean)...")
+    browse = [{'timestamp': 1705600000 + i*random.uniform(2, 30),
+               'dst_ip': f"{random.randint(1,200)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}",
+               'dst_port': 443, 'bytes_sent': random.randint(500, 3000),
+               'bytes_recv': random.randint(10000, 500000)} for i in range(20)]
+    r = sentinel.analyze(browse)
+    print(f"  {'✗ C2 DETECTED (FP!)' if r.is_c2 else '✓ Clean'} (prob={r.c2_probability:.2%})")
+
+    # Test 5: SSH keepalive
+    print("\n[5] SSH Keepalive (should be clean)...")
+    ssh = [{'timestamp': 1705600000 + i*30, 'dst_ip': '192.168.1.50', 'dst_port': 22,
+            'bytes_sent': 48, 'bytes_recv': 48} for i in range(15)]
+    r = sentinel.analyze(ssh)
+    print(f"  {'✗ C2 DETECTED (FP!)' if r.is_c2 else '✓ Clean'} (prob={r.c2_probability:.2%})")
+    print(f"  Pattern: {r.matched_legitimate_pattern}")
+
+    # Test 6: API calls (should be benign)
+    print("\n[6] API Calls (should be clean)...")
+    api = [{'timestamp': 1705600000 + i*random.uniform(0.5, 10),
+            'dst_ip': '52.85.132.99', 'dst_port': 443,
+            'bytes_sent': random.randint(100, 2000),
+            'bytes_recv': random.randint(500, 80000)} for i in range(25)]
+    r = sentinel.analyze(api)
+    print(f"  {'✗ C2 DETECTED (FP!)' if r.is_c2 else '✓ Clean'} (prob={r.c2_probability:.2%})")
+
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--epochs', type=int, default=100)
+    parser.add_argument('--samples', type=int, default=20000)
+    parser.add_argument('--batch-size', type=int, default=32)
+    parser.add_argument('--lr', type=float, default=0.0001)
+    parser.add_argument('--test-only', action='store_true')
+    args = parser.parse_args()
+
+    if args.test_only:
+        test_model()
+    else:
+        train_model(num_epochs=args.epochs, batch_size=args.batch_size,
+                   learning_rate=args.lr, num_samples=args.samples)
+        test_model()