This commit is contained in:
parent
660ca42149
commit
bce667300a
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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']:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Reference in New Issue