Update app.py

app.py

@@ -1,239 +1,180 @@
-import requests, uuid, time, threading, torch, json
-from fastapi import FastAPI
-from bs4 import BeautifulSoup
-from urllib.parse import urljoin
-from datetime import datetime
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-# ==================================================
-# CONFIG
-# ==================================================
-MODEL_NAME = "UCSB-SURFI/VulnLLM-R-7B"
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-UA = "RedTeam-Bot/1.0"
-MAX_HISTORY = 5  # her hedef için saklanacak en fazla tarama sayısı
-
-# ==================================================
-# APP + STATE
-# ==================================================
-app = FastAPI(title="Ultra Red-Team SaaS Engine")
-SCANS = {}          # scan_id -> result
-SCAN_HISTORY = {}   # target -> [scan_ids]
-
-# ==================================================
-# LOAD VULNLLM (ONCE, OFFLOAD)
-# ==================================================
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
-    device_map="auto",
-    offload_folder="offload",  # memory-safe offload
-    trust_remote_code=True
-)
-model.eval()
-
-def vulnllm_enrich(finding, context):
-    prompt = f"""
-You are an elite red-team security researcher.
-
-Finding:
-{finding}
-
-Context:
-{json.dumps(context, indent=2)}
-
-Return STRICTLY:
-- CWE
-- Risk level (Low/Medium/High/Critical)
-- Realistic exploit scenario
-- Clear remediation
-"""
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)  # daha kısa
-    if DEVICE == "cuda":
-        inputs = inputs.to(DEVICE)
-
-    with torch.no_grad():
-        out = model.generate(
-            **inputs,
-            max_new_tokens=200,  # memory-safe token sayısı
-            temperature=0.25,
-            top_p=0.95
-        )
-
-    result = tokenizer.decode(out[0], skip_special_tokens=True)
-    del inputs
-    torch.cuda.empty_cache()
-    return result
-
-# ==================================================
-# RECON AGENT
-# ==================================================
-def recon(url):
-    r = requests.get(url, timeout=10, headers={"User-Agent": UA})
-    soup = BeautifulSoup(r.text, "html.parser")
-
-    js_files = [urljoin(url, s["src"]) for s in soup.find_all("script", src=True)]
-    forms = []
-    for f in soup.find_all("form"):
-        forms.append({
-            "action": urljoin(url, f.get("action", "")),
-            "method": f.get("method", "GET").upper(),
-            "inputs": [i.get("name") for i in f.find_all("input")]
-        })
+# ... (önceki import'ların sonuna ekle)
+import networkx as nx
+import plotly.graph_objects as go
+import matplotlib.pyplot as plt
+from io import BytesIO
+import base64
+
+# ────────────────────────────────────────────────
+# Attack Graph Görselleştirme Fonksiyonları
+# ────────────────────────────────────────────────
+def create_attack_graph_data(recon_data: Dict) -> Dict:
+    """Graph verisini hazırlar (nodes, edges)"""
+    nodes = ["User", "Browser"]
+    edges = [("User", "Browser")]
+
+    forms_count = recon_data.get("forms_count", 0)
+    js_count = recon_data.get("js_files_count", 0)
+
+    if forms_count > 0:
+        nodes.append("Form Submission")
+        edges.append(("User", "Form Submission"))
+        edges.append(("Form Submission", "Backend"))
+
+    if js_count > 0:
+        nodes.append("Client JS")
+        edges.append(("Browser", "Client JS"))
+
+        # Örnek endpoint'ler (gerçekte recon'dan gelebilir)
+        for i in range(min(js_count, 4)):  # max 4 örnek göster
+            ep_name = f"API/Endpoint {i+1}"
+            nodes.append(ep_name)
+            edges.append(("Client JS", ep_name))
+
+    # Riskli noktaları vurgula (örnek)
+    risky_nodes = []
+    if forms_count > 2:
+        risky_nodes.append("Form Submission")
+    if js_count > 8:
+        risky_nodes.append("Client JS")
 
     return {
-        "headers": dict(r.headers),
-        "js_files": js_files,
-        "forms": forms,
-        "html_size": len(r.text)  # full HTML yerine boyut
-    }
-
-# ==================================================
-# JS SURFACE (AST-LIKE HEURISTIC)
-# ==================================================
-def js_surface(js_url):
-    try:
-        r = requests.get(js_url, timeout=5, headers={"User-Agent": UA})
-        endpoints = []
-        for line in r.text.splitlines():
-            if "fetch(" in line or "axios" in line:
-                endpoints.append(line.strip()[:200])
-        return endpoints
-    except:
-        return []
-
-# ==================================================
-# RED AGENT (ATTACK THINKING)
-# ==================================================
-def red_agent(recon_data, js_endpoints):
-    findings = []
-
-    headers = {k.lower(): v for k, v in recon_data["headers"].items()}
-    if "x-powered-by" in headers:
-        findings.append({
-            "title": "Technology stack disclosure",
-            "context": headers
-        })
-
-    if recon_data["forms"]:
-        findings.append({
-            "title": "User-controlled input surface",
-            "context": recon_data["forms"]
-        })
-
-    if len(js_endpoints) > 5:
-        findings.append({
-            "title": "Exposed client-side API surface",
-            "context": js_endpoints[:5]
-        })
-
-    return findings
-
-# ==================================================
-# RISK ENGINE
-# ==================================================
-def risk_score(enriched_findings):
-    score = 0
-    for e in enriched_findings:
-        txt = e.lower()
-        if "critical" in txt: score += 40
-        elif "high" in txt: score += 25
-        elif "medium" in txt: score += 10
-    return min(score, 100)
-
-# ==================================================
-# ATTACK GRAPH
-# ==================================================
-def attack_graph(js_eps, forms):
-    nodes = ["User", "Browser", "JS"]
-    edges = [{"from":"User","to":"Browser"},{"from":"Browser","to":"JS"}]
-
-    for ep in js_eps:
-        nodes.append(ep)
-        edges.append({"from":"JS","to":ep})
-
-    if forms:
-        nodes.append("FormInput")
-        edges.append({"from":"User","to":"FormInput"})
-
-    return {"nodes": list(set(nodes)), "edges": edges}
-
-# ==================================================
-# CORE SCAN PIPELINE
-# ==================================================
-def run_scan(target):
-    scan_id = str(uuid.uuid4())
-
-    recon_data = recon(target)
-
-    js_eps = []
-    for js in recon_data["js_files"]:
-        js_eps += js_surface(js)
-
-    raw_findings = red_agent(recon_data, js_eps)
-
-    enriched = []
-    for f in raw_findings:
-        enriched.append(vulnllm_enrich(f["title"], f["context"]))
-
-    risk = risk_score(enriched)
-    graph = attack_graph(js_eps, recon_data["forms"])
-
-    result = {
-        "scan_id": scan_id,
-        "target": target,
-        "time": datetime.utcnow().isoformat(),
-        "risk_score": risk,
-        "attack_graph": graph,
-        "findings_enriched": enriched,
-        "html_size": recon_data["html_size"]
+        "nodes": nodes,
+        "edges": edges,
+        "risky": risky_nodes
     }
 
-    SCANS[scan_id] = result
-    SCAN_HISTORY.setdefault(target, []).append(scan_id)
-
-    # memory-safe: eski taramaları sil
-    if len(SCAN_HISTORY[target]) > MAX_HISTORY:
-        oldest = SCAN_HISTORY[target].pop(0)
-        del SCANS[oldest]
-
-    return result
-
-# ==================================================
-# CONTINUOUS SCAN
-# ==================================================
-def continuous(target, interval):
-    while True:
-        run_scan(target)
-        time.sleep(interval)
-
-# ==================================================
-# API
-# ==================================================
-@app.post("/scan")
-def scan(target: str):
-    return run_scan(target)
-
-@app.post("/scan/continuous")
-def scan_continuous(target: str, interval: int = 3600):
-    threading.Thread(target=continuous, args=(target,interval), daemon=True).start()
-    return {"status":"scheduled","interval":interval}
-
-@app.get("/dashboard/{target}")
-def dashboard(target: str):
-    ids = SCAN_HISTORY.get(target, [])
-    return [SCANS[i] for i in ids]
-
-@app.get("/")
-def health():
-    return {
-        "status":"running",
-        "mode":"ULTRA B-MODE",
-        "model":MODEL_NAME,
-        "features":[
-            "recon","js-surface","red-agent",
-            "vulnllm-enrichment","attack-graph",
-            "risk-engine","continuous-scan"
-        ]
-    }
+def visualize_attack_graph_plotly(graph_data: Dict) -> go.Figure:
+    """Plotly ile interaktif graph"""
+    G = nx.DiGraph()
+    G.add_edges_from(graph_data["edges"])
+
+    pos = nx.spring_layout(G, seed=42)  # reproducible layout
+
+    edge_x = []
+    edge_y = []
+    for edge in G.edges():
+        x0, y0 = pos[edge[0]]
+        x1, y1 = pos[edge[1]]
+        edge_x.extend([x0, x1, None])
+        edge_y.extend([y0, y1, None])
+
+    edge_trace = go.Scatter(
+        x=edge_x, y=edge_y,
+        line=dict(width=2, color='#888'),
+        hoverinfo='none',
+        mode='lines'
+    )
+
+    node_x = []
+    node_y = []
+    node_text = []
+    node_color = []
+    for node in G.nodes():
+        x, y = pos[node]
+        node_x.append(x)
+        node_y.append(y)
+        node_text.append(node)
+        if node in graph_data["risky"]:
+            node_color.append('#ff4444')  # kırmızı = riskli
+        else:
+            node_color.append('#1f77b4')  # mavi = normal
+
+    node_trace = go.Scatter(
+        x=node_x, y=node_y,
+        mode='markers+text',
+        hoverinfo='text',
+        text=node_text,
+        textposition="top center",
+        marker=dict(
+            showscale=False,
+            color=node_color,
+            size=30,
+            line_width=2
+        )
+    )
+
+    fig = go.Figure(data=[edge_trace, node_trace],
+                    layout=go.Layout(
+                        title='Attack Graph Visualization',
+                        titlefont_size=16,
+                        showlegend=False,
+                        hovermode='closest',
+                        margin=dict(b=20, l=5, r=5, t=40),
+                        xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+                        yaxis=dict(showgrid=False, zeroline=False, showticklabels=False)
+                    ))
+    return fig
+
+def visualize_attack_graph_matplotlib(graph_data: Dict) -> str:
+    """Fallback: Matplotlib → base64 PNG"""
+    G = nx.DiGraph()
+    G.add_edges_from(graph_data["edges"])
+
+    fig, ax = plt.subplots(figsize=(8, 6))
+    pos = nx.spring_layout(G, seed=42)
+
+    node_colors = ['red' if n in graph_data["risky"] else 'lightblue' for n in G.nodes()]
+
+    nx.draw(G, pos, with_labels=True,
+            node_color=node_colors,
+            node_size=2200,
+            font_size=10,
+            font_weight='bold',
+            arrows=True,
+            arrowstyle='->',
+            arrowsize=20,
+            ax=ax)
+
+    ax.set_title("Attack Graph (Static)")
+
+    buf = BytesIO()
+    plt.savefig(buf, format='png', bbox_inches='tight')
+    buf.seek(0)
+    img_base64 = base64.b64encode(buf.read()).decode('utf-8')
+    plt.close(fig)
+    return f"data:image/png;base64,{img_base64}"
+
+# ────────────────────────────────────────────────
+# full_vuln_scan fonksiyonunu güncelle (graph kısmı)
+# ────────────────────────────────────────────────
+def full_vuln_scan(target_url: str, progress=gr.Progress(track_tqdm=True)):
+    # ... (önceki kod aynı, recon_data kısmından sonra ekle)
+
+    progress(0.75, desc="Attack Graph oluşturuluyor...")
+
+    graph_data = create_attack_graph_data(recon_data)
+    plotly_fig = visualize_attack_graph_plotly(graph_data)
+    # matplotlib_fallback = visualize_attack_graph_matplotlib(graph_data)  # istersen fallback ekle
+
+    # ... (diğer sonuçlar aynı)
+
+    return (
+        result_summary,
+        json.dumps(enriched_findings, indent=2, ensure_ascii=False),
+        plotly_fig,   # ← Plotly Figure direkt Plot component'e gider
+        history_md
+    )
+
+# ────────────────────────────────────────────────
+# Gradio Blocks güncellemesi (Attack Graph Tab)
+# ────────────────────────────────────────────────
+with gr.Blocks(...) as demo:
+    # ... önceki kısımlar aynı
+
+    with gr.Tabs():
+        # ... diğer tab'lar aynı
+        with gr.Tab("Attack Graph"):
+            gr.Markdown("### Potansiyel Saldırı Yolu Görselleştirmesi")
+            gr.Markdown("(Kırmızı node'lar yüksek riskli alanları gösterir)")
+            graph_plot = gr.Plot(label="Interactive Attack Graph (Plotly)")
+
+    # Events güncelle
+    scan_button.click(
+        fn=full_vuln_scan,
+        inputs=target_input,
+        outputs=[summary_output, json_output, graph_plot, history_output],
+        # ...
+    )
+
+    # ... kalan kısım aynı