Spaces:
Sleeping
Sleeping
File size: 2,395 Bytes
c6e6f10 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | import torch
import torch.nn as nn
import torch.optim as optim
import random
from models.gnn import PhishingGNN_Model
from pipeline.graph_engine import TopologicalGraphEngine
from config import *
def generate_synthetic_warfare(batch_size=100):
"""Generates synthetic network logs simulating both normal traffic and advanced attacks."""
logs = []
labels = []
for _ in range(batch_size):
is_attack = random.random() > 0.5
if is_attack:
logs.append({
'ip': f"{random.randint(1, 255)}.{random.randint(1,255)}.0.0",
'domain': None,
'asn': random.choice([666, 9999, 5555])
})
labels.append([1.0])
else:
logs.append({
'ip': f"104.21.{random.randint(1,100)}.{random.randint(1,255)}",
'domain': f"safe-service-{random.randint(1,100)}.com",
'asn': 13335
})
labels.append([0.0])
return logs, torch.tensor(labels, dtype=torch.float32)
def run_war_games():
print("[*] Initiating Autonomous Threat War Games...")
in_channels_dict = {'ip': 16, 'domain': 32, 'asn': 8, 'cert': 16}
model = PhishingGNN_Model(
metadata=GRAPH_METADATA,
in_channels_dict=in_channels_dict,
hidden_channels=HIDDEN_CHANNELS,
num_heads=NUM_HEADS,
num_layers=NUM_LAYERS,
dropout_rate=DROPOUT_RATE
)
optimizer = optim.AdamW(model.parameters(), lr=0.001)
criterion = nn.BCELoss()
engine = TopologicalGraphEngine()
for wave in range(50):
logs, labels = generate_synthetic_warfare(batch_size=200)
x_dict, edge_index_dict = engine.extract_and_build(logs)
optimizer.zero_grad()
predictions = model(x_dict, edge_index_dict)
valid_preds = predictions[:len(labels)]
loss = criterion(valid_preds, labels)
loss.backward()
optimizer.step()
if (wave + 1) % 10 == 0:
print(f"Attack Wave {wave+1:02d}/50 Defeated | Model Penetration Loss: {loss.item():.4f}")
model.safe_save(MODEL_SAVE_PATH)
print(f"[+] War games complete. Apex model hardened and saved to: {MODEL_SAVE_PATH}")
if __name__ == "__main__":
run_war_games() |