samarthu78/s-python-compiler
0
1from fastapi import FastAPI, WebSocket, WebSocketDisconnect2from fastapi.middleware.cors import CORSMiddleware3from contextlib import asynccontextmanager4import asyncio5import os6import sys7import uuid8import time9 10# ==========================================11# CONFIGURATION12# ==========================================13MAX_CONCURRENT_SCRIPTS = 4 14IDLE_TIMEOUT_SECONDS = 30.0 # Kill process if completely idle for 30s15 16# Global State17execution_semaphore = None 18active_sessions = {} # { session_id: { "process": proc, "last_active": float, "websocket": ws, "temp_path": str, "holds_semaphore": bool } }19 20# Anti-Spam Tracking dictionaries21ip_cooldowns = {} # { "ip_address": float(timestamp) }22ip_active_sessions = {} # { "ip_address": "session_id" }23 24# ==========================================25# PATCHES (Injected Code)26# ==========================================27 28PATCH_LIMITS = """29import resource30def _restrict_resources():31 # 1. Limit RAM to 512 MB32 max_mem = 512 * 1024 * 102433 try:34 resource.setrlimit(resource.RLIMIT_AS, (max_mem, max_mem))35 except ValueError:36 pass37 38 # 2. Limit actual CPU execution time to 3 seconds39 try:40 resource.setrlimit(resource.RLIMIT_CPU, (3, 3))41 except ValueError:42 pass43_restrict_resources()44"""45 46PATCH_HEADER = """47import sys48import base6449def _print_image_tag(image_bytes, format='png'):50 try:51 img_b64 = base64.b64encode(image_bytes).decode('utf-8')52 sys.stdout.write(f"\\n@@@IMAGE_START@@@{img_b64}@@@IMAGE_END@@@\\n")53 sys.stdout.flush()54 except Exception:55 pass56"""57 58PATCH_MATPLOTLIB = """59import io60try:61 import matplotlib62 matplotlib.use('Agg')63 import matplotlib.pyplot as plt64 def _plt_show_patch(*args, **kwargs):65 buf = io.BytesIO()66 try:67 plt.savefig(buf, format='png')68 buf.seek(0)69 _print_image_tag(buf.read(), 'png')70 finally:71 plt.clf()72 buf.close()73 plt.show = _plt_show_patch74except ImportError:75 pass76"""77 78PATCH_PIL = """79import io80try:81 from PIL import Image82 def _pil_show_patch(self, title=None, command=None):83 buf = io.BytesIO()84 self.save(buf, format='PNG') 85 _print_image_tag(buf.getvalue(), 'png')86 Image.Image.show = _pil_show_patch87except ImportError:88 pass89"""90 91PATCH_INPUT = """92import builtins93import sys94def _custom_input(prompt=""):95 if prompt:96 sys.stdout.write(str(prompt))97 sys.stdout.write("@@@WAITING_FOR_INPUT@@@")98 sys.stdout.flush()99 ans = sys.stdin.readline()100 return ans.rstrip('\\r\\n')101builtins.input = _custom_input102"""103 104# ==========================================105# BACKGROUND JANITOR (The 30-Second Kill Switch & Memory Sweep)106# ==========================================107async def session_janitor():108 global execution_semaphore109 while True:110 await asyncio.sleep(5)111 current_time = time.time()112 113 # 🟢 NEW FIX: Sweep old IP cooldowns to prevent memory leaks!114 for ip in list(ip_cooldowns.keys()):115 # If the user hasn't made a request in the last 2 seconds, they are clear.116 if current_time - ip_cooldowns.get(ip, 0) > 2.0:117 del ip_cooldowns[ip]118 119 for session_id in list(active_sessions.keys()):120 session = active_sessions.get(session_id)121 if not session:122 continue123 124 if current_time - session["last_active"] > IDLE_TIMEOUT_SECONDS:125 # 1. Kill Process126 proc = session.get("process")127 if proc and proc.returncode is None:128 try: proc.kill()129 except ProcessLookupError: pass130 131 # 2. Drop WebSocket132 ws = session.get("websocket")133 if ws:134 try: await ws.close(code=1008, reason="Idle timeout exceeded.")135 except: pass136 137 # 3. Cleanup File138 temp_path = session.get("temp_path")139 if temp_path and os.path.exists(temp_path):140 try: os.remove(temp_path)141 except: pass142 143 # 4. Release Locks & Memory144 if session.get("holds_semaphore"):145 execution_semaphore.release()146 147 del active_sessions[session_id]148 149# ==========================================150# LIFESPAN 151# ==========================================152@asynccontextmanager153async def lifespan(app: FastAPI):154 global execution_semaphore155 execution_semaphore = asyncio.Semaphore(MAX_CONCURRENT_SCRIPTS)156 janitor_task = asyncio.create_task(session_janitor())157 yield158 janitor_task.cancel()159 160app = FastAPI(lifespan=lifespan)161 162app.add_middleware(163 CORSMiddleware,164 allow_origins=["*"],165 allow_methods=["*"],166 allow_headers=["*"],167)168 169@app.get("/")170async def home():171 return {"status": "active", "active_websockets": len(active_sessions)}172 173# ==========================================174# WEBSOCKET ENDPOINT175# ==========================================176@app.websocket("/ws")177async def websocket_endpoint(websocket: WebSocket):178 await websocket.accept()179 180 # Extract real IP (protects against Cloudflare/Nginx proxies returning local IPs)181 raw_ip = websocket.headers.get("x-forwarded-for", websocket.headers.get("x-real-ip", websocket.client.host))182 client_ip = raw_ip.split(",")[0].strip() if isinstance(raw_ip, str) else raw_ip183 184 # --- 1. SPAM COOLDOWN (3 Seconds) ---185 current_time = time.time()186 if current_time - ip_cooldowns.get(client_ip, 0) < 2.0:187 await websocket.send_json({"type": "error", "data": "Please wait 2 seconds before running again."})188 await websocket.close()189 return190 ip_cooldowns[client_ip] = current_time191 192 # --- 2. PREVIOUS SESSION Auto-Kill ---193 # If the user restarts rapidly, kill their old process to free up one of the semaphore slots instantly.194 old_session_id = ip_active_sessions.get(client_ip)195 if old_session_id and old_session_id in active_sessions:196 old_session = active_sessions[old_session_id]197 old_proc = old_session.get("process")198 if old_proc and old_proc.returncode is None:199 try: old_proc.kill()200 except ProcessLookupError: pass201 202 # Give the asyncio event loop a tiny moment to run the old session's 'finally' block and release the lock203 await asyncio.sleep(0.15) 204 205 session_id = str(uuid.uuid4())206 ip_active_sessions[client_ip] = session_id207 208 # Store initial connection state before code execution209 active_sessions[session_id] = {210 "process": None,211 "last_active": time.time(),212 "websocket": websocket,213 "temp_path": None,214 "holds_semaphore": False215 }216 217 try:218 # 1. Wait for the initial code payload from the client219 initial_data = await websocket.receive_json()220 user_code = initial_data.get("code", "")221 222 if not user_code:223 await websocket.send_json({"type": "error", "data": "No code provided"})224 await websocket.close()225 return226 227 # 2. Check Capacity limits228 global execution_semaphore229 if execution_semaphore.locked():230 await websocket.send_json({"type": "error", "data": "Server busy. Try again."})231 await websocket.close()232 return233 234 await execution_semaphore.acquire()235 active_sessions[session_id]["holds_semaphore"] = True236 237 # 3. Build & Write File (Injected PATCH_INPUT here)238 final_code = PATCH_LIMITS + PATCH_HEADER + PATCH_INPUT239 if "matplotlib" in user_code or "plt." in user_code: final_code += PATCH_MATPLOTLIB240 if "PIL" in user_code or "Image" in user_code: final_code += PATCH_PIL241 final_code += "\n" + user_code242 243 temp_dir = "/dev/shm" if os.path.exists("/dev/shm") else "/tmp"244 temp_path = os.path.join(temp_dir, f"script_{session_id}.py")245 active_sessions[session_id]["temp_path"] = temp_path246 247 with open(temp_path, "w", encoding="utf-8") as f:248 f.write(final_code)249 250 # 4. Spawn Process (Unbuffered output is critical: "-u")251 proc = await asyncio.create_subprocess_exec(252 sys.executable, "-u", temp_path,253 stdout=asyncio.subprocess.PIPE,254 stderr=asyncio.subprocess.STDOUT, 255 stdin=asyncio.subprocess.PIPE,256 cwd=temp_dir257 )258 active_sessions[session_id]["process"] = proc259 260 # 5. Define Output Reader Loop (With Safe Buffer & Image Detection)261 async def stream_output():262 buffer = ""263 is_collecting_image = False264 265 try:266 while True:267 chunk = await proc.stdout.read(1024)268 if not chunk:269 if buffer:270 await websocket.send_json({"type": "output", "data": buffer})271 break272 273 text = chunk.decode('utf-8', errors='replace')274 buffer += text275 276 # --- SCENARIO A: We are currently collecting a massive image string ---277 if is_collecting_image:278 if "@@@IMAGE_END@@@" in buffer:279 # The image is fully collected! Send the whole buffer.280 await websocket.send_json({"type": "output", "data": buffer})281 buffer = ""282 is_collecting_image = False283 # If the END tag isn't here yet, just keep looping and collecting284 continue285 286 # --- SCENARIO B: Check if an image just STARTED ---287 if "@@@IMAGE_START@@@" in buffer:288 is_collecting_image = True289 if "@@@IMAGE_END@@@" in buffer:290 # It started and finished in the same chunk291 await websocket.send_json({"type": "output", "data": buffer})292 buffer = ""293 is_collecting_image = False294 continue295 296 # --- SCENARIO C: Handle Interactive Inputs ---297 if "@@@WAITING_FOR_INPUT@@@" in buffer:298 parts = buffer.split("@@@WAITING_FOR_INPUT@@@")299 300 if parts[0]:301 await websocket.send_json({"type": "output", "data": parts[0]})302 303 await websocket.send_json({"type": "input_request"})304 buffer = parts[1] if len(parts) > 1 else ""305 306 else:307 # --- SCENARIO D: Standard safe text streaming ---308 if len(buffer) > 30:309 safe_to_send = buffer[:-30]310 buffer = buffer[-30:]311 await websocket.send_json({"type": "output", "data": safe_to_send})312 except Exception:313 pass314 315 reader_task = asyncio.create_task(stream_output())316 317 # 6. Define Input Listener Loop318 async def listen_for_input():319 try:320 while True:321 data = await websocket.receive_json()322 msg_type = data.get("type")323 324 if msg_type == "ping":325 # Ignore pings so the idle timer doesn't falsely reset326 pass327 328 elif msg_type == "input":329 # Reset idle timer on genuine input330 active_sessions[session_id]["last_active"] = time.time()331 332 user_input = data.get("data", "")333 if proc.returncode is None:334 proc.stdin.write((user_input + "\n").encode('utf-8'))335 await proc.stdin.drain()336 except WebSocketDisconnect:337 pass338 339 listener_task = asyncio.create_task(listen_for_input())340 341 # 7. Wait for execution to finish naturally342 await proc.wait()343 344 # Ensure we read any final trailing output345 await reader_task 346 listener_task.cancel()347 348 await websocket.send_json({"type": "status", "data": "completed"})349 350 except WebSocketDisconnect:351 pass # Handle silent client disconnections securely352 except Exception as e:353 try: await websocket.send_json({"type": "error", "data": str(e)})354 except: pass355 356 finally:357 # 8. Complete Cleanup on Disconnect or Finish358 if session_id in active_sessions:359 session = active_sessions[session_id]360 proc = session.get("process")361 if proc and proc.returncode is None:362 try: proc.kill()363 except: pass364 365 if session.get("holds_semaphore"):366 execution_semaphore.release()367 368 temp_path = session.get("temp_path")369 if temp_path and os.path.exists(temp_path):370 try: os.remove(temp_path)371 except: pass372 373 del active_sessions[session_id]374 375 # Clean up the IP tracking to prevent memory leaks376 if ip_active_sessions.get(client_ip) == session_id:377 del ip_active_sessions[client_ip]378 379 try: await websocket.close()380 except: pass