Add phase 2 multi-task training script

train_model_phase2.py

@@ -0,0 +1,704 @@
+#!/usr/bin/env python3
+"""
+C2Sentinel Training Script - Phase 2: Multi-Task Learning & Adversarial Hardening
+
+Phase 2 adds:
+1. C2 type classification (10 framework types)
+2. Adversarial beacon patterns (jitter, domain fronting, etc.)
+3. Enhanced benign patterns to reduce false positives
+4. Confidence calibration
+"""
+
+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
+import os
+
+from c2sentinel import (
+    C2Sentinel, C2SentinelConfig, LogBERTC2Sentinel,
+    FeatureExtractor
+)
+from safetensors.torch import save_file, load_file
+
+
+class C2TrafficDatasetPhase2(Dataset):
+    """Enhanced dataset with adversarial patterns and multi-task labels."""
+
+    def __init__(self, num_samples=30000, normalize=True, norm_params_path=None):
+        self.samples = []
+        self.labels = []
+        self.c2_types = []  # 0=benign, 1-10=C2 framework types
+        self.feature_extractor = FeatureExtractor()
+
+        print(f"Generating {num_samples} phase 2 training samples...")
+
+        num_c2 = num_samples // 2
+        num_adversarial = num_c2 // 3  # 1/3 of C2 samples are adversarial
+        num_standard_c2 = num_c2 - num_adversarial
+        num_benign = num_samples - num_c2
+
+        # Generate standard C2 samples
+        print(f"\n[1/4] Standard C2 samples ({num_standard_c2})...")
+        for _ in tqdm(range(num_standard_c2), desc="Standard C2"):
+            connections, c2_type = self._generate_standard_c2()
+            features = self.feature_extractor.extract_features(connections)
+            self.samples.append(features)
+            self.labels.append(1)
+            self.c2_types.append(c2_type)
+
+        # Generate adversarial C2 samples (harder to detect)
+        print(f"\n[2/4] Adversarial C2 samples ({num_adversarial})...")
+        for _ in tqdm(range(num_adversarial), desc="Adversarial C2"):
+            connections, c2_type = self._generate_adversarial_c2()
+            features = self.feature_extractor.extract_features(connections)
+            self.samples.append(features)
+            self.labels.append(1)
+            self.c2_types.append(c2_type)
+
+        # Generate standard benign samples
+        num_standard_benign = num_benign * 2 // 3
+        print(f"\n[3/4] Standard benign samples ({num_standard_benign})...")
+        for _ in tqdm(range(num_standard_benign), desc="Standard Benign"):
+            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)
+
+        # Generate edge-case benign samples (look like C2 but aren't)
+        num_edge_benign = num_benign - num_standard_benign
+        print(f"\n[4/4] Edge-case benign samples ({num_edge_benign})...")
+        for _ in tqdm(range(num_edge_benign), desc="Edge Benign"):
+            connections = self._generate_edge_case_benign()
+            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)
+
+        # Load or compute normalization
+        if normalize:
+            if norm_params_path and os.path.exists(norm_params_path):
+                print(f"Loading normalization from {norm_params_path}")
+                params = np.load(norm_params_path)
+                self.mean = params['mean']
+                self.std = params['std']
+            else:
+                self.mean = np.mean(self.samples, axis=0)
+                self.std = np.std(self.samples, axis=0) + 1e-8
+                np.savez('normalization_params.npz', mean=self.mean, std=self.std)
+
+            self.samples = (self.samples - self.mean) / self.std
+            print(f"Feature stats - mean range: [{self.mean.min():.2f}, {self.mean.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]
+
+        # Report distribution
+        c2_count = np.sum(self.labels)
+        print(f"\nDataset: {len(self.labels)} samples")
+        print(f"  C2: {int(c2_count)} ({100*c2_count/len(self.labels):.1f}%)")
+        print(f"  Benign: {int(len(self.labels) - c2_count)} ({100*(1 - c2_count/len(self.labels)):.1f}%)")
+        print(f"  C2 type distribution: {np.bincount(self.c2_types[self.labels == 1].astype(int), minlength=11)[1:]}")
+
+    def _generate_standard_c2(self):
+        """Generate standard C2 beacon patterns."""
+        c2_type = random.randint(1, 10)
+
+        # C2 type characteristics
+        c2_profiles = {
+            1: {'name': 'Metasploit', 'interval': (2, 15), 'jitter': 0.1, 'ports': [4444, 4445, 5555]},
+            2: {'name': 'Cobalt Strike', 'interval': (30, 90), 'jitter': 0.2, 'ports': [443, 8443]},
+            3: {'name': 'Empire', 'interval': (5, 30), 'jitter': 0.15, 'ports': [443, 8080]},
+            4: {'name': 'Covenant', 'interval': (10, 60), 'jitter': 0.1, 'ports': [443, 80]},
+            5: {'name': 'Sliver', 'interval': (30, 120), 'jitter': 0.3, 'ports': [443, 8888]},
+            6: {'name': 'Brute Ratel', 'interval': (60, 180), 'jitter': 0.2, 'ports': [443]},
+            7: {'name': 'Mythic', 'interval': (15, 60), 'jitter': 0.15, 'ports': [443, 7443]},
+            8: {'name': 'PoshC2', 'interval': (10, 45), 'jitter': 0.2, 'ports': [443, 8000]},
+            9: {'name': 'Havoc', 'interval': (20, 90), 'jitter': 0.25, 'ports': [443, 40056]},
+            10: {'name': 'APT Custom', 'interval': (120, 600), 'jitter': 0.05, 'ports': [443, 80]},
+        }
+
+        profile = c2_profiles[c2_type]
+        interval = random.uniform(*profile['interval'])
+        jitter = profile['jitter']
+        port = random.choice(profile['ports'])
+
+        # Beacon characteristics
+        bytes_sent = random.randint(60, 200)
+        bytes_recv = random.randint(40, 150)
+
+        dst_ip = f"{random.randint(1,223)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}"
+        num_connections = random.randint(12, 50)
+
+        connections = []
+        timestamp = 1705600000
+
+        for _ in range(num_connections):
+            actual_interval = interval * (1 + random.uniform(-jitter, jitter))
+            timestamp += actual_interval
+            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_adversarial_c2(self):
+        """Generate adversarial C2 patterns that try to evade detection."""
+        c2_type = random.randint(1, 10)
+        evasion = random.choice(['high_jitter', 'variable_size', 'burst_pattern', 'domain_rotation', 'mixed'])
+
+        base_interval = random.uniform(30, 120)
+        dst_ip = f"{random.randint(1,223)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}"
+        num_connections = random.randint(15, 60)
+
+        connections = []
+        timestamp = 1705600000
+
+        if evasion == 'high_jitter':
+            # High jitter to look like normal traffic
+            jitter = random.uniform(0.4, 0.7)
+            for _ in range(num_connections):
+                actual_interval = base_interval * (1 + random.uniform(-jitter, jitter))
+                timestamp += max(5, actual_interval)
+                connections.append({
+                    'timestamp': timestamp,
+                    'dst_ip': dst_ip,
+                    'dst_port': 443,
+                    'bytes_sent': random.randint(80, 150),
+                    'bytes_recv': random.randint(50, 100),
+                    'protocol': 'tcp'
+                })
+
+        elif evasion == 'variable_size':
+            # Variable packet sizes but consistent timing
+            for _ in range(num_connections):
+                timestamp += base_interval * (1 + random.uniform(-0.1, 0.1))
+                # Sizes vary more but still bounded
+                connections.append({
+                    'timestamp': timestamp,
+                    'dst_ip': dst_ip,
+                    'dst_port': 443,
+                    'bytes_sent': random.randint(50, 500),
+                    'bytes_recv': random.randint(40, 400),
+                    'protocol': 'tcp'
+                })
+
+        elif evasion == 'burst_pattern':
+            # Beacon with occasional bursts (simulating commands)
+            for i in range(num_connections):
+                if i % 8 == 0:
+                    # Burst
+                    for _ in range(random.randint(2, 5)):
+                        timestamp += random.uniform(0.5, 3)
+                        connections.append({
+                            'timestamp': timestamp,
+                            'dst_ip': dst_ip,
+                            'dst_port': 443,
+                            'bytes_sent': random.randint(200, 2000),
+                            'bytes_recv': random.randint(500, 5000),
+                            'protocol': 'tcp'
+                        })
+                else:
+                    timestamp += base_interval * (1 + random.uniform(-0.15, 0.15))
+                    connections.append({
+                        'timestamp': timestamp,
+                        'dst_ip': dst_ip,
+                        'dst_port': 443,
+                        'bytes_sent': random.randint(80, 120),
+                        'bytes_recv': random.randint(50, 80),
+                        'protocol': 'tcp'
+                    })
+
+        elif evasion == 'domain_rotation':
+            # Multiple IPs (CDN-like) but same beacon pattern
+            ips = [f"{random.randint(1,223)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}"
+                   for _ in range(random.randint(2, 4))]
+            for _ in range(num_connections):
+                timestamp += base_interval * (1 + random.uniform(-0.2, 0.2))
+                connections.append({
+                    'timestamp': timestamp,
+                    'dst_ip': random.choice(ips),
+                    'dst_port': 443,
+                    'bytes_sent': random.randint(80, 150),
+                    'bytes_recv': random.randint(50, 100),
+                    'protocol': 'tcp'
+                })
+
+        else:  # mixed
+            # Mix of evasion techniques
+            jitter = random.uniform(0.25, 0.5)
+            for _ in range(num_connections):
+                actual_interval = base_interval * (1 + random.uniform(-jitter, jitter))
+                timestamp += max(3, actual_interval)
+                connections.append({
+                    'timestamp': timestamp,
+                    'dst_ip': dst_ip,
+                    'dst_port': random.choice([443, 8443]),
+                    'bytes_sent': random.randint(60, 300),
+                    'bytes_recv': random.randint(40, 200),
+                    'protocol': 'tcp'
+                })
+
+        return connections, c2_type
+
+    def _generate_benign_traffic(self):
+        """Generate standard benign traffic patterns."""
+        pattern = random.choice(['browsing', 'api', 'streaming', 'interactive', 'download', 'email'])
+
+        connections = []
+        timestamp = 1705600000
+
+        if pattern == 'browsing':
+            for _ in range(random.randint(10, 50)):
+                timestamp += random.uniform(0.5, 60)
+                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':
+            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, 60)):
+                timestamp += random.uniform(0.1, 30)
+                connections.append({
+                    'timestamp': timestamp,
+                    'dst_ip': dst_ip,
+                    'dst_port': 443,
+                    'bytes_sent': random.randint(100, 5000),
+                    'bytes_recv': random.randint(200, 200000),
+                    'protocol': 'tcp'
+                })
+
+        elif pattern == 'streaming':
+            dst_ip = f"{random.randint(1,223)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}"
+            for _ in range(random.randint(30, 100)):
+                timestamp += random.uniform(0.02, 2)
+                connections.append({
+                    'timestamp': timestamp,
+                    'dst_ip': dst_ip,
+                    'dst_port': 443,
+                    'bytes_sent': random.randint(30, 200),
+                    'bytes_recv': random.randint(5000, 200000),
+                    'protocol': 'tcp'
+                })
+
+        elif pattern == 'interactive':
+            dst_ip = f"192.168.{random.randint(0,255)}.{random.randint(1,254)}"
+            for _ in range(random.randint(20, 80)):
+                if random.random() < 0.3:
+                    timestamp += random.uniform(5, 60)
+                else:
+                    timestamp += random.uniform(0.1, 3)
+                connections.append({
+                    'timestamp': timestamp,
+                    'dst_ip': dst_ip,
+                    'dst_port': 22,
+                    'bytes_sent': random.randint(20, 1000),
+                    'bytes_recv': random.randint(50, 30000),
+                    'protocol': 'tcp'
+                })
+
+        elif pattern == 'download':
+            dst_ip = f"{random.randint(1,223)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}"
+            for _ in range(random.randint(50, 200)):
+                timestamp += random.uniform(0.01, 0.5)
+                connections.append({
+                    'timestamp': timestamp,
+                    'dst_ip': dst_ip,
+                    'dst_port': 443,
+                    'bytes_sent': random.randint(40, 100),
+                    'bytes_recv': random.randint(10000, 65000),
+                    'protocol': 'tcp'
+                })
+
+        else:  # email
+            for _ in range(random.randint(5, 20)):
+                timestamp += random.uniform(30, 300)
+                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([443, 993, 587]),
+                    'bytes_sent': random.randint(500, 50000),
+                    'bytes_recv': random.randint(1000, 500000),
+                    'protocol': 'tcp'
+                })
+
+        return connections
+
+    def _generate_edge_case_benign(self):
+        """Generate benign traffic that looks like C2 but isn't."""
+        pattern = random.choice([
+            'heartbeat', 'monitoring', 'sync', 'keepalive', 'polling', 'iot'
+        ])
+
+        connections = []
+        timestamp = 1705600000
+        dst_ip = f"{random.randint(1,223)}.{random.randint(0,255)}.{random.randint(0,255)}.{random.randint(1,254)}"
+
+        if pattern == 'heartbeat':
+            # Regular heartbeat but with large, variable responses
+            interval = random.uniform(30, 120)
+            for _ in range(random.randint(15, 40)):
+                timestamp += interval * (1 + random.uniform(-0.05, 0.05))
+                connections.append({
+                    'timestamp': timestamp,
+                    'dst_ip': dst_ip,
+                    'dst_port': 443,
+                    'bytes_sent': random.randint(50, 100),
+                    'bytes_recv': random.randint(1000, 50000),  # Large variable responses
+                    'protocol': 'tcp'
+                })
+
+        elif pattern == 'monitoring':
+            # Regular monitoring checks with status responses
+            interval = random.uniform(60, 300)
+            for _ in range(random.randint(10, 30)):
+                timestamp += interval * (1 + random.uniform(-0.1, 0.1))
+                connections.append({
+                    'timestamp': timestamp,
+                    'dst_ip': dst_ip,
+                    'dst_port': random.choice([443, 8443, 9090]),
+                    'bytes_sent': random.randint(100, 500),
+                    'bytes_recv': random.randint(500, 10000),
+                    'protocol': 'tcp'
+                })
+
+        elif pattern == 'sync':
+            # Periodic sync with variable data
+            interval = random.uniform(300, 900)
+            for _ in range(random.randint(8, 20)):
+                timestamp += interval * (1 + random.uniform(-0.15, 0.15))
+                connections.append({
+                    'timestamp': timestamp,
+                    'dst_ip': dst_ip,
+                    'dst_port': 443,
+                    'bytes_sent': random.randint(1000, 100000),
+                    'bytes_recv': random.randint(1000, 100000),
+                    'protocol': 'tcp'
+                })
+
+        elif pattern == 'keepalive':
+            # SSH/VPN keepalive - very regular but small
+            interval = random.uniform(15, 60)
+            for _ in range(random.randint(20, 60)):
+                timestamp += interval * (1 + random.uniform(-0.02, 0.02))
+                connections.append({
+                    'timestamp': timestamp,
+                    'dst_ip': f"192.168.{random.randint(0,255)}.{random.randint(1,254)}",
+                    'dst_port': random.choice([22, 1194, 443]),
+                    'bytes_sent': random.randint(40, 80),
+                    'bytes_recv': random.randint(40, 80),
+                    'protocol': 'tcp'
+                })
+
+        elif pattern == 'polling':
+            # API polling - regular but with variable responses
+            interval = random.uniform(10, 60)
+            for _ in range(random.randint(20, 50)):
+                timestamp += interval * (1 + random.uniform(-0.2, 0.2))
+                connections.append({
+                    'timestamp': timestamp,
+                    'dst_ip': dst_ip,
+                    'dst_port': 443,
+                    'bytes_sent': random.randint(80, 200),
+                    'bytes_recv': random.randint(100, 50000),  # Highly variable
+                    'protocol': 'tcp'
+                })
+
+        else:  # iot
+            # IoT device - regular small packets
+            interval = random.uniform(60, 300)
+            for _ in range(random.randint(15, 40)):
+                timestamp += interval * (1 + random.uniform(-0.1, 0.1))
+                connections.append({
+                    'timestamp': timestamp,
+                    'dst_ip': dst_ip,
+                    'dst_port': random.choice([443, 8883, 1883]),
+                    'bytes_sent': random.randint(50, 200),
+                    'bytes_recv': random.randint(50, 200),
+                    '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_phase2(
+    pretrained_path='c2_sentinel.safetensors',
+    num_epochs=50,
+    batch_size=32,
+    learning_rate=0.00005,  # Lower LR for fine-tuning
+    num_samples=30000
+):
+    """Phase 2 training with multi-task learning."""
+
+    print("=" * 70)
+    print("C2Sentinel Phase 2 Training - Multi-Task Learning")
+    print("=" * 70)
+
+    config = C2SentinelConfig()
+    model = LogBERTC2Sentinel(config)
+
+    # Load pretrained weights
+    if os.path.exists(pretrained_path):
+        print(f"Loading pretrained weights from {pretrained_path}")
+        state_dict = load_file(pretrained_path)
+        model.load_state_dict(state_dict)
+    else:
+        print("WARNING: No pretrained weights found, training from scratch")
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Device: {device}")
+    model.to(device)
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Model parameters: {total_params:,}")
+
+    # Create dataset with existing normalization params
+    norm_path = 'normalization_params.npz' if os.path.exists('normalization_params.npz') else None
+    dataset = C2TrafficDatasetPhase2(num_samples=num_samples, normalize=True, norm_params_path=norm_path)
+
+    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}")
+
+    # Multi-task loss: C2 detection + C2 type classification
+    c2_criterion = nn.BCEWithLogitsLoss()
+    type_criterion = nn.CrossEntropyLoss(ignore_index=0)  # Ignore benign (0)
+
+    # Lower LR for fine-tuning
+    optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.01)
+
+    # Cosine annealing
+    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)
+
+    best_val_acc = 0
+    patience = 10
+    patience_counter = 0
+
+    for epoch in range(num_epochs):
+        model.train()
+        train_loss = 0
+        train_c2_correct = 0
+        train_type_correct = 0
+        train_total = 0
+        train_c2_samples = 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)
+            c2_types = batch['c2_type'].to(device)
+
+            optimizer.zero_grad()
+            outputs = model(features)
+
+            # C2 detection loss
+            c2_loss = c2_criterion(outputs['c2_logits'].squeeze(), labels)
+
+            # C2 type classification loss (only for C2 samples)
+            c2_mask = labels == 1
+            if c2_mask.sum() > 0 and 'type_logits' in outputs:
+                type_loss = type_criterion(outputs['type_logits'][c2_mask], c2_types[c2_mask])
+                loss = c2_loss + 0.3 * type_loss  # Weighted combination
+            else:
+                loss = c2_loss
+
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+            optimizer.step()
+
+            train_loss += loss.item()
+
+            # C2 detection accuracy
+            c2_preds = (torch.sigmoid(outputs['c2_logits'].squeeze()) > 0.5).float()
+            train_c2_correct += (c2_preds == labels).sum().item()
+            train_total += labels.size(0)
+
+            # Type classification accuracy (for C2 samples only)
+            if c2_mask.sum() > 0 and 'type_logits' in outputs:
+                type_preds = outputs['type_logits'][c2_mask].argmax(dim=1)
+                train_type_correct += (type_preds == c2_types[c2_mask]).sum().item()
+                train_c2_samples += c2_mask.sum().item()
+
+        scheduler.step()
+
+        # Validation
+        model.eval()
+        val_c2_correct = 0
+        val_type_correct = 0
+        val_total = 0
+        val_c2_samples = 0
+        val_loss = 0
+
+        # Track per-class metrics
+        val_tp, val_fp, val_tn, val_fn = 0, 0, 0, 0
+
+        with torch.no_grad():
+            for batch in val_loader:
+                features = batch['features'].to(device)
+                labels = batch['label'].to(device)
+                c2_types = batch['c2_type'].to(device)
+
+                outputs = model(features)
+
+                c2_loss = c2_criterion(outputs['c2_logits'].squeeze(), labels)
+                val_loss += c2_loss.item()
+
+                c2_preds = (torch.sigmoid(outputs['c2_logits'].squeeze()) > 0.5).float()
+                val_c2_correct += (c2_preds == labels).sum().item()
+                val_total += labels.size(0)
+
+                # Confusion matrix stats
+                val_tp += ((c2_preds == 1) & (labels == 1)).sum().item()
+                val_fp += ((c2_preds == 1) & (labels == 0)).sum().item()
+                val_tn += ((c2_preds == 0) & (labels == 0)).sum().item()
+                val_fn += ((c2_preds == 0) & (labels == 1)).sum().item()
+
+                c2_mask = labels == 1
+                if c2_mask.sum() > 0 and 'type_logits' in outputs:
+                    type_preds = outputs['type_logits'][c2_mask].argmax(dim=1)
+                    val_type_correct += (type_preds == c2_types[c2_mask]).sum().item()
+                    val_c2_samples += c2_mask.sum().item()
+
+        train_c2_acc = 100 * train_c2_correct / train_total
+        train_type_acc = 100 * train_type_correct / train_c2_samples if train_c2_samples > 0 else 0
+        val_c2_acc = 100 * val_c2_correct / val_total
+        val_type_acc = 100 * val_type_correct / val_c2_samples if val_c2_samples > 0 else 0
+
+        precision = val_tp / (val_tp + val_fp) if (val_tp + val_fp) > 0 else 0
+        recall = val_tp / (val_tp + val_fn) if (val_tp + val_fn) > 0 else 0
+        f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
+
+        lr = optimizer.param_groups[0]['lr']
+
+        print(f"Epoch {epoch+1}: Loss={train_loss/len(train_loader):.4f}, "
+              f"Train C2={train_c2_acc:.1f}%, Val C2={val_c2_acc:.1f}%, "
+              f"Type Acc={val_type_acc:.1f}%, P={precision:.2f}, R={recall:.2f}, F1={f1:.2f}")
+
+        if val_c2_acc > best_val_acc:
+            best_val_acc = val_c2_acc
+            patience_counter = 0
+            save_file(model.state_dict(), 'c2_sentinel.safetensors')
+            print(f"  -> Saved (Val: {val_c2_acc:.1f}%, F1: {f1:.2f})")
+        else:
+            patience_counter += 1
+            if patience_counter >= patience:
+                print(f"Early stopping at epoch {epoch+1}")
+                break
+
+    print(f"\nBest validation C2 accuracy: {best_val_acc:.1f}%")
+    return model, config
+
+
+def test_adversarial():
+    """Test model on adversarial patterns."""
+    print("\n" + "=" * 70)
+    print("Adversarial Pattern Testing")
+    print("=" * 70)
+
+    sentinel = C2Sentinel.load('c2_sentinel')
+
+    test_cases = [
+        ("High-jitter Cobalt Strike", [
+            {'timestamp': 1705600000 + i * 60 * (1 + random.uniform(-0.5, 0.5)),
+             'dst_ip': '185.234.72.19', 'dst_port': 443,
+             'bytes_sent': 92, 'bytes_recv': 48}
+            for i in range(20)
+        ]),
+        ("Burst pattern beacon", [
+            {'timestamp': 1705600000 + (i * 60 if i % 5 != 0 else i * 60 + random.uniform(0, 5)),
+             'dst_ip': '45.33.32.156', 'dst_port': 443,
+             'bytes_sent': 100 if i % 5 != 0 else random.randint(500, 2000),
+             'bytes_recv': 60 if i % 5 != 0 else random.randint(1000, 5000)}
+            for i in range(25)
+        ]),
+        ("Variable size beacon", [
+            {'timestamp': 1705600000 + i * 45,
+             'dst_ip': '10.10.10.10', 'dst_port': 4444,
+             'bytes_sent': random.randint(50, 300),
+             'bytes_recv': random.randint(40, 250)}
+            for i in range(18)
+        ]),
+        ("SSH keepalive (should be clean)", [
+            {'timestamp': 1705600000 + i * 30,
+             'dst_ip': '192.168.1.50', 'dst_port': 22,
+             'bytes_sent': 48, 'bytes_recv': 48}
+            for i in range(20)
+        ]),
+        ("API polling (should be clean)", [
+            {'timestamp': 1705600000 + i * random.uniform(25, 35),
+             'dst_ip': '52.85.132.99', 'dst_port': 443,
+             'bytes_sent': 150, 'bytes_recv': random.randint(500, 50000)}
+            for i in range(25)
+        ]),
+    ]
+
+    for name, connections in test_cases:
+        result = sentinel.analyze(connections)
+        status = "C2 DETECTED" if result.is_c2 else "Clean"
+        print(f"\n{name}:")
+        print(f"  {status} (prob={result.c2_probability:.2%})")
+        if result.risk_factors:
+            for rf in result.risk_factors[:3]:
+                print(f"    - {rf}")
+
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--epochs', type=int, default=50)
+    parser.add_argument('--samples', type=int, default=30000)
+    parser.add_argument('--batch-size', type=int, default=32)
+    parser.add_argument('--lr', type=float, default=0.00005)
+    parser.add_argument('--pretrained', type=str, default='c2_sentinel.safetensors')
+    parser.add_argument('--test-only', action='store_true')
+    args = parser.parse_args()
+
+    if args.test_only:
+        test_adversarial()
+    else:
+        train_phase2(
+            pretrained_path=args.pretrained,
+            num_epochs=args.epochs,
+            batch_size=args.batch_size,
+            learning_rate=args.lr,
+            num_samples=args.samples
+        )
+        test_adversarial()