From f5830c1cd829681ddd2752643e462887dc81b121 Mon Sep 17 00:00:00 2001 From: jeremystretch Date: Fri, 24 Mar 2023 10:11:15 -0400 Subject: [PATCH] Enable resolving scripts/reports from module class --- netbox/extras/models/models.py | 32 ++++++++++++++++++++++++++++++++ netbox/extras/reports.py | 8 +------- netbox/extras/scripts.py | 10 +--------- netbox/extras/temp.py | 20 ++++++++++++++++++++ netbox/extras/views.py | 9 ++++----- 5 files changed, 58 insertions(+), 21 deletions(-) create mode 100644 netbox/extras/temp.py diff --git a/netbox/extras/models/models.py b/netbox/extras/models/models.py index 02d71538e..58de803a5 100644 --- a/netbox/extras/models/models.py +++ b/netbox/extras/models/models.py @@ -1,5 +1,7 @@ +import inspect import json import uuid +from functools import cached_property from pkgutil import ModuleInfo, get_importer from django.conf import settings @@ -31,6 +33,7 @@ from netbox.models.features import ( CloningMixin, CustomFieldsMixin, CustomLinksMixin, ExportTemplatesMixin, JobResultsMixin, SyncedDataMixin, TagsMixin, WebhooksMixin, ) +from ..temp import is_report, is_script from utilities.querysets import RestrictedQuerySet from utilities.rqworker import get_queue_for_model from utilities.utils import render_jinja2 @@ -827,6 +830,11 @@ class PythonModuleMixin: ispkg=False ) + def get_module(self): + importer, module_name, _ = self.get_module_info() + module = importer.find_module(module_name).load_module(module_name) + return module + class Script(JobResultsMixin, WebhooksMixin, models.Model): """ @@ -862,6 +870,18 @@ class ScriptModule(JobResultsMixin, WebhooksMixin, PythonModuleMixin, ManagedFil def name(self): return self.file_path + @cached_property + def scripts(self): + module = self.get_module() + + scripts = {} + for name, cls in inspect.getmembers(module, is_script): + # For child objects in submodules use the full import path w/o the root module as the name + child_name = cls.full_name.split(".", maxsplit=1)[1] + scripts[child_name] = cls + + return scripts + # # Reports @@ -896,3 +916,15 @@ class ReportModule(JobResultsMixin, WebhooksMixin, PythonModuleMixin, ManagedFil def get_absolute_url(self): return reverse('extras:report_list') + + @cached_property + def reports(self): + module = self.get_module() + + reports = {} + for name, cls in inspect.getmembers(module, is_report): + # For child objects in submodules use the full import path w/o the root module as the name + child_name = cls().full_name.split(".", maxsplit=1)[1] + reports[child_name] = cls + + return reports diff --git a/netbox/extras/reports.py b/netbox/extras/reports.py index fafe04634..c028a0a98 100644 --- a/netbox/extras/reports.py +++ b/netbox/extras/reports.py @@ -7,18 +7,12 @@ from django_rq import job from .choices import JobResultStatusChoices, LogLevelChoices from .models import JobResult, ReportModule +from .temp import is_report from .utils import get_modules logger = logging.getLogger(__name__) -def is_report(obj): - """ - Returns True if the given object is a Report. - """ - return obj in Report.__subclasses__() - - def get_reports(): return get_modules(ReportModule.objects.all(), is_report, 'report_order') diff --git a/netbox/extras/scripts.py b/netbox/extras/scripts.py index ad51280b6..347438f9f 100644 --- a/netbox/extras/scripts.py +++ b/netbox/extras/scripts.py @@ -22,6 +22,7 @@ from utilities.exceptions import AbortScript, AbortTransaction from utilities.forms import add_blank_choice, DynamicModelChoiceField, DynamicModelMultipleChoiceField from .context_managers import change_logging from .forms import ScriptForm +from .temp import is_script from .utils import get_modules __all__ = [ @@ -423,15 +424,6 @@ class Script(BaseScript): # Functions # -def is_script(obj): - """ - Returns True if the object is a Script. - """ - try: - return issubclass(obj, Script) and obj != Script - except TypeError: - return False - def is_variable(obj): """ diff --git a/netbox/extras/temp.py b/netbox/extras/temp.py new file mode 100644 index 000000000..5b07403b0 --- /dev/null +++ b/netbox/extras/temp.py @@ -0,0 +1,20 @@ +def is_script(obj): + """ + Returns True if the object is a Script. + """ + from .scripts import Script + try: + return issubclass(obj, Script) and obj != Script + except TypeError: + return False + + +def is_report(obj): + """ + Returns True if the given object is a Report. + """ + from .reports import Report + try: + return issubclass(obj, Report) and obj != Report + except TypeError: + return False diff --git a/netbox/extras/views.py b/netbox/extras/views.py index 4118a5d04..8495914ba 100644 --- a/netbox/extras/views.py +++ b/netbox/extras/views.py @@ -849,10 +849,8 @@ class ReportView(ContentTypePermissionRequiredMixin, View): return 'extras.view_report' def get(self, request, module, name): - - report = get_report(module, name) - if report is None: - raise Http404 + module = get_object_or_404(ReportModule.objects.restrict(request.user), file_path=f'{module}.py') + report = module.reports[name]() report_content_type = ContentType.objects.get(app_label='extras', model='report') report.result = JobResult.objects.filter( @@ -1001,7 +999,8 @@ class ScriptView(ContentTypePermissionRequiredMixin, GetScriptMixin, View): return 'extras.view_script' def get(self, request, module, name): - script = self._get_script(name, module) + module = get_object_or_404(ScriptModule.objects.restrict(request.user), file_path=f'{module}.py') + script = module.scripts[name]() form = script.as_form(initial=normalize_querydict(request.GET)) # Look for a pending JobResult (use the latest one by creation timestamp)