netbox/netbox/extras/tests/test_utils.py

from types import SimpleNamespace
from unittest.mock import patch

from django.contrib.contenttypes.models import ContentType
from django.core.files.storage import Storage
from django.test import TestCase

from extras.models import ExportTemplate, ImageAttachment
from extras.utils import _build_image_attachment_path, filename_from_model, image_upload
from tenancy.models import ContactGroup, TenantGroup
from wireless.models import WirelessLANGroup


class FilenameFromModelTestCase(TestCase):
    def test_expected_output(self):
        cases = (
            (ExportTemplate, 'netbox_export_templates'),
            (ContactGroup, 'netbox_contact_groups'),
            (TenantGroup, 'netbox_tenant_groups'),
            (WirelessLANGroup, 'netbox_wireless_lan_groups'),
        )

        for model, expected in cases:
            self.assertEqual(filename_from_model(model), expected)


class OverwriteStyleStorage(Storage):
    """
    Mimic an overwrite-style backend (for example, S3 with file_overwrite=True),
    where get_available_name() returns the incoming name unchanged.
    """

    def __init__(self, existing_names=None):
        self.existing_names = set(existing_names or [])

    def exists(self, name):
        return name in self.existing_names

    def get_available_name(self, name, max_length=None):
        return name

    def get_alternative_name(self, file_root, file_ext):
        return f'{file_root}_sdmmer4{file_ext}'


class ImageUploadTestCase(TestCase):
    @classmethod
    def setUpTestData(cls):
        # We only need a ContentType with model="rack" for the prefix;
        # this doesn't require creating a Rack object.
        cls.ct_rack = ContentType.objects.get_by_natural_key('dcim', 'rack')

    def _stub_instance(self, object_id=12, name=None):
        """
        Creates a minimal stub for use with image attachment path generation.
        """
        return SimpleNamespace(object_type=self.ct_rack, object_id=object_id, name=name)

    def _bound_instance(self, *, storage, object_id=12, name=None, max_length=100):
        return SimpleNamespace(
            object_type=self.ct_rack,
            object_id=object_id,
            name=name,
            image=SimpleNamespace(field=SimpleNamespace(storage=storage, max_length=max_length)),
        )

    def _second_segment(self, path: str):
        """
        Extracts and returns the portion of the input string after the
        first '/' character.
        """
        return path.split('/', 1)[1]

    def test_windows_fake_path_and_extension_lowercased(self):
        """
        Tests handling of a Windows file path with a fake directory and extension.
        """
        inst = self._stub_instance(name=None)
        path = _build_image_attachment_path(inst, r'C:\fake_path\MyPhoto.JPG')
        # Base directory and single-level path
        seg2 = self._second_segment(path)
        self.assertTrue(path.startswith('image-attachments/rack_12_'))
        self.assertNotIn('/', seg2, 'should not create nested directories')
        # Extension from the uploaded file, lowercased
        self.assertTrue(seg2.endswith('.jpg'))

    def test_name_with_slashes_is_flattened_no_subdirectories(self):
        """
        Tests that a name with slashes is flattened and does not
        create subdirectories.
        """
        inst = self._stub_instance(name='5/31/23')
        path = _build_image_attachment_path(inst, 'image.png')
        seg2 = self._second_segment(path)
        self.assertTrue(seg2.startswith('rack_12_'))
        self.assertNotIn('/', seg2)
        self.assertNotIn('\\', seg2)
        self.assertTrue(seg2.endswith('.png'))

    def test_name_with_backslashes_is_flattened_no_subdirectories(self):
        """
        Tests that a name including backslashes is correctly flattened
        into a single directory name without creating subdirectories.
        """
        inst = self._stub_instance(name=r'5\31\23')
        path = _build_image_attachment_path(inst, 'image_name.png')

        seg2 = self._second_segment(path)
        self.assertTrue(seg2.startswith('rack_12_'))
        self.assertNotIn('/', seg2)
        self.assertNotIn('\\', seg2)
        self.assertTrue(seg2.endswith('.png'))

    def test_prefix_format_is_as_expected(self):
        """
        Tests the output path format generated by the `image_upload` function.
        """
        inst = self._stub_instance(object_id=99, name='label')
        path = _build_image_attachment_path(inst, 'a.webp')
        # The second segment must begin with "rack_99_"
        seg2 = self._second_segment(path)
        self.assertTrue(seg2.startswith('rack_99_'))
        self.assertTrue(seg2.endswith('.webp'))

    def test_unsupported_file_extension(self):
        """
        Test that when the file extension is not allowed, the extension
        is omitted.
        """
        inst = self._stub_instance(name='test')
        path = _build_image_attachment_path(inst, 'document.txt')

        seg2 = self._second_segment(path)
        self.assertTrue(seg2.startswith('rack_12_test'))
        self.assertFalse(seg2.endswith('.txt'))
        # When not allowed, no extension should be appended
        self.assertNotRegex(seg2, r'\.txt$')

    def test_instance_name_with_whitespace_and_special_chars(self):
        """
        Test that an instance name with leading/trailing whitespace and
        special characters is sanitized properly.
        """
        # Suppose the instance name has surrounding whitespace and
        # extra slashes.
        inst = self._stub_instance(name='  my/complex\\name  ')
        path = _build_image_attachment_path(inst, 'irrelevant.png')

        # The output should be flattened and sanitized.
        # We expect the name to be transformed into a valid filename without
        # path separators.
        seg2 = self._second_segment(path)
        self.assertNotIn(' ', seg2)
        self.assertNotIn('/', seg2)
        self.assertNotIn('\\', seg2)
        self.assertTrue(seg2.endswith('.png'))

    def test_separator_variants_with_subTest(self):
        """
        Tests that both forward slash and backslash in file paths are
        handled consistently by the `image_upload` function and
        processed into a sanitized uniform format.
        """
        for name in ['2025/09/12', r'2025\09\12']:
            with self.subTest(name=name):
                inst = self._stub_instance(name=name)
                path = _build_image_attachment_path(inst, 'x.jpeg')
                seg2 = self._second_segment(path)
                self.assertTrue(seg2.startswith('rack_12_'))
                self.assertNotIn('/', seg2)
                self.assertNotIn('\\', seg2)
                self.assertTrue(seg2.endswith('.jpeg') or seg2.endswith('.jpg'))

    def test_fallback_on_suspicious_file_operation(self):
        """
        Test that when default_storage.get_valid_name() raises a
        SuspiciousFileOperation, the fallback default is used.
        """
        inst = self._stub_instance(name=' ')
        path = _build_image_attachment_path(inst, 'sample.png')
        # Expect the fallback name 'unnamed' to be used.
        self.assertIn('unnamed', path)
        self.assertTrue(path.startswith('image-attachments/rack_12_'))

    def test_image_upload_preserves_original_name_when_available(self):
        inst = self._bound_instance(
            storage=OverwriteStyleStorage(),
            name='action-buttons',
        )

        path = image_upload(inst, 'action-buttons.png')

        self.assertEqual(path, 'image-attachments/rack_12_action-buttons.png')

    def test_image_upload_uses_base_collision_handling_with_overwrite_style_storage(self):
        inst = self._bound_instance(
            storage=OverwriteStyleStorage(existing_names={'image-attachments/rack_12_action-buttons.png'}),
            name='action-buttons',
        )

        path = image_upload(inst, 'action-buttons.png')

        self.assertEqual(
            path,
            'image-attachments/rack_12_action-buttons_sdmmer4.png',
        )

    def test_image_field_generate_filename_uses_image_upload_collision_handling(self):
        field = ImageAttachment._meta.get_field('image')
        instance = ImageAttachment(
            object_type=self.ct_rack,
            object_id=12,
        )

        with patch.object(
            field,
            'storage',
            OverwriteStyleStorage(existing_names={'image-attachments/rack_12_action-buttons.png'}),
        ):
            path = field.generate_filename(instance, 'action-buttons.png')

        self.assertEqual(
            path,
            'image-attachments/rack_12_action-buttons_sdmmer4.png',
        )