Spaces:
Sleeping
Sleeping
| """ | |
| Web Exploitation generator (Red Team) — produces diverse web-attack scenarios | |
| with realistic HTTP request/response pairs, vulnerability descriptions, and | |
| a flag-based answer. | |
| The generator cycles through 10 attack types (XSS, SQLi, CSRF, SSRF, IDOR, | |
| LFI, XXE, CMDi, Auth Bypass, File Upload) using creative seeds and AI calls. | |
| """ | |
| ## Generators AGENTS.md contract: | |
| ## POOL_TARGET / POOL_THRESHOLD / POOL_BATCH (module-level) | |
| ## async get_pool_count(team_role) -> int | |
| ## async refill_pool(team_role, count) -> int | |
| ## async start_pool_watcher(team_role) # long-running asyncio Task | |
| import asyncio | |
| import hashlib | |
| import json | |
| import logging | |
| import random | |
| import time | |
| from dataclasses import dataclass | |
| from typing import Optional | |
| import httpx | |
| from app.core.config import ( | |
| SUPABASE_URL, SUPABASE_ANON_KEY, | |
| CLOUDFLARE_API_TOKEN, CLOUDFLARE_ACCOUNT_ID, CLOUDFLARE_URL, | |
| GROQ_API_KEY, GROQ_API_URL, | |
| NVIDIA_API_KEY, NVIDIA_MODEL, NVIDIA_URL, | |
| MISTRAL_API_KEY, MISTRAL_MODEL, MISTRAL_API_URL, | |
| ) | |
| # ── Pool tuning ─────────────────────────────────────────────────────── | |
| POOL_TARGET = 5 | |
| POOL_THRESHOLD = 2 | |
| POOL_BATCH = 3 | |
| # ── AI backoff ───────────────────────────────────────────────────────── | |
| _AI_BACKOFF_UNTIL: dict[str, float] = {} | |
| # Concurrency control | |
| _POOL_LOCKS: dict[str, asyncio.Lock] = {} | |
| def _get_pool_lock(team_role: str) -> asyncio.Lock: | |
| if team_role not in _POOL_LOCKS: | |
| _POOL_LOCKS[team_role] = asyncio.Lock() | |
| return _POOL_LOCKS[team_role] | |
| # ── Topic rotation ───────────────────────────────────────────────────── | |
| WEB_EXPLOIT_TOPICS = [ | |
| "xss", "sqli", "csrf", "ssrf", "idor", | |
| "lfi", "xxe", "cmdi", "auth", "upload", | |
| ] | |
| _TOPIC_INDEX: dict[str, int] = {} # team_role -> last index | |
| # Cloudflare model fallback chain (same as code_fixing.py) | |
| CLOUDFLARE_MODEL_FALLBACKS = [ | |
| "@cf/qwen/qwen2.5-coder-32b-instruct", | |
| "@cf/meta/llama-3.3-70b-instruct-fp8-fast", | |
| "@cf/meta/llama-3.1-70b-instruct", | |
| "@cf/mistralai/mistral-small-3.1-24b-instruct", | |
| "@cf/openai/gpt-oss-120b", | |
| "@cf/openai/gpt-oss-20b", | |
| "@cf/meta/llama-3.1-8b-instruct", | |
| ] | |
| _CURATED_SEEDS: list[dict] = [] | |
| # ── Data model ───────────────────────────────────────────────────────── | |
| class WebExploitChallenge: | |
| title: str | |
| story: str | |
| task_outline: str | |
| vulnerability_type: str | |
| vulnerability_class: str | |
| vulnerability_description: str | |
| http_request: str | |
| http_response: str | |
| flag_preview: str | |
| flag_hash: str | |
| hints: list[dict] | |
| difficulty: str | |
| xp_reward: int = 200 | |
| topic: str = "" | |
| team_role: str = "red" | |
| module: str = "web-exploitation" | |
| # ── Helpers ──────────────────────────────────────────────────────────── | |
| def _sb_headers(content_type: bool = False) -> dict: | |
| h = { | |
| "apikey": SUPABASE_ANON_KEY, | |
| "Authorization": f"Bearer {SUPABASE_ANON_KEY}", | |
| } | |
| if content_type: | |
| h["Content-Type"] = "application/json" | |
| h["Prefer"] = "return=representation" | |
| return h | |
| async def _supabase_select(table: str, query: str) -> list[dict]: | |
| url = f"{SUPABASE_URL}/rest/v1/{table}?{query}" | |
| async with httpx.AsyncClient(timeout=20) as client: | |
| r = await client.get(url, headers=_sb_headers()) | |
| if r.status_code == 200: | |
| return r.json() | |
| return [] | |
| async def _supabase_insert(table: str, row: dict) -> Optional[dict]: | |
| url = f"{SUPABASE_URL}/rest/v1/{table}" | |
| async with httpx.AsyncClient(timeout=20) as client: | |
| r = await client.post(url, json=row, headers=_sb_headers(content_type=True)) | |
| if r.status_code in (200, 201): | |
| rows = r.json() | |
| return rows[0] if rows else None | |
| logging.error(f"[web_exploit] insert error {r.status_code}: {r.text[:300]}") | |
| return None | |
| def _compute_flag_hash(flag: str) -> str: | |
| return hashlib.sha256(flag.encode("utf-8")).hexdigest() | |
| def _generate_flag() -> str: | |
| """Generate a unique CyberArena{...} flag.""" | |
| rand_hex = hashlib.sha256(str(random.getrandbits(256)).encode()).hexdigest()[:16] | |
| return f"CyberArena{{{rand_hex}}}" | |
| def _extract_json(text: str) -> Optional[dict]: | |
| """Robust JSON extraction from LLM response.""" | |
| import re as _re | |
| text = text.strip() | |
| # Try direct parse | |
| try: | |
| return json.loads(text) | |
| except json.JSONDecodeError: | |
| pass | |
| # Try to find JSON block | |
| m = _re.search(r"\{[\s\S]*\}", text) | |
| if m: | |
| try: | |
| return json.loads(m.group()) | |
| except json.JSONDecodeError: | |
| pass | |
| return None | |
| # ── Curated Seeds (6 diverse scenarios, one per method) ─────────────── | |
| def _build_seeds() -> list[dict]: | |
| seeds = [] | |
| # 1. SQLi | |
| flag1 = _generate_flag() | |
| seeds.append({ | |
| "title": "ثغرة حقن SQL في لوحة تحكم المخزون", | |
| "story": "أثناء اختبار اختراق لموقع متجر إلكتروني، اكتشفت أن صفحة عرض المنتجات تستخدم معامل `id` في عنوان URL بشكل مباشر في استعلام SQL دون تنقية. يمكنك استغلال هذه الثغرة لسحب جميع المنتجات المخفية والبيانات الحساسة.", | |
| "task_outline": "حاول استغلال ثغرة SQL Injection في معامل id بالرابط /api/products?id=1. استخدم تقنية UNION لعرض جميع المنتجات المخفية واستخرج العَلم.", | |
| "vulnerability_type": "sql-injection", | |
| "vulnerability_class": "Injection", | |
| "vulnerability_description": "ثغرة SQL Injection تسمح للمهاجم بحقن استعلامات SQL خبيثة عبر مدخلات المستخدم، مما قد يؤدي إلى تسرب البيانات أو تدمير قاعدة البيانات.", | |
| "http_request": "GET /api/products?id=1 UNION SELECT 1,2,3,4,5,6-- HTTP/1.1\nHost: target.apex-train.com\nUser-Agent: Mozilla/5.0\nAccept: application/json", | |
| "http_response": "HTTP/1.1 200 OK\nContent-Type: application/json\n\n[{\"id\":1,\"name\":\"لابتوب\",\"price\":2500},{\"id\":2,\"name\":\"هاتف\",\"price\":1200},{\"id\":999,\"name\":\"FLAG_PLACEHOLDER\",\"price\":0}]", | |
| "flag_preview": flag1, | |
| "flag_hash": _compute_flag_hash(flag1), | |
| "hints": [ | |
| {"level": 1, "text": "الحقن موجود في معامل id. جرب إضافة علامات اقتباس بسيطة لتكسير الاستعلام.", "xp_cost": 10}, | |
| {"level": 2, "text": "استخدم UNION SELECT لتوحيد نتائج استعلامك مع الاستعلام الأصلي. جرب: 1 UNION SELECT 1,2,3,4,5,6--", "xp_cost": 20}, | |
| {"level": 3, "text": "قد تحتاج إلى تحديد عدد الأعمدة أولاً باستخدام ORDER BY. العلم موجود في أحد الصفوف المخفية.", "xp_cost": 30}, | |
| ], | |
| "difficulty": "متوسط", | |
| }) | |
| # 2. XSS | |
| flag2 = _generate_flag() | |
| seeds.append({ | |
| "title": "ثغرة XSS في منصة التعليقات", | |
| "story": "تتضمن منتديات النقاش العامة ثغرة من نوع Cross-Site Scripting حيث يتم عرض محتوى التعليقات دون تعقيم. يمكنك حقن حمولة XSS لسرقة جلسة المسؤول.", | |
| "task_outline": "أرسل تعليقاً يحتوي على حمولة XSS لتنبيه نافذة تحتوي على العَلم. استخدم <script> أو <img onerror>.", | |
| "vulnerability_type": "xss", | |
| "vulnerability_class": "Injection", | |
| "vulnerability_description": "ثغرة XSS تسمح للمهاجم بحقن نصوص برمجية في صفحات الويب التي يشاهدها مستخدمون آخرون، مما قد يؤدي إلى سرقة الجلسات أو إعادة التوجيه لمواقع ضارة.", | |
| "http_request": "POST /api/comments HTTP/1.1\nHost: forum.apex-train.com\nContent-Type: application/json\n\n{\"comment\":\"<script>alert('XSS')</script>\"}", | |
| "http_response": "HTTP/1.1 201 Created\nContent-Type: application/json\n\n{\"status\":\"ok\",\"comment\":\"<script>alert('XSS')</script>\"}", | |
| "flag_preview": flag2, | |
| "flag_hash": _compute_flag_hash(flag2), | |
| "hints": [ | |
| {"level": 1, "text": "ابحث عن حقل الإدخال الذي يرسل النص دون تنقية. جرب <script>alert(1)</script>.", "xp_cost": 10}, | |
| {"level": 2, "text": "إذا كان <script> محظوراً، استخدم <img src=x onerror=alert(1)>.", "xp_cost": 20}, | |
| {"level": 3, "text": "قد تحتاج إلى استعمال حمولة <svg onload=alert(1)> إذا تم حظر img.", "xp_cost": 30}, | |
| ], | |
| "difficulty": "سهل", | |
| }) | |
| # 3. IDOR | |
| flag3 = _generate_flag() | |
| seeds.append({ | |
| "title": "ثغرة IDOR في ملفات المستخدمين", | |
| "story": "يتيح نظام إدارة الملفات للمستخدمين تحميل واستعراض ملفاتهم عبر رابط مباشر. لاحظت أن معرّفات الملفات هي أرقام متسلسلة يمكن التكهن بها بسهولة.", | |
| "task_outline": "جرب الوصول إلى ملفات مستخدمين آخرين عبر تغيير معرف الملف في الرابط /files/1. ابحث عن الملف الذي يحتوي على العَلم.", | |
| "vulnerability_type": "idor", | |
| "vulnerability_class": "Broken Access Control", | |
| "vulnerability_description": "ثغرة IDOR (Insecure Direct Object Reference) تحدث عندما يعرض التطبيق مرجعاً مباشراً لكائن داخلي (مثل معرف ملف) دون التحقق من صلاحية الوصول.", | |
| "http_request": "GET /files/42 HTTP/1.1\nHost: cloud.apex-train.com\nCookie: session=abc123", | |
| "http_response": "HTTP/1.1 200 OK\nContent-Type: text/plain\n\nFLAG_CONTENT_HERE", | |
| "flag_preview": flag3, | |
| "flag_hash": _compute_flag_hash(flag3), | |
| "hints": [ | |
| {"level": 1, "text": "لاحظ أن معرفات الملفات هي أرقام متسلسلة (1, 2, 3...).", "xp_cost": 10}, | |
| {"level": 2, "text": "جرب الوصول للملف رقم 42 أو 99 أو 100.", "xp_cost": 20}, | |
| {"level": 3, "text": "قد تجد العلم في ملف نصي باسم flag.txt أو secret.txt.", "xp_cost": 30}, | |
| ], | |
| "difficulty": "سهل", | |
| }) | |
| # 4. SSRF | |
| flag4 = _generate_flag() | |
| seeds.append({ | |
| "title": "ثغرة SSRF في خدمة جلب الصور", | |
| "story": "تطبيق ويب يسمح لك بإدخال رابط صورة لعرضها. يقوم الخادم بجلب الصورة من الرابط نيابة عنك. يمكنك استغلال هذه الميزة لجلب موارد داخلية.", | |
| "task_outline": "استغل ثغرة Server-Side Request Forgery لجلب الملف الداخلي http://localhost:8080/admin/flag الذي يحتوي على العَلم.", | |
| "vulnerability_type": "ssrf", | |
| "vulnerability_class": "Server-Side Request Forgery", | |
| "vulnerability_description": "ثغرة SSRF تسمح للمهاجم بإجبار الخادم على إرسال طلبات إلى عناوين داخلية أو خارجية غير مقصودة، مما قد يكشف خدمات داخلية محمية.", | |
| "http_request": "POST /api/fetch-image HTTP/1.1\nHost: img.apex-train.com\nContent-Type: application/json\n\n{\"url\":\"http://localhost:8080/admin/flag\"}", | |
| "http_response": "HTTP/1.1 200 OK\nContent-Type: text/plain\n\nFLAG_CONTENT_HERE", | |
| "flag_preview": flag4, | |
| "flag_hash": _compute_flag_hash(flag4), | |
| "hints": [ | |
| {"level": 1, "text": "الخدمة تجلب الصور من أي URL تعطيه إياه. جرب إدخال http://localhost/.", "xp_cost": 10}, | |
| {"level": 2, "text": "الهدف موجود على منفذ داخلي. جرب http://localhost:8080/admin.", "xp_cost": 20}, | |
| {"level": 3, "text": "جرب http://localhost:8080/admin/flag.", "xp_cost": 30}, | |
| ], | |
| "difficulty": "متوسط", | |
| }) | |
| # 5. Command Injection (CMDi) | |
| flag5 = _generate_flag() | |
| seeds.append({ | |
| "title": "ثغرة حقن أوامر في أداة ping", | |
| "story": "توفر لوحة تحكم الشبكة أداة ping لاختبار الاتصال بالخوادم. المدخل لا يتم تنقيته بشكل صحيح، مما يسمح بتنفيذ أوامر نظام إضافية.", | |
| "task_outline": "استغل ثغرة Command Injection في حقل ping بتنفيذ أمر ls أو dir للعثور على ملف flag في النظام.", | |
| "vulnerability_type": "command-injection", | |
| "vulnerability_class": "Injection", | |
| "vulnerability_description": "ثغرة Command Injection تحدث عندما يمرر التطبيق مدخلات المستخدم مباشرة إلى shell النظام دون تنقية، مما يسمح بتنفيذ أوامر تعسفية.", | |
| "http_request": "POST /api/ping HTTP/1.1\nHost: admin.apex-train.com\nContent-Type: application/json\n\n{\"host\":\"8.8.8.8; ls -la\"}", | |
| "http_response": "HTTP/1.1 200 OK\nContent-Type: application/json\n\n{\"output\":\"PING 8.8.8.8...\\nflag.txt\\nindex.php\\n\"}", | |
| "flag_preview": flag5, | |
| "flag_hash": _compute_flag_hash(flag5), | |
| "hints": [ | |
| {"level": 1, "text": "جرب إضافة (;) أو (&&) بعد الأمر الأصلي لتنفيذ أمر إضافي.", "xp_cost": 10}, | |
| {"level": 2, "text": "جرب: 8.8.8.8; ls", "xp_cost": 20}, | |
| {"level": 3, "text": "بعد رؤية الملفات، اقرأ محتوى العلم بـ: 8.8.8.8; cat flag.txt", "xp_cost": 30}, | |
| ], | |
| "difficulty": "متوسط", | |
| }) | |
| # 6. CSRF | |
| flag6 = _generate_flag() | |
| seeds.append({ | |
| "title": "ثغرة CSRF في تغيير كلمة المرور", | |
| "story": "يسمح موقع إدارة المستخدمين بتغيير كلمة المرور عبر طلب GET دون رمز حماية CSRF. يمكنك إنشاء رابط خبيث يغير كلمة مرور المسؤول.", | |
| "task_outline": "قم بإنشاء حمولة CSRF تُغيِّر كلمة مرور المسؤول إلى 'hacked123' عبر إرسال طلب GET إلى /change-password.", | |
| "vulnerability_type": "csrf", | |
| "vulnerability_class": "Broken Access Control", | |
| "vulnerability_description": "ثغرة CSRF (Cross-Site Request Forgery) تسمح للمهاجم بإجبار ضحية على تنفيذ إجراءات غير مرغوب فيها في تطبيق ويب تمت مصادقته.", | |
| "http_request": "GET /change-password?newPass=hacked123&confirm=hacked123 HTTP/1.1\nHost: admin.apex-train.com\nCookie: session=admin_session", | |
| "http_response": "HTTP/1.1 200 OK\nContent-Type: application/json\n\n{\"status\":\"password changed successfully\"}", | |
| "flag_preview": flag6, | |
| "flag_hash": _compute_flag_hash(flag6), | |
| "hints": [ | |
| {"level": 1, "text": "التطبيق لا يستخدم رمز CSRF. أي طلب تغيير كلمة مرور يتم تنفيذه فوراً.", "xp_cost": 10}, | |
| {"level": 2, "text": "جرب إرسال طلب GET مباشرة إلى /change-password?newPass=test123&confirm=test123", "xp_cost": 20}, | |
| {"level": 3, "text": "الحمولة قد تحتاج إلى تضمينها في <img> أو <form> مخفي.", "xp_cost": 30}, | |
| ], | |
| "difficulty": "سهل", | |
| }) | |
| for s in seeds: | |
| flag = s["flag_preview"] | |
| s["flag_hash"] = _compute_flag_hash(flag) | |
| return seeds | |
| _CURATED_SEEDS = _build_seeds() | |
| # ── Pool info ───────────────────────────────────────────────────────── | |
| async def get_pool_count(team_role: str) -> int: | |
| rows = await _supabase_select( | |
| "web_exploitation_challenges", | |
| f"select=id&team_role=eq.{team_role}", | |
| ) | |
| return len(rows) | |
| # ── AI Generation ────────────────────────────────────────────────────── | |
| _WEB_EXPLOIT_PROMPT = ( | |
| "You are a senior penetration tester specializing in web application security.\n" | |
| "\n" | |
| "Your task: generate a unique and creative web exploitation challenge. " | |
| "The challenge MUST be different from any previous one.\n" | |
| "\n" | |
| "Strict requirements:\n" | |
| "- Topic: {topic}\n" | |
| "- Difficulty: {difficulty}\n" | |
| "- Write title, story, task_outline, and vulnerability_description in ARABIC only\n" | |
| "- Pick vulnerability_type from: xss, sqli, csrf, ssrf, idor, lfi, xxe, cmdi, auth, upload\n" | |
| "- Pick vulnerability_class from: Injection, Broken Access Control, ...\n" | |
| "\n" | |
| "The challenge must contain:\n" | |
| "1. **http_request**: A realistic HTTP request showing how the exploit works (with Host header, body, etc.)\n" | |
| "2. **http_response**: A realistic HTTP response from the server\n" | |
| "3. **flag_preview**: A unique flag in format CyberArena{{<hex>}} - MUST be different every time\n" | |
| "\n" | |
| "Return ONLY valid JSON with this exact shape:\n" | |
| '{{\n' | |
| ' "title": "Arabic challenge title",\n' | |
| ' "story": "Arabic scenario story",\n' | |
| ' "task_outline": "Arabic task description for the student",\n' | |
| ' "vulnerability_type": "xss",\n' | |
| ' "vulnerability_class": "Injection",\n' | |
| ' "vulnerability_description": "Arabic description of the vulnerability",\n' | |
| ' "http_request": "GET /path HTTP/1.1\\nHost: target.com\\n...",\n' | |
| ' "http_response": "HTTP/1.1 200 OK\\nContent-Type: ...\\n\\nBody",\n' | |
| ' "flag_preview": "CyberArena{{unique_hex_here}}",\n' | |
| ' "difficulty": "{difficulty}",\n' | |
| ' "hints": [\n' | |
| ' {{"level":1,"text":"First hint in Arabic","xp_cost":10}},\n' | |
| ' {{"level":2,"text":"Second hint in Arabic","xp_cost":20}},\n' | |
| ' {{"level":3,"text":"Third hint in Arabic","xp_cost":30}}\n' | |
| " ]\n" | |
| "}}\n" | |
| "\n" | |
| "IMPORTANT: Do NOT repeat any previously generated challenge. " | |
| "Every call must produce a fresh, unique scenario with a different vulnerability type, " | |
| "different HTTP traffic, and a different context." | |
| ) | |
| async def _call_cloudflare(prompt: str, timeout: int = 45) -> Optional[str]: | |
| """Try Cloudflare models in order, cycling through fallbacks.""" | |
| if not CLOUDFLARE_API_TOKEN or not CLOUDFLARE_ACCOUNT_ID: | |
| return None | |
| for model in CLOUDFLARE_MODEL_FALLBACKS: | |
| url = f"https://api.cloudflare.com/client/v4/accounts/{CLOUDFLARE_ACCOUNT_ID}/ai/run/{model}" | |
| payload = { | |
| "messages": [ | |
| {"role": "system", "content": "أنت خبير أمن سيبراني. أعد JSON فقط."}, | |
| {"role": "user", "content": prompt}, | |
| ], | |
| } | |
| try: | |
| async with httpx.AsyncClient(timeout=timeout) as client: | |
| resp = await client.post( | |
| url, | |
| json={**payload, "response_format": {"type": "json_object"}}, | |
| headers={"Authorization": f"Bearer {CLOUDFLARE_API_TOKEN}"}, | |
| ) | |
| if resp.status_code == 200: | |
| data = resp.json() | |
| if data.get("success") and data.get("result"): | |
| return data["result"]["response"] | |
| if resp.status_code in (400, 404, 422): | |
| # Retry without response_format | |
| resp2 = await client.post( | |
| url, | |
| json=payload, | |
| headers={"Authorization": f"Bearer {CLOUDFLARE_API_TOKEN}"}, | |
| ) | |
| if resp2.status_code == 200: | |
| data = resp2.json() | |
| if data.get("success") and data.get("result"): | |
| return data["result"]["response"] | |
| if resp.status_code == 429: | |
| return None # Hard backoff | |
| except Exception as e: | |
| logging.warning(f"[web_exploit] CF model {model} error: {e}") | |
| continue | |
| return None | |
| async def _call_groq(prompt: str, timeout: int = 60) -> Optional[str]: | |
| if not GROQ_API_KEY: | |
| return None | |
| try: | |
| async with httpx.AsyncClient(timeout=timeout) as client: | |
| resp = await client.post( | |
| GROQ_API_URL, | |
| json={ | |
| "model": "llama-3.3-70b-versatile", | |
| "messages": [ | |
| {"role": "system", "content": "أنت خبير أمن سيبراني. أعد JSON فقط."}, | |
| {"role": "user", "content": prompt}, | |
| ], | |
| "temperature": 0.8, | |
| "max_tokens": 2048, | |
| "response_format": {"type": "json_object"}, | |
| }, | |
| headers={ | |
| "Content-Type": "application/json", | |
| "Authorization": f"Bearer {GROQ_API_KEY}", | |
| }, | |
| ) | |
| if resp.status_code == 200: | |
| return resp.json()["choices"][0]["message"]["content"] | |
| except Exception as e: | |
| logging.warning(f"[web_exploit] Groq error: {e}") | |
| return None | |
| async def _call_nvidia(prompt: str, timeout: int = 60) -> Optional[str]: | |
| if not NVIDIA_API_KEY: | |
| return None | |
| try: | |
| async with httpx.AsyncClient(timeout=timeout) as client: | |
| resp = await client.post( | |
| "https://integrate.api.nvidia.com/v1/chat/completions", | |
| json={ | |
| "model": NVIDIA_MODEL or "deepseek-ai/deepseek-v4-pro", | |
| "messages": [ | |
| {"role": "system", "content": "أنت خبير أمن سيبراني. أعد JSON فقط."}, | |
| {"role": "user", "content": prompt}, | |
| ], | |
| "temperature": 0.8, | |
| "max_tokens": 2048, | |
| "response_format": {"type": "json_object"}, | |
| }, | |
| headers={ | |
| "Content-Type": "application/json", | |
| "Authorization": f"Bearer {NVIDIA_API_KEY}", | |
| }, | |
| ) | |
| if resp.status_code == 200: | |
| return resp.json()["choices"][0]["message"]["content"] | |
| except Exception as e: | |
| logging.warning(f"[web_exploit] NVIDIA error: {e}") | |
| return None | |
| async def _call_mistral(prompt: str) -> Optional[str]: | |
| if not MISTRAL_API_KEY: | |
| print("[mistral] no API key set, skipping") | |
| return None | |
| headers = { | |
| "Content-Type": "application/json", | |
| "Authorization": f"Bearer {MISTRAL_API_KEY}", | |
| } | |
| payload = { | |
| "model": MISTRAL_MODEL, | |
| "messages": [{"role": "user", "content": prompt}], | |
| "max_tokens": 2048, | |
| "temperature": 0.7, | |
| } | |
| print(f"[mistral] calling {MISTRAL_MODEL}...") | |
| try: | |
| async with httpx.AsyncClient(timeout=60) as client: | |
| resp = await client.post(MISTRAL_API_URL, json=payload, headers=headers) | |
| if resp.status_code == 200: | |
| text = resp.json()["choices"][0]["message"]["content"] | |
| print("[mistral] success") | |
| return text | |
| print(f"[mistral] HTTP {resp.status_code}: {resp.text[:200]}") | |
| return None | |
| except Exception as e: | |
| print(f"[mistral] error: {e}") | |
| return None | |
| async def _generate_one_challenge(team_role: str, topic: str, difficulty: str) -> Optional[WebExploitChallenge]: | |
| """Try AI providers in order, fall back to curated seed.""" | |
| prompt = _WEB_EXPLOIT_PROMPT.format(topic=topic, difficulty=difficulty) | |
| # Tier 1: Mistral | |
| content = await _call_mistral(prompt) | |
| # Tier 2: Cloudflare | |
| if not content: | |
| print("[gen] Mistral failed, trying Cloudflare...") | |
| content = await _call_cloudflare(prompt) | |
| # Tier 3: Groq | |
| if not content: | |
| print("[gen] Cloudflare failed, trying Groq...") | |
| content = await _call_groq(prompt) | |
| # Tier 4: NVIDIA | |
| if not content: | |
| print("[gen] Groq failed, trying NVIDIA...") | |
| content = await _call_nvidia(prompt) | |
| if content: | |
| data = _extract_json(content) | |
| if data: | |
| return _dict_to_challenge(data, topic) | |
| print("[gen] extracted JSON is None, falling back to seed") | |
| # Tier 5: Curated seed | |
| print("[gen] all AI tiers exhausted, using curated seed") | |
| seed = random.choice(_CURATED_SEEDS) | |
| seed_copy = dict(seed) | |
| flag = _generate_flag() | |
| seed_copy["flag_preview"] = flag | |
| seed_copy["flag_hash"] = _compute_flag_hash(flag) | |
| seed_copy["topic"] = topic | |
| seed_copy["team_role"] = team_role | |
| return WebExploitChallenge(**seed_copy) | |
| def _dict_to_challenge(data: dict, topic: str) -> Optional[WebExploitChallenge]: | |
| try: | |
| hints = data.get("hints", []) | |
| if not isinstance(hints, list): | |
| hints = [] | |
| flag = data.get("flag_preview", _generate_flag()) | |
| return WebExploitChallenge( | |
| title=str(data.get("title", ""))[:200], | |
| story=str(data.get("story", ""))[:1000], | |
| task_outline=str(data.get("task_outline", ""))[:1000], | |
| vulnerability_type=str(data.get("vulnerability_type", "xss"))[:50], | |
| vulnerability_class=str(data.get("vulnerability_class", "Injection"))[:50], | |
| vulnerability_description=str(data.get("vulnerability_description", ""))[:1000], | |
| http_request=str(data.get("http_request", ""))[:2000], | |
| http_response=str(data.get("http_response", ""))[:2000], | |
| flag_preview=flag, | |
| flag_hash=_compute_flag_hash(flag), | |
| hints=hints[:3], | |
| difficulty=str(data.get("difficulty", "متوسط"))[:20], | |
| topic=topic, | |
| ) | |
| except Exception as e: | |
| logging.error(f"[web_exploit] dict_to_challenge error: {e}") | |
| return None | |
| # ── DB Insert ───────────────────────────────────────────────────────── | |
| async def _insert_to_db(challenge: WebExploitChallenge) -> Optional[dict]: | |
| from app.services.insert_guard import atomic_insert | |
| from app.services.dedup import is_duplicate_web_exploit | |
| row = { | |
| "team_role": "red", | |
| "module": "web-exploitation", | |
| "topic": challenge.topic, | |
| "title": challenge.title, | |
| "story": challenge.story, | |
| "task_outline": challenge.task_outline, | |
| "vulnerability_type": challenge.vulnerability_type, | |
| "vulnerability_class": challenge.vulnerability_class, | |
| "vulnerability_description": challenge.vulnerability_description, | |
| "http_request": challenge.http_request, | |
| "http_response": challenge.http_response, | |
| "flag_preview": challenge.flag_preview, | |
| "flag_hash": challenge.flag_hash, | |
| "hints": json.dumps(challenge.hints, ensure_ascii=False), | |
| "difficulty": challenge.difficulty, | |
| "xp_reward": challenge.xp_reward, | |
| } | |
| ok = await atomic_insert( | |
| table="web_exploitation_challenges", | |
| team_role="red", | |
| row=row, | |
| dedup_func=is_duplicate_web_exploit, | |
| dedup_args=[ | |
| row.get("title", ""), | |
| row.get("http_request", ""), | |
| row.get("task_outline", ""), | |
| ], | |
| dedup_kwargs={"role_filter": "red"}, | |
| ) | |
| return row if ok else None | |
| # ── Pool refill ─────────────────────────────────────────────────────── | |
| async def _refill_pool_inner(team_role: str, count: int) -> int: | |
| """Unlocked inner — caller must hold _get_pool_lock(team_role).""" | |
| now = time.time() | |
| backoff = _AI_BACKOFF_UNTIL.get(team_role, 0) | |
| if now < backoff: | |
| logging.info(f"[web_exploit] backoff active for {team_role}, skipping") | |
| return 0 | |
| inserted = 0 | |
| idx = _TOPIC_INDEX.get(team_role, 0) | |
| difficulties = ["سهل", "سهل", "متوسط", "متوسط", "صعب"] | |
| for i in range(count): | |
| topic = WEB_EXPLOIT_TOPICS[idx % len(WEB_EXPLOIT_TOPICS)] | |
| idx += 1 | |
| difficulty = random.choice(difficulties) | |
| challenge = await _generate_one_challenge(team_role, topic, difficulty) | |
| if challenge: | |
| result = await _insert_to_db(challenge) | |
| if result: | |
| inserted += 1 | |
| await asyncio.sleep(1) | |
| if i < count - 1: | |
| await asyncio.sleep(0.5) | |
| _TOPIC_INDEX[team_role] = idx | |
| return inserted | |
| async def refill_pool(team_role: str, count: int) -> int: | |
| async with _get_pool_lock(team_role): | |
| return await _refill_pool_inner(team_role, count) | |
| # ── Pool watcher ────────────────────────────────────────────────────── | |
| _WATCHER_STARTED: set[tuple[str, str]] = set() | |
| async def start_pool_watcher(team_role: str) -> None: | |
| key = ("web-exploitation", team_role) | |
| if key in _WATCHER_STARTED: | |
| print(f"[web-exploit:{team_role}] watcher already running — skipping duplicate start.") | |
| return | |
| _WATCHER_STARTED.add(key) | |
| try: | |
| label = f"[web-exploit:{team_role}]" | |
| print(f"{label} watcher started (target={POOL_TARGET}).") | |
| while True: | |
| try: | |
| sleep_secs = 30 | |
| async with _get_pool_lock(team_role): | |
| count = await get_pool_count(team_role) | |
| if count < POOL_TARGET: | |
| needed = POOL_TARGET - count | |
| print(f"{label} pool below target ({count}/{POOL_TARGET}) — refilling {needed}…") | |
| added = await _refill_pool_inner(team_role, needed) | |
| new_count = count + added | |
| print(f"{label} refilled: {count} → {new_count} (target {POOL_TARGET}).") | |
| sleep_secs = 5 | |
| await asyncio.sleep(sleep_secs) | |
| except Exception as e: | |
| import traceback | |
| print(f"{label} watcher error: {e}") | |
| traceback.print_exc() | |
| await asyncio.sleep(10) | |
| finally: | |
| _WATCHER_STARTED.discard(key) | |
| # ── CLI entry (for testing) ─────────────────────────────────────────── | |
| if __name__ == "__main__": | |
| import sys | |
| async def _main(): | |
| team_role = sys.argv[2] if len(sys.argv) > 2 else "red" | |
| if "--refill" in sys.argv: | |
| n = int(sys.argv[sys.argv.index("--refill") + 1]) if "--refill" in sys.argv and len(sys.argv) > sys.argv.index("--refill") + 1 else 3 | |
| inserted = await refill_pool(team_role, n) | |
| print(f"Inserted {inserted} challenges for {team_role}") | |
| elif "--seed-only" in sys.argv: | |
| for s in _CURATED_SEEDS: | |
| c = WebExploitChallenge( | |
| title=s["title"], story=s["story"], task_outline=s["task_outline"], | |
| vulnerability_type=s["vulnerability_type"], vulnerability_class=s["vulnerability_class"], | |
| vulnerability_description=s["vulnerability_description"], | |
| http_request=s["http_request"], http_response=s["http_response"], | |
| flag_preview=s["flag_preview"], flag_hash=s["flag_hash"], | |
| hints=s["hints"], difficulty=s["difficulty"], topic=s.get("vulnerability_type", "xss"), | |
| ) | |
| r = await _insert_to_db(c) | |
| if r: | |
| print(f"Inserted seed: {c.title}") | |
| else: | |
| print(f"Skipped (dup): {c.title}") | |
| else: | |
| print("Usage: python -m app.generators.web_exploitation --refill N [--team red]") | |
| print(" python -m app.generators.web_exploitation --seed-only [--team red]") | |
| asyncio.run(_main()) | |