Просмотр исходного кода

Fixes #21737: Check that uploaded custom scripts are valid Python modules before saving (#21920)

Jeremy Stretch 2 месяцев назад
Родитель
Сommit
bce667300a

+ 10 - 0
netbox/extras/api/serializers_/scripts.py

@@ -9,6 +9,7 @@ from rest_framework import serializers
 from core.api.serializers_.jobs import JobSerializer
 from core.choices import ManagedFileRootPathChoices
 from extras.models import Script, ScriptModule
+from extras.utils import validate_script_content
 from netbox.api.serializers import ValidatedModelSerializer
 from utilities.datetime import local_now
 
@@ -39,6 +40,15 @@ class ScriptModuleSerializer(ValidatedModelSerializer):
         data = super().validate(data)
         data.pop('file_root', None)
         if file is not None:
+            # Validate that the uploaded script can be loaded as a Python module
+            content = file.read()
+            file.seek(0)
+            try:
+                validate_script_content(content, file.name)
+            except Exception as e:
+                raise serializers.ValidationError(
+                    _("Error loading script: {error}").format(error=e)
+                )
             data['file'] = file
         return data
 

+ 17 - 0
netbox/extras/forms/scripts.py

@@ -4,6 +4,7 @@ from django.utils.translation import gettext_lazy as _
 
 from core.choices import JobIntervalChoices
 from core.forms import ManagedFileForm
+from extras.utils import validate_script_content
 from utilities.datetime import local_now
 from utilities.forms.widgets import DateTimePicker, NumberWithOptions
 
@@ -64,6 +65,22 @@ class ScriptFileForm(ManagedFileForm):
     """
     ManagedFileForm with a custom save method to use django-storages.
     """
+    def clean(self):
+        super().clean()
+
+        if upload_file := self.cleaned_data.get('upload_file'):
+            # Validate that the uploaded script can be loaded as a Python module
+            content = upload_file.read()
+            upload_file.seek(0)
+            try:
+                validate_script_content(content, upload_file.name)
+            except Exception as e:
+                raise forms.ValidationError(
+                    _("Error loading script: {error}").format(error=e)
+                )
+
+        return self.cleaned_data
+
     def save(self, *args, **kwargs):
         # If a file was uploaded, save it to disk
         if self.cleaned_data['upload_file']:

+ 15 - 0
netbox/extras/tests/test_api.py

@@ -1450,6 +1450,21 @@ class ScriptModuleTest(APITestCase):
         self.assertTrue(ScriptModule.objects.filter(file_path='test_upload.py').exists())
         self.assertTrue(Script.objects.filter(module__file_path='test_upload.py', name='TestScript').exists())
 
+    def test_upload_faulty_script_module(self):
+        """Uploading a script with an import error should return 400 and not create a DB record."""
+        self.add_permissions('extras.add_scriptmodule', 'core.add_managedfile')
+        # 'extras.script' is invalid; the correct module is 'extras.scripts'
+        script_content = b"from extras.script import Script\nclass TestScript(Script):\n    pass\n"
+        upload_file = SimpleUploadedFile('test_faulty.py', script_content, content_type='text/plain')
+        response = self.client.post(
+            self.url,
+            {'file': upload_file},
+            format='multipart',
+            **self.header,
+        )
+        self.assertHttpStatus(response, status.HTTP_400_BAD_REQUEST)
+        self.assertFalse(ScriptModule.objects.filter(file_path='test_faulty.py').exists())
+
     def test_upload_script_module_without_file_fails(self):
         self.add_permissions('extras.add_scriptmodule', 'core.add_managedfile')
         response = self.client.post(self.url, {}, format='json', **self.header)

+ 13 - 0
netbox/extras/utils.py

@@ -1,4 +1,5 @@
 import importlib
+import types
 from pathlib import Path
 
 from django.core.exceptions import ImproperlyConfigured, SuspiciousFileOperation
@@ -21,6 +22,7 @@ __all__ = (
     'is_script',
     'is_taggable',
     'run_validators',
+    'validate_script_content',
 )
 
 
@@ -134,6 +136,17 @@ def is_script(obj):
         return False
 
 
+def validate_script_content(content, filename):
+    """
+    Validate that the given content can be loaded as a Python module by compiling
+    and executing it. Raises an exception if the script cannot be loaded.
+    """
+    code = compile(content, filename, 'exec')
+    module_name = Path(filename).stem
+    module = types.ModuleType(module_name)
+    exec(code, module.__dict__)
+
+
 def is_report(obj):
     """
     Returns True if the given object is a Report.