""" code_fixing_generator.py ======================== Generates fully-realized code-fixing challenges for the CyberArena Blue Team pool. Architecture: main.py watcher (polls code_fixing_challenges count) | v refill_pool(team_role, count) --> ai_generate_challenge() | | | v | AI generates vulnerable code | | | v | ChallengeBuilder produces DB row v insert_to_db() <----- Challenge row | v public.code_fixing_challenges Public API (used by main.py): - POOL_TARGET = 5 - POOL_THRESHOLD = 2 - POOL_BATCH = 3 - get_pool_count(team_role) -> int - async refill_pool(team_role, count) -> int - async start_pool_watcher(team_role) CLI: python code_fixing_generator.py --team blue --seed-only python code_fixing_generator.py --team blue --ai --count 1 """ from __future__ import annotations import argparse import asyncio import hashlib import json import os import random import re import sys import time from dataclasses import dataclass, field from typing import Optional, Literal # Load .env early so the CLI works without manual export try: from dotenv import load_dotenv _env_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".env") if os.path.exists(_env_path): load_dotenv(_env_path) except ImportError: pass import httpx # --------------------------------------------------------------------------- # # 0. Constants # --------------------------------------------------------------------------- # ALLOWED_LANGUAGES = ("C++", "JAVA", "PYTHON", "JAVASCRIPT", "PHP", "RUST") ALLOWED_DIFFICULTIES = ("مبتدئ", "متوسط", "قوي") ALLOWED_TEAMS = ("blue",) # Vulnerability types per language VULN_TYPES_BY_LANGUAGE = { "C++": { "buffer-overflow": "Hijacking stack-based buffer overflow", "use-after-free": "Use-After-Free dangling pointer", "format-string": "Format string vulnerability", "integer-overflow": "Integer overflow leading to buffer size miscalculation", }, "JAVA": { "sql-injection": "SQL Injection via string concatenation", "path-traversal": "Path Traversal via unsanitized user input", "unsafe-deserialization": "Unsafe Deserialization of untrusted data", "xss": "Cross-Site Scripting via unsanitized output", }, "PYTHON": { "sql-injection": "SQL Injection via f-string / format", "command-injection": "Command Injection via os.system / subprocess", "pickle-deserialization": "Pickle Deserialization of untrusted data", "path-traversal": "Path Traversal via open() with user input", }, "JAVASCRIPT": { "xss": "Cross-Site Scripting via innerHTML / eval", "prototype-pollution": "Prototype Pollution via merge/extend", "redos": "Regular Expression Denial of Service (ReDoS)", "path-traversal": "Path Traversal in file system operations", }, "PHP": { "sql-injection": "SQL Injection via mysql_query / string concat", "file-inclusion": "Local File Inclusion via include/require", "command-injection": "Command Injection via exec / system / shell_exec", "type-juggling": "Type Juggling loose comparison bypass", }, "RUST": { "unsafe-block": "Unsafe block bypassing borrow checker safety", "unwrap-panic": "Unwrap panic on None/Err causing DoS", "integer-overflow": "Integer overflow in release mode", }, } # Module names for each vulnerability MODULE_BY_VULN = { "buffer-overflow": "systems-security", "use-after-free": "systems-security", "format-string": "systems-security", "integer-overflow": "systems-security", "sql-injection": "web-security", "path-traversal": "web-security", "unsafe-deserialization": "web-security", "xss": "web-security", "pickle-deserialization": "web-security", "command-injection": "web-security", "prototype-pollution": "web-security", "redos": "web-security", "file-inclusion": "web-security", "type-juggling": "web-security", "unsafe-block": "systems-security", "unwrap-panic": "systems-security", } # Rotation: cycles through (language, vuln_type) pairs VULN_ROTATION = [] for lang, vulns in VULN_TYPES_BY_LANGUAGE.items(): for vuln_key in vulns: VULN_ROTATION.append((lang, vuln_key)) # Pool constants POOL_TARGET = 5 POOL_THRESHOLD = 2 POOL_BATCH = 3 from app.core.config import ( SUPABASE_URL, SUPABASE_ANON_KEY, CLOUDFLARE_API_TOKEN, CLOUDFLARE_ACCOUNT_ID, CLOUDFLARE_MODEL, CLOUDFLARE_URL, GROQ_API_KEY, GROQ_MODEL, GROQ_API_URL, NVIDIA_API_KEY, NVIDIA_MODEL, NVIDIA_URL, MISTRAL_API_KEY, MISTRAL_MODEL, MISTRAL_API_URL, ) 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", ] TABLE_NAME = "code_fixing_challenges" # Per-team backoff tracker _AI_BACKOFF_UNTIL: dict[str, float] = {} # Concurrency control — one asyncio.Lock per team serialises pool operations _POOL_LOCKS: dict[str, asyncio.Lock] = {} _WATCHER_STARTED: set[str] = set() # --------------------------------------------------------------------------- # # 1. Helpers # --------------------------------------------------------------------------- # 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] def supabase_headers(content_type: bool = False) -> dict: headers = { "apikey": SUPABASE_ANON_KEY, "Authorization": f"Bearer {SUPABASE_ANON_KEY}", } if content_type: headers["Content-Type"] = "application/json" return headers def _extract_string_value(text: str, key: str) -> Optional[str]: """Extract a JSON string value for a given key, handling multi-line content and embedded quotes.""" # Match "key": "value" with proper string boundary detection # Value can contain escaped quotes (\"), newlines (\n), tabs, etc. pattern = rf'"{re.escape(key)}"\s*:\s*"((?:[^"\\]|\\.)*)"' match = re.search(pattern, text, re.DOTALL) if match: value = match.group(1) # Unescape common sequences value = value.replace('\\n', '\n').replace('\\t', '\t').replace('\\"', '"').replace('\\\\', '\\') return value return None def _find_matching_bracket(text: str, open_idx: int, open_ch: str, close_ch: str) -> Optional[int]: """Find the matching closing bracket, accounting for nested brackets and strings.""" depth = 0 i = open_idx in_string = False escape_next = False while i < len(text): c = text[i] if escape_next: escape_next = False i += 1 continue if c == '\\' and in_string: escape_next = True i += 1 continue if c == '"': in_string = not in_string elif not in_string: if c == open_ch: depth += 1 elif c == close_ch: depth -= 1 if depth == 0: return i i += 1 return None def parse_json_safe(raw) -> dict: """Robustly parse LLM JSON output, tolerating markdown fences, trailing commas, etc. Accepts both string and dict — if a dict is passed (some providers auto-parse), returns it directly after validation. """ if raw is None or raw == "": raise ValueError("Empty response from model") if isinstance(raw, dict): return raw if not isinstance(raw, str): raise ValueError(f"Expected str or dict, got {type(raw).__name__}") cleaned = raw.strip() cleaned = raw.strip() # 1) Strip markdown code fences fence_match = re.search(r"```(?:json)?\s*([\s\S]*?)```", cleaned, re.IGNORECASE) if fence_match: cleaned = fence_match.group(1).strip() # 2) Locate outermost JSON object start = cleaned.find("{") if start == -1: raise ValueError(f"No JSON object found. Raw: {cleaned[:200]}") end = _find_matching_bracket(cleaned, start, "{", "}") if end is None: # Fallback: take last } end = cleaned.rfind("}") if end == -1 or end <= start: raise ValueError(f"Unbalanced JSON braces. Raw: {cleaned[:200]}") cleaned = cleaned[start:end + 1] # 3) Remove control characters cleaned = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F\x7F]", "", cleaned) # 4) Try strict parse first try: return json.loads(cleaned, strict=False) except json.JSONDecodeError: pass # 5) Fix common LLM JSON bugs and retry fixed = cleaned # Remove trailing commas fixed = re.sub(r',\s*([}\]])', r'\1', fixed) # Fix unquoted keys (word followed by :) fixed = re.sub(r'(? tuple[str, str]: return VULN_ROTATION[i % len(VULN_ROTATION)] # --------------------------------------------------------------------------- # # 2. Curated Seeds # --------------------------------------------------------------------------- # def _build_seeds() -> list[dict]: """Return a list of curated seed challenges for the pool.""" seeds = [ # --- PYTHON --- { "language": "PYTHON", "vulnerability_type": "sql-injection", "module": "web-security", "title": "إصلاح ثغرة SQL Injection في كود Python", "story": "لديك تطبيق ويب يستخدم قاعدة بيانات SQLite لتخزين بيانات المستخدمين. المطور استخدم تنسيق النصوص مباشرة في الاستعلام بدون parameterized queries، مما يسمح بحقن SQL خبيث.", "task_outline": "راجع الكود المصدري أدناه وحدد ثغرة SQL Injection. صحّح الكود باستخدام parameterized queries (مثلاً cursor.execute('SELECT * FROM users WHERE username = ?', (username,))) بدلاً من f-string أو format().", "vulnerable_code": """import sqlite3 def get_user(username): conn = sqlite3.connect('app.db') cursor = conn.cursor() # !! SQL Injection vulnerability: user input is directly interpolated query = f"SELECT * FROM users WHERE username = '{username}'" cursor.execute(query) result = cursor.fetchone() conn.close() return result # Example usage user_input = input("Enter username: ") user = get_user(user_input) print(user)""", "vulnerability_type": "sql-injection", "vulnerability_description": "The query is built using f-string interpolation, allowing an attacker to inject SQL like `' OR '1'='1' --` to bypass authentication or dump the database.", "difficulty": "مبتدئ", "xp_reward": 150, "hints": [ {"level": 1, "text": "استخدم cursor.execute() مع parameterized query بدلاً من f-string", "xp_cost": 20}, {"level": 2, "text": "النمط الصحيح: cursor.execute('SELECT ... WHERE col = ?', (value,))", "xp_cost": 40}, ], }, { "language": "PYTHON", "vulnerability_type": "command-injection", "module": "web-security", "title": "إصلاح ثغرة Command Injection في Python", "story": "أداة إدارية تقوم بعمل ping لعنوان IP يدخله المستخدم لفحص الاتصال. المطور استخدم os.system() مع دمج النص مباشرة.", "task_outline": "اكتب كوداً آمناً يستخدم subprocess.run() بدلاً من os.system()، ويتحقق من أن المدخل IP صالح (أرقام ونقاط فقط) قبل التنفيذ.", "vulnerable_code": """import os def ping_host(ip_address): # !! Command Injection: unsanitized input passed to shell command = f"ping -c 4 {ip_address}" os.system(command) # User input target = input("Enter IP to ping: ") ping_host(target)""", "vulnerability_type": "command-injection", "vulnerability_description": "os.system() executes the command via the shell. An attacker can inject `; rm -rf /` or `$(malicious_command)` to execute arbitrary commands.", "difficulty": "متوسط", "xp_reward": 200, "hints": [ {"level": 1, "text": "استخدم subprocess.run() مع shell=False", "xp_cost": 20}, {"level": 2, "text": "تحقق من صحة المدخل: re.match(r'^[0-9.]+$', ip)", "xp_cost": 40}, ], }, { "language": "PYTHON", "vulnerability_type": "pickle-deserialization", "module": "web-security", "title": "إصلاح ثغرة Pickle Deserialization", "story": "تطبيق يُخزّن جلسات المستخدمين في ملفات pickle. عند قراءة الجلسة، يتم استخدام pickle.loads() مباشرة على البيانات المحملة من ملف غير موثوق.", "task_outline": "استبدل pickle.loads() بآمن JSON أو استخدم hashlib للتحقق من سلامة البيانات قبل التحميل.", "vulnerable_code": """import pickle import os def load_session(session_file): with open(session_file, 'rb') as f: # !! Unsafe deserialization: pickle.loads on untrusted data data = pickle.loads(f.read()) return data def save_session(session_file, data): with open(session_file, 'wb') as f: pickle.dump(data, f) # Load session session = load_session('user_session.pkl') print(session)""", "vulnerability_type": "pickle-deserialization", "vulnerability_description": "pickle.loads() can execute arbitrary code during deserialization. An attacker can craft a malicious pickle file that runs OS commands when loaded.", "difficulty": "قوي", "xp_reward": 250, "hints": [ {"level": 1, "text": "استبدل pickle بـ json.dumps/loads لتخزين آمن", "xp_cost": 20}, {"level": 2, "text": "إذا كنت تحتاج pickle، تحقق من HMAC قبل التحميل", "xp_cost": 40}, ], }, # --- JAVA --- { "language": "JAVA", "vulnerability_type": "sql-injection", "module": "web-security", "title": "إصلاح ثغرة SQL Injection في Java", "story": "تطبيق ويب جافا يستخدم JDBC للاتصال بقاعدة البيانات. الاستعلام يُبنى بـ StringBuilder مع محاذاة النص مباشرة.", "task_outline": "استبدل String concatenation بـ PreparedStatement مع parameterized query.", "vulnerable_code": """import java.sql.*; public class UserManager { public static User findUser(String username, String password) { Connection conn = null; try { conn = DriverManager.getConnection("jdbc:mysql://localhost/app", "root", "pass"); Statement stmt = conn.createStatement(); // !! SQL Injection: user input directly concatenated into query String query = "SELECT * FROM users WHERE username = '" + username + "' AND password = '" + password + "'"; ResultSet rs = stmt.executeQuery(query); if (rs.next()) { return new User(rs.getString("username"), rs.getString("email")); } } catch (Exception e) { e.printStackTrace(); } return null; } }""", "vulnerability_type": "sql-injection", "vulnerability_description": "String concatenation in SQL query allows injection. An attacker can bypass authentication with `' OR '1'='1'` as password.", "difficulty": "متوسط", "xp_reward": 200, "hints": [ {"level": 1, "text": "استخدم PreparedStatement بدلاً من Statement", "xp_cost": 20}, {"level": 2, "text": "النمط: PreparedStatement ps = conn.prepareStatement('SELECT ... WHERE username = ? AND password = ?')", "xp_cost": 40}, ], }, { "language": "JAVA", "vulnerability_type": "path-traversal", "module": "web-security", "title": "إصلاح ثغرة Path Traversal في Java", "story": "خادم ويب يسمح بتحميل ملفات من مجلد معين باستخدام اسم الملف من الطلب.", "task_outline": "أضف تحقق من أن المسار الناتج لا يتجاوز المجلد المحدد باستخدام Path.normalize() و startsWith().", "vulnerable_code": """import java.io.*; import java.nio.file.*; public class FileServer { public static byte[] getFile(String filename) throws IOException { // !! Path Traversal: no validation on filename String basePath = "/var/www/uploads/"; String fullPath = basePath + filename; return Files.readAllBytes(Paths.get(fullPath)); } // Example: getFile("../../../etc/passwd") reads system files }""", "vulnerability_type": "path-traversal", "vulnerability_description": "No validation on filename allows `../../../etc/passwd` to read system files outside the uploads directory.", "difficulty": "مبتدئ", "xp_reward": 150, "hints": [ {"level": 1, "text": "استخدم Paths.get(basePath, filename).normalize()", "xp_cost": 20}, {"level": 2, "text": "تحقق: resolved.startsWith(Paths.get(basePath))", "xp_cost": 40}, ], }, # --- C++ --- { "language": "C++", "vulnerability_type": "buffer-overflow", "module": "systems-security", "title": "إصلاح ثغرة Buffer Overflow في C++", "story": "دالة تقرأ مدخلات المستخدم إلى مخزن ثابت الحجم بدون التحقق من الطول.", "task_outline": "استبدل المخزن الثابت بـ std::string أو أضف تحقق من طول المدخلات قبل النسخ.", "vulnerable_code": """#include #include void process_input() { char buffer[64]; // Fixed-size buffer std::cout << "Enter your name: "; // !! Buffer Overflow: no bounds checking on input std::cin.getline(buffer, 256); // Can write up to 256 bytes into 64-byte buffer std::cout << "Hello, " << buffer << std::endl; } int main() { process_input(); return 0; }""", "vulnerability_type": "buffer-overflow", "vulnerability_description": "std::cin.getline with limit 256 but buffer is only 64 bytes. Input longer than 63 chars overflows the buffer, potentially overwriting return address.", "difficulty": "متوسط", "xp_reward": 200, "hints": [ {"level": 1, "text": "استخدم std::string بدلاً من char[]", "xp_cost": 20}, {"level": 2, "text": "أو عدّل الحد الأقصى: std::cin.getline(buffer, sizeof(buffer))", "xp_cost": 40}, ], }, { "language": "C++", "vulnerability_type": "format-string", "module": "systems-security", "title": "إصلاح ثغرة Format String في C++", "story": "دالة تسجل رسائل باستخدام printf مع نص من المستخدم مباشرة.", "task_outline": "غيّر printf(user_input) إلى printf('%s', user_input) لمنع حقن Formatters.", "vulnerable_code": """#include #include void log_message(char *user_msg) { // !! Format String vulnerability: user input used as format string printf(user_msg); } int main() { char input[256]; std::cout << "Enter log message: "; std::cin.getline(input, sizeof(input)); log_message(input); return 0; }""", "vulnerability_type": "format-string", "vulnerability_description": "printf with user-controlled format string allows `%x`, `%n` to read/write stack memory. Attacker can leak stack data or overwrite memory.", "difficulty": "قوي", "xp_reward": 250, "hints": [ {"level": 1, "text": "غيّر printf(user_msg) إلى printf('%s', user_msg)", "xp_cost": 20}, {"level": 2, "text": "أو استخدم std::cout بدلاً من printf", "xp_cost": 40}, ], }, # --- JAVASCRIPT --- { "language": "JAVASCRIPT", "vulnerability_type": "xss", "module": "web-security", "title": "إصلاح ثغرة XSS في JavaScript", "story": "تطبيق ويب يعرض تعليقات المستخدمين باستخدام innerHTML مباشرة.", "task_outline": "استبدل innerHTML بـ textContent لمنع تنفيذ JavaScript في التعليقات.", "vulnerable_code": """function displayComment(comment) { const container = document.getElementById('comments'); // !! XSS vulnerability: user input rendered as HTML const div = document.createElement('div'); div.innerHTML = comment; // executes! container.appendChild(div); } // Example: comment = '' displayComment(userComment);""", "vulnerability_type": "xss", "vulnerability_description": "innerHTML parses HTML tags. An attacker can inject `