From de30f551348c08ba333a53c0101ccf20ff661d10 Mon Sep 17 00:00:00 2001
From: Arthur <worldnomad@gmail.com>
Date: Fri, 13 Oct 2023 13:34:40 -0700
Subject: [PATCH] 14025 add file extension validation and simplify get logic

---
 netbox/core/forms/model_forms.py | 26 ++++++++++++++++++++++++--
 netbox/extras/views.py           | 17 ++---------------
 2 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/netbox/core/forms/model_forms.py b/netbox/core/forms/model_forms.py
index 01d5474c6..c2c8ad9d5 100644
--- a/netbox/core/forms/model_forms.py
+++ b/netbox/core/forms/model_forms.py
@@ -1,12 +1,14 @@
 import copy
 
 from django import forms
+from django.core.exceptions import ValidationError
 from django.utils.translation import gettext_lazy as _
 
 from core.forms.mixins import SyncedDataMixin
 from core.models import *
 from netbox.forms import NetBoxModelForm
 from netbox.registry import registry
+from pathlib import Path
 from utilities.forms import get_field_value
 from utilities.forms.fields import CommentField
 from utilities.forms.widgets import HTMXSelect
@@ -88,14 +90,34 @@ class ManagedFileForm(SyncedDataMixin, NetBoxModelForm):
         model = ManagedFile
         fields = ('data_source', 'data_file', 'auto_sync_enabled')
 
+    def _validate_file_extension(self, file):
+        if file:
+            extension = Path(file.name).suffix[1:].lower()
+
+            if extension != "py":
+                raise ValidationError(
+                    "File extension “%(extension)s” is not allowed .py extension is required",
+                    code="invalid_extension",
+                    params={
+                        "extension": extension,
+                    },
+                )
+
     def clean(self):
         super().clean()
 
-        if self.cleaned_data.get('upload_file') and self.cleaned_data.get('data_file'):
+        upload_file = self.cleaned_data['upload_file']
+        data_file = self.cleaned_data['data_file']
+
+        if upload_file and data_file:
+            extension = Path(upload_file.name).suffix[1:].lower()
             raise forms.ValidationError("Cannot upload a file and sync from an existing file")
-        if not self.cleaned_data.get('upload_file') and not self.cleaned_data.get('data_file'):
+        if not upload_file and not data_file:
             raise forms.ValidationError("Must upload a file or select a data file to sync")
 
+        self._validate_file_extension(upload_file)
+        self._validate_file_extension(data_file)
+
         return self.cleaned_data
 
     def save(self, *args, **kwargs):
diff --git a/netbox/extras/views.py b/netbox/extras/views.py
index 04ba8c774..666de71ac 100644
--- a/netbox/extras/views.py
+++ b/netbox/extras/views.py
@@ -3,7 +3,7 @@ from django.contrib.auth.mixins import LoginRequiredMixin
 from django.contrib.contenttypes.models import ContentType
 from django.core.paginator import EmptyPage
 from django.db.models import Count, Q
-from django.http import HttpResponseBadRequest, HttpResponseForbidden, HttpResponse, Http404
+from django.http import HttpResponseBadRequest, HttpResponseForbidden, HttpResponse
 from django.shortcuts import get_object_or_404, redirect, render
 from django.urls import reverse
 from django.utils.translation import gettext as _
@@ -1152,20 +1152,7 @@ class ScriptListView(ContentTypePermissionRequiredMixin, View):
 
 
 def get_script_module(module, request):
-    modules = ScriptModule.objects.restrict(request.user).filter(file_path__startswith=module)
-    if not modules:
-        raise Http404
-
-    if len(modules) == 1:
-        return modules[0]
-
-    # module is without the ".py" so startswith=module can return multiple if the file_path has
-    # two modules starting with the same characters "test.py" and "testfile.py"
-    for obj in modules:
-        if obj.file_path == module or obj.file_path == module + ".py":
-            return obj
-
-    raise Http404
+    return get_object_or_404(ScriptModule.objects.restrict(request.user), file_path=f"{module}.py")
 
 
 class ScriptView(ContentTypePermissionRequiredMixin, View):