Remove training files from public view

train_model.py

@@ -1,399 +0,0 @@
-#!/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()