app.py

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 "<h1>CyberLog-GPT</h1>"

@app.route("/app")
def index():
    return open("ui.html").read() if os.path.exists("ui.html") else "<h1>ui.html missing</h1>"

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 7860)), debug=False)