Add support loading module as python package or standalone file

Author: Omripresent

netbox/extras/scripts.py

@@ -10,6 +10,7 @@
 from django.conf import settings
 from django.core.validators import RegexValidator
 from django.db import transaction
+from django.db.models import Q
 from django.utils.functional import classproperty
 
 from core.choices import JobStatusChoices
@@ -462,7 +463,7 @@ def is_variable(obj):
 
 
 def get_module_and_script(module_name, script_name):
-    module = ScriptModule.objects.get(file_path=f'{module_name}.py')
+    module = ScriptModule.objects.get(Q(file_path=f'{module_name}.py') | Q(file_path=f'{module_name}/__init__.py'))
     script = module.scripts.get(script_name)
     return module, script
 

netbox/extras/views.py

@@ -855,6 +855,7 @@ class ReportListView(ContentTypePermissionRequiredMixin, View):
     """
     Retrieve all the available reports from disk and the recorded Job (if any) for each.
     """
+
     def get_required_permission(self):
         return 'extras.view_report'
 
@@ -871,11 +872,15 @@ class ReportView(ContentTypePermissionRequiredMixin, View):
     """
     Display a single Report and its associated Job (if any).
     """
+
     def get_required_permission(self):
         return 'extras.view_report'
 
     def get(self, request, module, name):
-        module = get_object_or_404(ReportModule.objects.restrict(request.user), file_path__startswith=module)
+        module = get_object_or_404(
+            ReportModule.objects.restrict(request.user),
+            Q(file_path=f'{module}.py') | Q(file_path=f'{module}/__init__.py')
+        )
         report = module.reports[name]()
 
         object_type = ContentType.objects.get(app_label='extras', model='reportmodule')
@@ -896,7 +901,10 @@ def post(self, request, module, name):
         if not request.user.has_perm('extras.run_report'):
             return HttpResponseForbidden()
 
-        module = get_object_or_404(ReportModule.objects.restrict(request.user), file_path__startswith=module)
+        module = get_object_or_404(
+            ReportModule.objects.restrict(request.user),
+            Q(file_path=f'{module}.py') | Q(file_path=f'{module}/__init__.py'),
+        )
         report = module.reports[name]()
         form = ReportForm(request.POST, scheduling_enabled=report.scheduling_enabled)
 
@@ -935,7 +943,10 @@ def get_required_permission(self):
         return 'extras.view_report'
 
     def get(self, request, module, name):
-        module = get_object_or_404(ReportModule.objects.restrict(request.user), file_path__startswith=module)
+        module = get_object_or_404(
+            ReportModule.objects.restrict(request.user),
+            Q(file_path=f'{module}.py') | Q(file_path=f'{module}/__init__.py')
+        )
         report = module.reports[name]()
 
         return render(request, 'extras/report/source.html', {
@@ -951,7 +962,10 @@ def get_required_permission(self):
         return 'extras.view_report'
 
     def get(self, request, module, name):
-        module = get_object_or_404(ReportModule.objects.restrict(request.user), file_path__startswith=module)
+        module = get_object_or_404(
+            ReportModule.objects.restrict(request.user),
+            Q(file_path=f'{module}.py') | Q(file_path=f'{module}/__init__.py'),
+        )
         report = module.reports[name]()
 
         object_type = ContentType.objects.get(app_label='extras', model='reportmodule')
@@ -980,6 +994,7 @@ class ReportResultView(ContentTypePermissionRequiredMixin, View):
     """
     Display a Job pertaining to the execution of a Report.
     """
+
     def get_required_permission(self):
         return 'extras.view_report'
 
@@ -1046,7 +1061,10 @@ def get_required_permission(self):
         return 'extras.view_script'
 
     def get(self, request, module, name):
-        module = get_object_or_404(ScriptModule.objects.restrict(request.user), file_path__startswith=module)
+        module = get_object_or_404(
+            ScriptModule.objects.restrict(request.user),
+            Q(file_path=f'{module}.py') | Q(file_path=f'{module}/__init__.py'),
+        )
         script = module.scripts[name]()
         form = script.as_form(initial=normalize_querydict(request.GET))
 
@@ -1070,7 +1088,10 @@ def post(self, request, module, name):
         if not request.user.has_perm('extras.run_script'):
             return HttpResponseForbidden()
 
-        module = get_object_or_404(ScriptModule.objects.restrict(request.user), file_path__startswith=module)
+        module = get_object_or_404(
+            ScriptModule.objects.restrict(request.user),
+            Q(file_path=f'{module}.py') | Q(file_path=f'{module}/__init__.py'),
+        )
         script = module.scripts[name]()
         form = script.as_form(request.POST, request.FILES)
 
@@ -1107,7 +1128,10 @@ def get_required_permission(self):
         return 'extras.view_script'
 
     def get(self, request, module, name):
-        module = get_object_or_404(ScriptModule.objects.restrict(request.user), file_path__startswith=module)
+        module = get_object_or_404(
+            ScriptModule.objects.restrict(request.user),
+            Q(file_path=f'{module}.py') | Q(file_path=f'{module}/__init__.py'),
+        )
         script = module.scripts[name]()
 
         return render(request, 'extras/script/source.html', {
@@ -1123,7 +1147,10 @@ def get_required_permission(self):
         return 'extras.view_script'
 
     def get(self, request, module, name):
-        module = get_object_or_404(ScriptModule.objects.restrict(request.user), file_path__startswith=module)
+        module = get_object_or_404(
+            ScriptModule.objects.restrict(request.user),
+            Q(file_path=f'{module}.py') | Q(file_path=f'{module}/__init__.py'),
+        )
         script = module.scripts[name]()
 
         object_type = ContentType.objects.get(app_label='extras', model='scriptmodule')