修改为同步数据库写入

Author: EITS403

app/tasks/sqlmap_scheduler.py

@@ -1,27 +1,20 @@
-import asyncio
-from celery import shared_task
 from sqlalchemy import select
-
-from app.database.database import AsyncSessionLocal
+from app.middleware.celery_app import celery_app
+from app.database.celery_sync_database import SessionLocal
 from app.models.sqlmap_result import SqlmapScanPayload, ScanStatus
 from app.tasks.sqlmap_worker import poll_single_sqlmap_task
 
 
-# 轮询查询数据库中的正在运行状态的数据
-@shared_task
+@celery_app.task(name="app.tasks.sqlmap_scheduler.poll_active_sqlmap_tasks")
 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]
-                    )
-                )
+    with SessionLocal() as session:
+        tasks = (
+            session.query(SqlmapScanPayload)
+            .filter(
+                SqlmapScanPayload.status.in_([ScanStatus.pending, ScanStatus.running])
             )
-            tasks = result.scalars().all()
-
-            for task in tasks:
-                poll_single_sqlmap_task.delay(task.task_id)
+            .all()
+        )
 
-    asyncio.run(_run())
+        for task in tasks:
+            poll_single_sqlmap_task.delay(task.task_id)

app/tasks/sqlmap_worker.py

@@ -1,90 +1,90 @@
-import asyncio
 import requests
-from celery import shared_task
 from datetime import datetime
-from sqlalchemy import select
+from sqlalchemy.exc import SQLAlchemyError
 
-from app.database.database import AsyncSessionLocal
+from app.middleware.celery_app import celery_app
+from app.database.celery_sync_database import SessionLocal
 from app.models.sqlmap_result import (
     SqlmapScanPayload,
     SqlmapScanLog,
     ScanStatus,
+    SqlmapScanResult,
 )
 import os
 
 SQLMAP_API = os.getenv("SQLMAP_API")
-AUTH = (os.getenv("SQLMAP_USERNAME"), os.getenv("SQLMAP_PASSWORD"))
+AUTH = (os.getenv("SQLMAP_USERNAME"), os.getenv("SQLMAP_PASSWORD"))  # Basic Auth
 
 
-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(
+@celery_app.task(
     bind=True,
     autoretry_for=(Exception,),
     retry_backoff=5,
     retry_kwargs={"max_retries": 3},
+    name="app.tasks.sqlmap_worker.poll_single_sqlmap_task",
 )
 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())
+    session = SessionLocal()
+    try:
+        task = (
+            session.query(SqlmapScanPayload)
+            .filter(SqlmapScanPayload.task_id == task_id)
+            .first()
+        )
+
+        if not task:
+            return
+
+        # 查询 sqlmap task 状态
+        resp = requests.get(
+            f"{SQLMAP_API}/scan/{task_id}/status",
+            timeout=10,
+            auth=AUTH,
+        )
+        resp.raise_for_status()
+        status_data = resp.json()
+
+        status = status_data.get("status")
+
+        if status == "running":
+            task.status = ScanStatus.running
+            session.commit()
+            return
+
+        if status != "terminated":
+            return
+
+        # 获取扫描结果
+        result_resp = requests.get(
+            f"{SQLMAP_API}/scan/{task_id}/data",
+            timeout=30,
+            auth=AUTH,
+        )
+        result_resp.raise_for_status()
+        data = result_resp.json()
+
+        # 解析 sqlmap 返回
+        scan_result = SqlmapScanResult(
+            target_url=task.scan_url,
+            dbms=data.get("dbms"),
+            vulnerable=bool(data.get("data")),
+            injection_points=data.get("data"),
+            dump_data=data.get("dump"),
+            raw_output=data.get("raw"),
+            command=data.get("command", ""),
+            started_at=datetime.utcnow(),
+            finished_at=datetime.utcnow(),
+        )
+
+        session.add(scan_result)
+        task.status = ScanStatus.success
+
+        session.commit()
+
+    except Exception:
+        session.rollback()
+        task.status = ScanStatus.failed
+        session.commit()
+        raise
+    finally:
+        session.close()