mirror of
https://github.com/yogeshojha/rengine.git
synced 2026-10-20 23:31:31 +02:00
89 lines
2.6 KiB
Python
89 lines
2.6 KiB
Python
from celery import Celery
|
|
|
|
from shared.config import BaseAppSettings
|
|
from shared.definitions.constants import SCANS_QUEUE
|
|
from shared.logging import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
_celery_client: Celery | None = None
|
|
|
|
|
|
def get_celery_client() -> Celery:
|
|
global _celery_client # noqa: PLW0603
|
|
if _celery_client is None:
|
|
_settings = BaseAppSettings()
|
|
_celery_client = Celery(broker=_settings.celery_broker_url)
|
|
_celery_client.conf.update(
|
|
task_serializer="json",
|
|
accept_content=["json"],
|
|
)
|
|
return _celery_client
|
|
|
|
|
|
def dispatch_whois_lookups(target_ids: list[str]) -> None:
|
|
if not target_ids:
|
|
return
|
|
logger.info("Dispatching WHOIS lookups for %d targets", len(target_ids))
|
|
get_celery_client().send_task(
|
|
"app.tasks.whois.perform_whois_lookups",
|
|
kwargs={"target_ids": target_ids},
|
|
queue="default",
|
|
)
|
|
|
|
|
|
def dispatch_ripestat_enrichment(target_ids: list[str]) -> None:
|
|
if not target_ids:
|
|
return
|
|
|
|
logger.info("Dispatching RIPEstat enrichment for %d targets", len(target_ids))
|
|
get_celery_client().send_task(
|
|
"app.tasks.ripestat.enrich_targets_bgp",
|
|
kwargs={"target_ids": target_ids},
|
|
queue="default",
|
|
)
|
|
|
|
|
|
def dispatch_dns_lookups(target_ids: list[str]) -> None:
|
|
if not target_ids:
|
|
return
|
|
|
|
logger.info(f"Dispatching DNS lookup for {len(target_ids)} target(s)")
|
|
get_celery_client().send_task(
|
|
"app.tasks.dns.perform_dns_lookups",
|
|
kwargs={"target_ids": target_ids},
|
|
queue="default",
|
|
)
|
|
|
|
|
|
def dispatch_scan_run(scan_id: str) -> None:
|
|
logger.info("Dispatching scan run %s", scan_id)
|
|
get_celery_client().send_task(
|
|
"app.tasks.scan.run_scan",
|
|
kwargs={"scan_id": scan_id},
|
|
queue=SCANS_QUEUE,
|
|
)
|
|
|
|
|
|
def revoke_scan_tasks(task_ids: list[str]) -> None:
|
|
"""SIGKILL-revoke a scan's celery tasks so an in-flight scan stops promptly."""
|
|
if not task_ids:
|
|
return
|
|
logger.info("Revoking %d scan task(s)", len(task_ids))
|
|
try:
|
|
get_celery_client().control.revoke(
|
|
list(task_ids), terminate=True, signal="SIGKILL"
|
|
)
|
|
except Exception:
|
|
logger.warning("scan task revoke failed", exc_info=True)
|
|
|
|
|
|
def dispatch_template_sync() -> bool:
|
|
"""Kick off a library refresh. Returns whether the queue accepted it."""
|
|
try:
|
|
get_celery_client().send_task("app.tasks.vuln_templates.sync", queue="default")
|
|
except Exception:
|
|
logger.warning("template sync dispatch failed", exc_info=True)
|
|
return False
|
|
return True
|