add utils.restricted_exec

add app.tasks

2 files changed, 45 insertions, 0 deletions

diff --git a/app/main.py b/app/main.py
index 59e7a8b..b90fef9 100644
--- a/app/main.py
+++ b/app/main.py
@@ -1,11 +1,15 @@
 from typing import Annotated
+import asyncio
+from datetime import datetime, timedelta
 
 from fastapi import FastAPI, Request, Depends
 from fastapi.responses import HTMLResponse
 from fastapi.staticfiles import StaticFiles
 from starlette.templating import Jinja2Templates
+from pydantic import BaseModel
 
 from utils import get_avatar_urls, get_client_geo
+from app.tasks import restricted_exec_task
 
 templates = Jinja2Templates(directory="templates")
 
@@ -33,3 +37,27 @@ async def client_addr(
         "partials/client_geo.html",
         {"request": request, "client_geo": client_geo},
     )
+
+
+class RestrictedExecBase(BaseModel):
+    code: str
+
+
+@app.post('/restricted_exec')
+async def post_restricted_exec(
+    body: RestrictedExecBase
+):
+    result = restricted_exec_task.delay(body.code)
+
+    start_time = datetime.now()
+    while datetime.now() - start_time < timedelta(seconds=5):
+        if result.ready():
+            return result.get()
+
+        await asyncio.sleep(0.1)
+
+    result.revoke(terminate=True, signal='SIGTERM')
+    return [
+        'Execution timeout, task revoked',
+        result.status
+    ]
diff --git a/app/tasks.py b/app/tasks.py
new file mode 100644
index 0000000..bba3fae
--- /dev/null
+++ b/app/tasks.py
@@ -0,0 +1,17 @@
+from celery import Celery
+
+from utils import restricted_exec
+
+app = Celery(
+    'tasks',
+    broker='redis://localhost:6379/0',
+    backend='redis://localhost:6379/0',
+    task_send_sent_event=True,
+    worker_send_task_events=True,
+    worker_enable_remote_control=True,
+)
+
+
+@app.task
+def restricted_exec_task(code: str):
+    return restricted_exec.getoutput(code)