--

src/nercone_website/tools/tls_test/__init__.py

@@ -1,5 +1,13 @@
 from .runner import TlsJobQueue
 from .db import TlsTestDB
-from .views import tls_submit, tls_api_submit, tls_results_context
+from .engine import run_full_scan
+from .views import tls_submit, tls_api_submit, tls_results_context, tls_websocket_handler
 
-__all__ = ["TlsJobQueue", "TlsTestDB", "tls_submit", "tls_api_submit", "tls_results_context"]
+tls_test_db = TlsTestDB()
+tls_test_queue = TlsJobQueue(tls_test_db, run_full_scan)
+
+__all__ = [
+    "TlsJobQueue", "TlsTestDB",
+    "tls_submit", "tls_api_submit", "tls_results_context", "tls_websocket_handler",
+    "tls_test_db", "tls_test_queue",
+]

src/nercone_website/tools/tls_test/views.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
+import json
 import datetime
-from fastapi import Request
+from fastapi import Request, WebSocket, WebSocketDisconnect
 from fastapi.templating import Jinja2Templates
 from fastapi.responses import Response
 from .engine import validate_tls_target
@@ -101,3 +102,39 @@ def tls_results_context(job: dict, test_id: str, request: Request, tls_test_db)
         "log_entries": log_entries,
         "history": history,
     }
+
+
+async def tls_websocket_handler(websocket: WebSocket, test_id: str, tls_test_db, tls_test_queue) -> None:
+    job = tls_test_db.get_job(test_id)
+    if not job:
+        await websocket.close(code=4404)
+        return
+    await websocket.accept()
+    tls_test_queue.add_subscriber(test_id, websocket)
+    try:
+        await websocket.send_text(json.dumps({
+            "type": "history",
+            "status": job.get("status"),
+            "target": job.get("target"),
+            "entries": tls_test_db.get_progress(test_id),
+        }))
+        if job.get("status") == "done":
+            await websocket.send_text(json.dumps({
+                "type": "done",
+                "redirect": f"/tools/tls-test/results/{test_id}/",
+                "rank": job.get("rank"),
+                "score": job.get("score"),
+            }))
+            await websocket.close()
+            return
+        while True:
+            try:
+                await websocket.receive_text()
+            except WebSocketDisconnect:
+                break
+    except WebSocketDisconnect:
+        pass
+    except Exception:
+        pass
+    finally:
+        tls_test_queue.remove_subscriber(test_id, websocket)