import os, time, json, random
from datetime import datetime
import torch, torch.nn as nn, torch.nn.functional as F
from flask import Flask, request, jsonify
app = Flask(__name__)
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MAX_LINES = 500
MODEL, ENCODE, DECODE, CKPT, BLOCK_SIZE = None, None, None, None, 256
ATTACK_PROMPTS = {
"random":"","ssh":"2024-01-15 03:22:11 auth-server sshd",
"portscan":"2024-01-15 02:11:04 SNORT[3]: [1:1000001:1] PORT SCAN",
"firewall":"2024-01-15 14:33:07 FW01 kernel: [BLOCK] IN=eth0",
"webattack":"2024-01-15 11:44:22 web01 apache2:",
"malware":"THREAT_INTEL: C2_BEACON_DETECTED",
"privesc":"2024-01-15 04:12:09 db01 sudo:",
"exfil":"DLP_ALERT: [CRITICAL] Large data transfer",
"ransomware":"SIEM_ALERT: [CRITICAL] RANSOMWARE",
"exploit":"IDS_ALERT: [HIGH] Exploit attempt detected",
"siem":"SIEM_ALERT: [HIGH]",
}
RANDOM_POOL = [("ssh",.15),("portscan",.12),("firewall",.18),("webattack",.15),
("malware",.10),("privesc",.07),("exfil",.06),("ransomware",.05),
("exploit",.08),("siem",.04)]
def get_random_prompt():
types, weights = zip(*RANDOM_POOL)
t = random.choices(types, weights=weights, k=1)[0]
return ATTACK_PROMPTS[t], t
# ── Model — key names MUST match the original training notebook exactly ──
def build_model(vocab_size, n_embd, n_head, n_layer, block_size):
class Head(nn.Module):
def __init__(self, head_size):
super().__init__()
self.query = nn.Linear(n_embd, head_size, bias=False)
self.key = nn.Linear(n_embd, head_size, bias=False)
self.value = nn.Linear(n_embd, head_size, bias=False)
self.register_buffer("tril", torch.tril(torch.ones(block_size, block_size)))
self.dropout = nn.Dropout(0.0)
def forward(self, x):
B, T, C = x.shape
q, k, v = self.query(x), self.key(x), self.value(x)
w = q @ k.transpose(-2,-1) * (k.shape[-1]**-0.5)
w = w.masked_fill(self.tril[:T,:T]==0, float("-inf"))
return self.dropout(F.softmax(w, dim=-1)) @ v
class MultiHeadAttention(nn.Module):
def __init__(self, num_heads, head_size):
super().__init__()
self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
self.proj = nn.Linear(head_size * num_heads, n_embd)
self.dropout = nn.Dropout(0.0)
def forward(self, x):
return self.dropout(self.proj(torch.cat([h(x) for h in self.heads], dim=-1)))
class FeedForward(nn.Module):
def __init__(self, n_embd):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_embd, 4*n_embd), nn.GELU(),
nn.Linear(4*n_embd, n_embd), nn.Dropout(0.0)
)
def forward(self, x): return self.net(x)
class Block(nn.Module):
def __init__(self, n_embd, n_head):
super().__init__()
head_size = n_embd // n_head
self.sa = MultiHeadAttention(n_head, head_size)
self.ff = FeedForward(n_embd)
self.ln1 = nn.LayerNorm(n_embd)
self.ln2 = nn.LayerNorm(n_embd)
def forward(self, x):
x = x + self.sa(self.ln1(x))
return x + self.ff(self.ln2(x))
# ← These names MUST match what the notebook saved
class CyberLogGPT(nn.Module):
def __init__(self):
super().__init__()
self.token_embedding_table = nn.Embedding(vocab_size, n_embd)
self.position_embedding_table = nn.Embedding(block_size, n_embd)
self.blocks = nn.Sequential(*[Block(n_embd, n_head) for _ in range(n_layer)])
self.ln_f = nn.LayerNorm(n_embd)
self.lm_head = nn.Linear(n_embd, vocab_size)
def forward(self, idx, targets=None):
B, T = idx.shape
tok_emb = self.token_embedding_table(idx)
pos_emb = self.position_embedding_table(torch.arange(T, device=DEVICE))
x = self.ln_f(self.blocks(tok_emb + pos_emb))
logits = self.lm_head(x)
if targets is None: return logits, None
B, T, C = logits.shape
return logits, F.cross_entropy(logits.view(B*T,C), targets.view(B*T))
@torch.no_grad()
def generate(self, idx, n, temperature=1.0, top_k=None):
for _ in range(n):
ic = idx[:, -block_size:]
logits, _ = self(ic)
logits = logits[:, -1, :] / temperature
if top_k:
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < v[:, [-1]]] = float("-inf")
idx = torch.cat((idx, torch.multinomial(F.softmax(logits,-1), 1)), dim=1)
return idx
return CyberLogGPT()
def load_model():
global MODEL, ENCODE, DECODE, CKPT, BLOCK_SIZE
if not os.path.exists("cyberlog_gpt.pt"):
print("ERROR: cyberlog_gpt.pt not found")
return False
try:
ckpt = torch.load("cyberlog_gpt.pt", map_location=DEVICE)
cfg = ckpt["config"]
stoi, itos = ckpt["stoi"], ckpt["itos"]
BLOCK_SIZE = cfg["block_size"]
ENCODE = lambda s: [stoi[c] for c in s if c in stoi]
DECODE = lambda l: "".join([itos[i] for i in l])
m = build_model(ckpt["vocab_size"], cfg["n_embd"], cfg["n_head"],
cfg["n_layer"], cfg["block_size"]).to(DEVICE)
m.load_state_dict(ckpt["model_state_dict"])
m.eval()
MODEL, CKPT = m, ckpt
total = sum(p.numel() for p in m.parameters())
print(f"✅ Model loaded: {total/1e6:.2f}M params | "
f"train={ckpt['final_train_loss']:.4f} val={ckpt['final_val_loss']:.4f}")
return True
except Exception as e:
print(f"❌ Load error: {e}")
return False
load_model()
@app.route("/api/generate", methods=["POST"])
def api_generate():
if MODEL is None: return jsonify({"error": "Model not loaded"}), 503
d = request.get_json(silent=True) or {}
attack_type = d.get("attack_type", "random")
n_lines = max(1, min(int(d.get("n_lines", 20)), MAX_LINES))
temperature = max(0.3, min(float(d.get("temperature", 0.7)), 1.5))
top_k = max(5, min(int(d.get("top_k", 40)), 100))
fmt = d.get("format", "log")
custom_prompt = str(d.get("custom_prompt", "")).strip()[:200]
actual_type = attack_type
if custom_prompt:
prompt = custom_prompt
elif attack_type == "random":
prompt, actual_type = get_random_prompt()
else:
prompt = ATTACK_PROMPTS.get(attack_type, "")
try:
t0 = time.time()
ctx = torch.tensor(ENCODE(prompt), dtype=torch.long, device=DEVICE).unsqueeze(0) \
if prompt else torch.zeros((1,1), dtype=torch.long, device=DEVICE)
ids = MODEL.generate(ctx, min(n_lines*150, 75000),
temperature=temperature, top_k=top_k)
raw = DECODE(ids[0].tolist())
lines = [l for l in raw.split("\n") if l.strip()][:n_lines]
elapsed = round((time.time()-t0)*1000)
if fmt == "json":
entries = [{"id":i+1,"raw":l,"attack_type":actual_type,
"generated_at":datetime.utcnow().isoformat()+"Z"}
for i,l in enumerate(lines)]
output = json.dumps({"logs":entries,"count":len(entries),
"model":"CyberLog-GPT"},indent=2)
elif fmt == "csv":
rows = ["id,timestamp,raw_log,attack_type"]
for i,l in enumerate(lines):
ts = l[:19] if len(l)>19 else datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S")
safe = l.replace(',',';').replace('"','\\"')
rows.append('{},{},"{}",{}'.format(i+1, ts, safe, actual_type))
output = "\n".join(rows)
else:
output = "\n".join(lines)
return jsonify({"logs":output,"lines_count":len(lines),"chars_count":len(output),
"tokens_used":len(ids[0]),"elapsed_ms":elapsed,
"attack_type":actual_type,"format":fmt})
except Exception as e:
return jsonify({"error": str(e)}), 500
@app.route("/api/info")
def api_info():
if MODEL is None: return jsonify({"status":"not_loaded"}), 503
total = sum(p.numel() for p in MODEL.parameters())
return jsonify({"status":"ready","parameters_M":round(total/1e6,2),
"train_loss":round(CKPT["final_train_loss"],4),
"val_loss":round(CKPT["final_val_loss"],4),
"vocab_size":CKPT["vocab_size"],"block_size":BLOCK_SIZE,
"device":str(DEVICE),"max_lines":MAX_LINES})
@app.route("/api/health")
def health():
return jsonify({"status":"ok","model_loaded":MODEL is not None})
@app.route("/")
def landing():
if os.path.exists("landing.html"): return open("landing.html").read()
return open("ui.html").read() if os.path.exists("ui.html") else "CyberLog-GPT
"
@app.route("/app")
def index():
return open("ui.html").read() if os.path.exists("ui.html") else "ui.html missing
"
if __name__ == "__main__":
app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 7860)), debug=False)