增加定时和轮询任务

Author: EITS403

app/middleware/celery_beat.py

@@ -0,0 +1,9 @@
+from celery.schedules import crontab
+from app.middleware import celery_app
+
+celery_app.conf.beat_schedule = {
+    "poll-sqlmap-tasks-every-5-seconds": {
+        "task": "app.tasks.sqlmap_scheduler.poll_active_sqlmap_tasks",
+        "schedule": 5.0,
+    }
+}

app/tasks/sqlmap_scheduler.py

@@ -0,0 +1,27 @@
+import asyncio
+from celery import shared_task
+from sqlalchemy import select
+
+from app.database.database import AsyncSessionLocal
+from app.models.sqlmap_result import SqlmapScanPayload, ScanStatus
+from app.tasks.sqlmap_worker import poll_single_sqlmap_task
+
+
+# 轮询查询数据库中的正在运行状态的数据
+@shared_task
+def poll_active_sqlmap_tasks():
+    async def _run():
+        async with AsyncSessionLocal() as session:
+            result = await session.execute(
+                select(SqlmapScanPayload).where(
+                    SqlmapScanPayload.status.in_(
+                        [ScanStatus.pending, ScanStatus.running]
+                    )
+                )
+            )
+            tasks = result.scalars().all()
+
+            for task in tasks:
+                poll_single_sqlmap_task.delay(task.task_id)
+
+    asyncio.run(_run())

app/tasks/sqlmap_worker.py

@@ -0,0 +1,90 @@
+import asyncio
+import requests
+from celery import shared_task
+from datetime import datetime
+from sqlalchemy import select
+
+from app.database.database import AsyncSessionLocal
+from app.models.sqlmap_result import (
+    SqlmapScanPayload,
+    SqlmapScanLog,
+    ScanStatus,
+)
+import os
+
+SQLMAP_API = os.getenv("SQLMAP_API")
+AUTH = (os.getenv("SQLMAP_USERNAME"), os.getenv("SQLMAP_PASSWORD"))
+
+
+def fetch_status(task_id: str) -> dict:
+    r = requests.get(f"{SQLMAP_API}/scan/{task_id}/status", auth=AUTH, timeout=10)
+    r.raise_for_status()
+    return r.json()
+
+
+def fetch_log(task_id: str) -> dict:
+    r = requests.get(f"{SQLMAP_API}/scan/{task_id}/log", auth=AUTH, timeout=10)
+    r.raise_for_status()
+    return r.json()
+
+
+@shared_task(
+    bind=True,
+    autoretry_for=(Exception,),
+    retry_backoff=5,
+    retry_kwargs={"max_retries": 3},
+)
+def poll_single_sqlmap_task(self, task_id: str):
+    """
+    Worker：轮询单个 sqlmap 任务
+    """
+
+    async def _run():
+        async with AsyncSessionLocal() as session:
+            task = await session.scalar(
+                select(SqlmapScanPayload).where(SqlmapScanPayload.task_id == task_id)
+            )
+
+            if not task:
+                return
+
+            # 已结束 → 不再轮询
+            if task.status in (
+                ScanStatus.success,
+                ScanStatus.failed,
+                ScanStatus.stopped,
+            ):
+                return
+
+            # --- 调 sqlmap ---
+            status_data = fetch_status(task_id)
+            log_data = fetch_log(task_id)
+
+            sqlmap_status = status_data.get("status")
+
+            # --- 状态映射 ---
+            if sqlmap_status == "running":
+                task.status = ScanStatus.running
+
+            elif sqlmap_status == "terminated":
+                task.status = ScanStatus.success
+                task.finished_at = datetime.utcnow()
+
+            elif sqlmap_status == "error":
+                task.status = ScanStatus.failed
+                task.finished_at = datetime.utcnow()
+
+            # --- 写日志（全量 or 增量）---
+            for item in log_data.get("log", []):
+                session.add(
+                    SqlmapScanLog(
+                        task_id=task_id,
+                        level=item.get("level", "INFO"),
+                        message=item.get("message", ""),
+                        log_time=item.get("time"),
+                    )
+                )
+
+            await session.commit()
+
+    asyncio.run(_run())