From 05aa008ce12b919c9cd0607e016a7fabc34e7f4c Mon Sep 17 00:00:00 2001
From: Jeremy Stretch <stretch@packetlife.net>
Date: Mon, 13 Jul 2020 15:11:58 -0400
Subject: [PATCH] #2006: Prevent script/report execution if no RQ worker is
 running

---
 netbox/extras/api/views.py     | 12 +++++++++++-
 netbox/extras/views.py         | 21 +++++++++++++++------
 netbox/utilities/exceptions.py | 14 ++++++++++++++
 3 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/netbox/extras/api/views.py b/netbox/extras/api/views.py
index e9e2b6c85..501182da5 100644
--- a/netbox/extras/api/views.py
+++ b/netbox/extras/api/views.py
@@ -3,11 +3,13 @@ from collections import OrderedDict
 from django.contrib.contenttypes.models import ContentType
 from django.db.models import Count
 from django.http import Http404
+from django_rq.queues import get_connection
 from rest_framework import status
 from rest_framework.decorators import action
 from rest_framework.exceptions import PermissionDenied
 from rest_framework.response import Response
 from rest_framework.viewsets import ReadOnlyModelViewSet, ViewSet
+from rq import Worker
 
 from extras import filters
 from extras.choices import JobResultStatusChoices
@@ -17,6 +19,7 @@ from extras.models import (
 from extras.reports import get_report, get_reports, run_report
 from extras.scripts import get_script, get_scripts, run_script
 from utilities.api import IsAuthenticatedOrLoginNotRequired, ModelViewSet
+from utilities.exceptions import RQWorkerNotRunningException
 from utilities.metadata import ContentTypeMetadata
 from utilities.utils import copy_safe_request
 from . import serializers
@@ -219,11 +222,14 @@ class ReportViewSet(ViewSet):
         """
         Run a Report identified as "<module>.<script>" and return the pending JobResult as the result
         """
-
         # Check that the user has permission to run reports.
         if not request.user.has_perm('extras.run_script'):
             raise PermissionDenied("This user does not have permission to run reports.")
 
+        # Check that at least one RQ worker is running
+        if not Worker.count(get_connection('default')):
+            raise RQWorkerNotRunningException()
+
         # Retrieve and run the Report. This will create a new JobResult.
         report = self._retrieve_report(pk)
         report_content_type = ContentType.objects.get(app_label='extras', model='report')
@@ -299,6 +305,10 @@ class ScriptViewSet(ViewSet):
         script = self._get_script(pk)()
         input_serializer = serializers.ScriptInputSerializer(data=request.data)
 
+        # Check that at least one RQ worker is running
+        if not Worker.count(get_connection('default')):
+            raise RQWorkerNotRunningException()
+
         if input_serializer.is_valid():
             data = input_serializer.data['data']
             commit = input_serializer.data['commit']
diff --git a/netbox/extras/views.py b/netbox/extras/views.py
index 10e711064..b38b39cd6 100644
--- a/netbox/extras/views.py
+++ b/netbox/extras/views.py
@@ -1,13 +1,14 @@
-import time
-
 from django import template
 from django.conf import settings
+from django.contrib import messages
 from django.contrib.contenttypes.models import ContentType
 from django.db.models import Count, Prefetch, Q
 from django.http import Http404, HttpResponseForbidden
 from django.shortcuts import get_object_or_404, redirect, render
 from django.views.generic import View
+from django_rq.queues import get_connection
 from django_tables2 import RequestConfig
+from rq import Worker
 
 from dcim.models import DeviceRole, Platform, Region, Site
 from tenancy.models import Tenant, TenantGroup
@@ -21,7 +22,7 @@ from utilities.views import (
 from virtualization.models import Cluster, ClusterGroup
 from . import filters, forms, tables
 from .choices import JobResultStatusChoices
-from .models import ConfigContext, ImageAttachment, ObjectChange, Report, JobResult, Script, Tag, TaggedItem
+from .models import ConfigContext, ImageAttachment, ObjectChange, JobResult, Tag
 from .reports import get_report, get_reports, run_report
 from .scripts import get_scripts, run_script
 
@@ -388,9 +389,13 @@ class ReportView(GetReportMixin, ContentTypePermissionRequiredMixin, View):
             return HttpResponseForbidden()
 
         report = self._get_report(name, module)
-
         form = ConfirmationForm(request.POST)
-        if form.is_valid():
+
+        # Allow execution only if RQ worker process is running
+        if not Worker.count(get_connection('default')):
+            messages.error(request, "Unable to run report: RQ worker process not running.")
+
+        elif form.is_valid():
 
             # Run the Report. A new JobResult is created.
             report_content_type = ContentType.objects.get(app_label='extras', model='report')
@@ -504,7 +509,11 @@ class ScriptView(ContentTypePermissionRequiredMixin, GetScriptMixin, View):
         script = self._get_script(name, module)
         form = script.as_form(request.POST, request.FILES)
 
-        if form.is_valid():
+        # Allow execution only if RQ worker process is running
+        if not Worker.count(get_connection('default')):
+            messages.error(request, "Unable to run script: RQ worker process not running.")
+
+        elif form.is_valid():
             commit = form.cleaned_data.pop('_commit')
 
             script_content_type = ContentType.objects.get(app_label='extras', model='script')
diff --git a/netbox/utilities/exceptions.py b/netbox/utilities/exceptions.py
index 5032aacee..77a915d9c 100644
--- a/netbox/utilities/exceptions.py
+++ b/netbox/utilities/exceptions.py
@@ -1,5 +1,19 @@
+from rest_framework import status
+from rest_framework.exceptions import APIException
+
+
 class AbortTransaction(Exception):
     """
     A dummy exception used to trigger a database transaction rollback.
     """
     pass
+
+
+class RQWorkerNotRunningException(APIException):
+    """
+    Indicates the temporary inability to enqueue a new task (e.g. custom script execution) because no RQ worker
+    processes are currently running.
+    """
+    status_code = status.HTTP_503_SERVICE_UNAVAILABLE
+    default_detail = 'Unable to process request: RQ worker process not running.'
+    default_code = 'rq_worker_not_running'