netbox/netbox/extras/tests/test_scripts.py

import io
import sys
from datetime import UTC, date, datetime
from decimal import Decimal
from unittest.mock import patch

from django.core.files.uploadedfile import SimpleUploadedFile
from django.test import TestCase
from netaddr import IPAddress, IPNetwork

from dcim.models import DeviceRole
from extras.constants import SCRIPT_MODULE_NAME_PREFIX
from extras.models import ScriptModule
from extras.scripts import *

CHOICES = (
    ('ff0000', 'Red'),
    ('00ff00', 'Green'),
    ('0000ff', 'Blue')
)

YAML_DATA = """
Foo: 123
Bar: 456
Baz:
 - A
 - B
 - C
"""

JSON_DATA = """
{
  "Foo": 123,
  "Bar": 456,
  "Baz": ["A", "B", "C"]
}
"""


class ScriptVariablesTestCase(TestCase):

    def test_stringvar(self):

        class TestScript(Script):

            var1 = StringVar(
                min_length=3,
                max_length=3,
                regex=r'[a-z]+'
            )

        # Validate min_length enforcement
        data = {'var1': 'xx'}
        form = TestScript().as_form(data, None)
        self.assertFalse(form.is_valid())
        self.assertIn('var1', form.errors)

        # Validate max_length enforcement
        data = {'var1': 'xxxx'}
        form = TestScript().as_form(data, None)
        self.assertFalse(form.is_valid())
        self.assertIn('var1', form.errors)

        # Validate regex enforcement
        data = {'var1': 'ABC'}
        form = TestScript().as_form(data, None)
        self.assertFalse(form.is_valid())
        self.assertIn('var1', form.errors)

        # Validate valid data
        data = {'var1': 'abc'}
        form = TestScript().as_form(data, None)
        self.assertTrue(form.is_valid())
        self.assertEqual(form.cleaned_data['var1'], data['var1'])

    def test_textvar(self):

        class TestScript(Script):

            var1 = TextVar()

        # Validate valid data
        data = {'var1': 'This is a test string'}
        form = TestScript().as_form(data, None)
        self.assertTrue(form.is_valid())
        self.assertEqual(form.cleaned_data['var1'], data['var1'])

    def test_integervar(self):

        class TestScript(Script):

            var1 = IntegerVar(
                min_value=5,
                max_value=10
            )

        # Validate min_value enforcement
        data = {'var1': 4}
        form = TestScript().as_form(data, None)
        self.assertFalse(form.is_valid())
        self.assertIn('var1', form.errors)

        # Validate max_value enforcement
        data = {'var1': 11}
        form = TestScript().as_form(data, None)
        self.assertFalse(form.is_valid())
        self.assertIn('var1', form.errors)

        # Validate valid data
        data = {'var1': 7}
        form = TestScript().as_form(data, None)
        self.assertTrue(form.is_valid())
        self.assertEqual(form.cleaned_data['var1'], data['var1'])

    def test_decimalvar(self):

        class TestScript(Script):

            var1 = DecimalVar(
                min_value=-100.500,
                max_value=100.500,
                max_digits=6,
                decimal_places=3,
                required=False
            )

            var2 = DecimalVar(
                max_digits=3,
                decimal_places=1,
                required=False
            )

        # Validate min_value enforcement
        data = {'var1': -100.501}
        form = TestScript().as_form(data, None)
        self.assertFalse(form.is_valid())
        self.assertIn('var1', form.errors)

        # Validate max_value enforcement
        data = {'var1': 100.501}
        form = TestScript().as_form(data, None)
        self.assertFalse(form.is_valid())
        self.assertIn('var1', form.errors)

        # Validate max_digits enforcement
        data = {'var2': 123.4}
        form = TestScript().as_form(data, None)
        self.assertFalse(form.is_valid())
        self.assertIn('var2', form.errors)

        # Validate decimal_places
        data = {'var2': 1.23}
        form = TestScript().as_form(data, None)
        self.assertFalse(form.is_valid())
        self.assertIn('var2', form.errors)

        # Validate valid data
        data = {'var1': '50.123'}
        form = TestScript().as_form(data, None)
        self.assertTrue(form.is_valid())
        self.assertEqual(form.cleaned_data['var1'], Decimal(data['var1']))

    def test_booleanvar(self):

        class TestScript(Script):

            var1 = BooleanVar()

        # Validate True
        data = {'var1': True}
        form = TestScript().as_form(data, None)
        self.assertTrue(form.is_valid())
        self.assertEqual(form.cleaned_data['var1'], True)

        # Validate False
        data = {'var1': False}
        form = TestScript().as_form(data, None)
        self.assertTrue(form.is_valid())
        self.assertEqual(form.cleaned_data['var1'], False)

    def test_choicevar(self):

        class TestScript(Script):

            var1 = ChoiceVar(
                choices=CHOICES
            )

        # Validate valid choice
        data = {'var1': 'ff0000'}
        form = TestScript().as_form(data)
        self.assertTrue(form.is_valid())
        self.assertEqual(form.cleaned_data['var1'], 'ff0000')

        # Validate invalid choice
        data = {'var1': 'taupe'}
        form = TestScript().as_form(data)
        self.assertFalse(form.is_valid())

    def test_multichoicevar(self):

        class TestScript(Script):

            var1 = MultiChoiceVar(
                choices=CHOICES
            )

        # Validate single choice
        data = {'var1': ['ff0000']}
        form = TestScript().as_form(data)
        self.assertTrue(form.is_valid())
        self.assertEqual(form.cleaned_data['var1'], ['ff0000'])

        # Validate multiple choices
        data = {'var1': ('ff0000', '00ff00')}
        form = TestScript().as_form(data)
        self.assertTrue(form.is_valid())
        self.assertEqual(form.cleaned_data['var1'], ['ff0000', '00ff00'])

        # Validate invalid choice
        data = {'var1': 'taupe'}
        form = TestScript().as_form(data)
        self.assertFalse(form.is_valid())

    def test_objectvar(self):

        class TestScript(Script):
            var1 = ObjectVar(model=DeviceRole)

        # Populate some objects
        for i in range(1, 6):
            DeviceRole(
                name='Device Role {}'.format(i),
                slug='device-role-{}'.format(i)
            ).save()

        # Validate valid data
        data = {'var1': DeviceRole.objects.first().pk}
        form = TestScript().as_form(data, None)
        self.assertTrue(form.is_valid())
        self.assertEqual(form.cleaned_data['var1'].pk, data['var1'])

    def test_multiobjectvar(self):

        class TestScript(Script):
            var1 = MultiObjectVar(model=DeviceRole)

        # Populate some objects
        for i in range(1, 6):
            DeviceRole(
                name='Device Role {}'.format(i),
                slug='device-role-{}'.format(i)
            ).save()

        # Validate valid data
        data = {'var1': [role.pk for role in DeviceRole.objects.all()[:3]]}
        form = TestScript().as_form(data, None)
        self.assertTrue(form.is_valid())
        self.assertEqual(form.cleaned_data['var1'][0].pk, data['var1'][0])
        self.assertEqual(form.cleaned_data['var1'][1].pk, data['var1'][1])
        self.assertEqual(form.cleaned_data['var1'][2].pk, data['var1'][2])

    def test_filevar(self):

        class TestScript(Script):

            var1 = FileVar()

        # Dummy file
        testfile = SimpleUploadedFile(
            name='test_file.txt',
            content=b'This is a dummy file for testing'
        )

        # Validate valid data
        file_data = {'var1': testfile}
        form = TestScript().as_form(None, file_data)
        self.assertTrue(form.is_valid())
        self.assertEqual(form.cleaned_data['var1'], testfile)

    def test_ipaddressvar(self):

        class TestScript(Script):

            var1 = IPAddressVar()

        # Validate IP network enforcement
        data = {'var1': '1.2.3'}
        form = TestScript().as_form(data, None)
        self.assertFalse(form.is_valid())
        self.assertIn('var1', form.errors)

        # Validate IP mask exclusion
        data = {'var1': '192.0.2.0/24'}
        form = TestScript().as_form(data, None)
        self.assertFalse(form.is_valid())
        self.assertIn('var1', form.errors)

        # Validate valid data
        data = {'var1': '192.0.2.1'}
        form = TestScript().as_form(data, None)
        self.assertTrue(form.is_valid())
        self.assertEqual(form.cleaned_data['var1'], IPAddress(data['var1']))

    def test_ipaddresswithmaskvar(self):

        class TestScript(Script):

            var1 = IPAddressWithMaskVar()

        # Validate IP network enforcement
        data = {'var1': '1.2.3'}
        form = TestScript().as_form(data, None)
        self.assertFalse(form.is_valid())
        self.assertIn('var1', form.errors)

        # Validate IP mask requirement
        data = {'var1': '192.0.2.0'}
        form = TestScript().as_form(data, None)
        self.assertFalse(form.is_valid())
        self.assertIn('var1', form.errors)

        # Validate valid data
        data = {'var1': '192.0.2.0/24'}
        form = TestScript().as_form(data, None)
        self.assertTrue(form.is_valid())
        self.assertEqual(form.cleaned_data['var1'], IPNetwork(data['var1']))

    def test_ipnetworkvar(self):

        class TestScript(Script):

            var1 = IPNetworkVar()

        # Validate IP network enforcement
        data = {'var1': '1.2.3'}
        form = TestScript().as_form(data, None)
        self.assertFalse(form.is_valid())
        self.assertIn('var1', form.errors)

        # Validate host IP check
        data = {'var1': '192.0.2.1/24'}
        form = TestScript().as_form(data, None)
        self.assertFalse(form.is_valid())
        self.assertIn('var1', form.errors)

        # Validate valid data
        data = {'var1': '192.0.2.0/24'}
        form = TestScript().as_form(data, None)
        self.assertTrue(form.is_valid())
        self.assertEqual(form.cleaned_data['var1'], IPNetwork(data['var1']))

    def test_datevar(self):

        class TestScript(Script):

            var1 = DateVar()
            var2 = DateVar(required=False)

        # Test date validation
        data = {'var1': 'not a date'}
        form = TestScript().as_form(data, None)
        self.assertFalse(form.is_valid())
        self.assertIn('var1', form.errors)

        # Validate valid data
        input_date = date(2024, 4, 1)
        data = {'var1': input_date}
        form = TestScript().as_form(data, None)
        self.assertTrue(form.is_valid())
        self.assertEqual(form.cleaned_data['var1'], input_date)
        # Validate required=False works for this Var type
        self.assertEqual(form.cleaned_data['var2'], None)

    def test_datetimevar(self):

        class TestScript(Script):

            var1 = DateTimeVar()
            var2 = DateTimeVar(required=False)

        # Test datetime validation
        data = {'var1': 'not a datetime'}
        form = TestScript().as_form(data, None)
        self.assertFalse(form.is_valid())
        self.assertIn('var1', form.errors)

        # Validate valid data
        input_datetime = datetime(2024, 4, 1, 8, 0, 0, 0, UTC)
        data = {'var1': input_datetime}
        form = TestScript().as_form(data, None)
        self.assertTrue(form.is_valid())
        self.assertEqual(form.cleaned_data['var1'], input_datetime)
        # Validate required=False works for this Var type
        self.assertEqual(form.cleaned_data['var2'], None)


class ScriptModuleLoadingTestCase(TestCase):

    def test_module_does_not_shadow_core_app(self):
        """
        Loading a custom script whose filename matches a core app label must not replace that
        app's package in sys.modules. Regression test for issue #22566.
        """
        import circuits  # The real core app package

        script_content = (
            b"from extras.scripts import Script\n\n\n"
            b"class TestScript(Script):\n    pass\n"
        )

        class _Storage:
            def open(self, name, mode='rb'):
                return io.BytesIO(script_content)

        module = ScriptModule(file_root='scripts', file_path='circuits.py')
        namespaced_key = f'{SCRIPT_MODULE_NAME_PREFIX}circuits'
        self.addCleanup(lambda: sys.modules.pop(namespaced_key, None))

        with patch('extras.models.mixins.storages') as mock_storages:
            mock_storages.__getitem__.return_value = _Storage()
            loaded = module.get_module()

            # The script module is registered under the private, namespaced key, and its own
            # __name__ matches that key (i.e. sys.modules[module.__name__] resolves to the module)
            self.assertIs(sys.modules[namespaced_key], loaded)
            self.assertEqual(loaded.__name__, namespaced_key)

            # The namespacing must not leak into the derived script name stored in the database
            self.assertEqual(next(iter(module.module_scripts)), 'TestScript')

            # Nor into the user-facing names exposed on the Script class (used for logger
            # namespaces, page headers, etc.): these must reflect the original filename.
            script_class = loaded.TestScript
            self.assertEqual(script_class.module, 'circuits')
            self.assertEqual(script_class.full_name, 'circuits.TestScript')
            self.assertEqual(script_class.root_module(), 'circuits')

        # The real circuits app must be untouched and remain an importable package
        self.assertIs(sys.modules['circuits'], circuits)
        self.assertTrue(hasattr(circuits, '__path__'))