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

Address code review: correctness fixes and tests for clean() re-encoding

- Move extension allowlist check before PillowImage.open() so SVG/unknown
  formats get a clear error rather than a generic 'unable to process' message.
- Replace hasattr(self.image, 'file') guard with not self.image._committed so
  already-stored images are not unnecessarily re-encoded on every save.
- Replace bare BytesIO assignment with ContentFile so Django's pre_save()
  reliably writes the sanitised content to storage.
- Tighten except clause to (OSError, UnidentifiedImageError, DecompressionBombError).
- Use ImageSequence.Iterator for animated GIF re-encoding to preserve all
  frames; pass duration/loop metadata through.
- Add four model tests covering polyglot stripping, animated GIF frame
  preservation, disallowed extension rejection, and committed-image skip.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Brian Tiemann 3 недель назад
Родитель
Сommit
0227be7ebe
2 измененных файлов с 118 добавлено и 16 удалено
  1. 33 15
      netbox/extras/models/models.py
  2. 85 1
      netbox/extras/tests/test_models.py

+ 33 - 15
netbox/extras/models/models.py

@@ -6,6 +6,7 @@ from pathlib import Path
 from django.conf import settings
 from django.contrib.contenttypes.fields import GenericForeignKey, GenericRelation
 from django.contrib.postgres.fields import ArrayField
+from django.core.files.base import ContentFile
 from django.core.validators import ValidationError
 from django.db import models
 from django.urls import reverse
@@ -743,30 +744,47 @@ class ImageAttachment(ChangeLoggedModel):
                 _("Image attachments cannot be assigned to this object type ({type}).").format(type=self.object_type)
             )
 
-        # Re-encode the image through Pillow to strip any embedded non-image
-        # payloads (e.g. polyglot HTML/PNG files). This must be done after the
-        # ImageField's own Pillow dimension check has already run.
-        if self.image and hasattr(self.image, 'file'):
+        # Re-encode new image uploads through Pillow to strip any embedded
+        # non-image payloads (e.g. polyglot HTML/PNG files).
+        # Only process uncommitted (new) uploads; skip already-stored images.
+        if self.image and not self.image._committed:
             try:
-                self.image.file.seek(0)
-                img = PillowImage.open(self.image.file)
-                img_format = img.format or 'PNG'
-                # Validate extension against allowed formats
+                # Validate extension before attempting to open with Pillow so
+                # that disallowed types (e.g. SVG) get a clear error message
+                # rather than a generic "unable to process" one.
                 ext = Path(self.image.name).suffix.lstrip('.').lower()
                 if ext not in IMAGE_ATTACHMENT_IMAGE_FORMATS:
                     raise ValidationError(
                         _("Unsupported image format: {ext}").format(ext=ext)
                     )
-                # Re-encode to a new buffer; this strips any non-image trailer data
-                buf = io.BytesIO()
-                img.save(buf, format=img_format)
-                buf.seek(0)
-                self.image.file = buf
                 self.image.file.seek(0)
+                img = PillowImage.open(self.image.file)
+                img_format = img.format or 'PNG'
+                # Re-encode into a fresh buffer; any non-image trailer bytes
+                # are discarded. For animated GIFs, iterate through all frames
+                # via ImageSequence so animation is preserved.
+                buf = io.BytesIO()
+                if img_format == 'GIF':
+                    from PIL import ImageSequence  # noqa: PLC0415
+                    frames = [frame.copy() for frame in ImageSequence.Iterator(img)]
+                    if frames:
+                        frames[0].save(
+                            buf, format='GIF', save_all=True,
+                            append_images=frames[1:] if len(frames) > 1 else [],
+                            duration=img.info.get('duration', 100),
+                            loop=img.info.get('loop', 0),
+                        )
+                    else:
+                        img.save(buf, format='GIF')
+                else:
+                    img.save(buf, format=img_format)
+                # Replace the in-memory upload with the sanitised content so
+                # Django writes the clean version to storage on model save.
+                self.image.file = ContentFile(buf.getvalue(), name=self.image.name)
             except ValidationError:
                 raise
-            except Exception:
-                raise ValidationError(_("Unable to process the uploaded image."))
+            except (OSError, PillowImage.UnidentifiedImageError, PillowImage.DecompressionBombError) as e:
+                raise ValidationError(_("Unable to process the uploaded image.")) from e
 
     def delete(self, *args, **kwargs):
 

+ 85 - 1
netbox/extras/tests/test_models.py

@@ -189,7 +189,91 @@ class ImageAttachmentTestCase(TestCase):
         self.assertEqual(first.filename, 'action-buttons.png')
         self.assertEqual(second.filename, 'action-buttons_sdmmer4.png')
 
-        self.assertCountEqual(storage.files.keys(), {base_name, suffixed_name})
+    # ------------------------------------------------------------------
+    # Security: polyglot upload sanitisation (SR-001 / VM-326)
+    # ------------------------------------------------------------------
+
+    def _make_polyglot_png(self, filename, payload=b'<script>alert(1)</script>'):
+        """Return a SimpleUploadedFile that is a valid PNG with extra bytes appended."""
+        buf = io.BytesIO()
+        Image.new('RGB', (1, 1), color='red').save(buf, format='PNG')
+        # Append a non-image payload after the IEND chunk
+        polyglot = buf.getvalue() + payload
+        return SimpleUploadedFile(name=filename, content=polyglot, content_type='image/png')
+
+    def _make_animated_gif(self, filename, frames=2):
+        """Return a SimpleUploadedFile containing a minimal animated GIF."""
+        buf = io.BytesIO()
+        # Use distinct RGB colours and a size > 1×1 so GIF optimisation
+        # does not collapse identical frames into a single frame.
+        images = [Image.new('RGB', (2, 2), color=(i * 80, 0, 0)) for i in range(frames)]
+        images[0].save(buf, format='GIF', save_all=True, append_images=images[1:], loop=0, duration=100)
+        return SimpleUploadedFile(name=filename, content=buf.getvalue(), content_type='image/gif')
+
+    def _make_attachment(self, uploaded_file):
+        return ImageAttachment(
+            object_type=self.ct_site,
+            object_id=self.site.pk,
+            image=uploaded_file,
+        )
+
+    def test_clean_strips_polyglot_payload_from_png(self):
+        """clean() re-encodes PNG uploads, stripping any non-image trailer data."""
+        payload = b'<script>alert(1)</script>'
+        ia = self._make_attachment(self._make_polyglot_png('test.png', payload=payload))
+        # The raw upload contains the payload
+        self.assertIn(payload, ia.image.file.read())
+        ia.image.file.seek(0)
+
+        ia.full_clean()
+
+        # After clean(), the sanitised ContentFile must not contain the payload
+        sanitised = ia.image.file.read()
+        self.assertNotIn(payload, sanitised)
+        # But it must still be a valid PNG
+        buf = io.BytesIO(sanitised)
+        img = Image.open(buf)
+        self.assertEqual(img.format, 'PNG')
+
+    def test_clean_preserves_animated_gif_frames(self):
+        """clean() re-encodes GIF uploads while keeping all animation frames."""
+        n_frames = 3
+        ia = self._make_attachment(self._make_animated_gif('anim.gif', frames=n_frames))
+        ia.full_clean()
+
+        sanitised = ia.image.file.read()
+        buf = io.BytesIO(sanitised)
+        img = Image.open(buf)
+        self.assertEqual(img.format, 'GIF')
+        # All frames must be preserved
+        frame_count = 0
+        try:
+            while True:
+                frame_count += 1
+                img.seek(img.tell() + 1)
+        except EOFError:
+            pass
+        self.assertEqual(frame_count, n_frames)
+
+    def test_clean_rejects_disallowed_extension(self):
+        """clean() raises ValidationError for file extensions not in IMAGE_ATTACHMENT_IMAGE_FORMATS."""
+        svg_content = b'<svg xmlns="http://www.w3.org/2000/svg"><script>alert(1)</script></svg>'
+        ia = self._make_attachment(
+            SimpleUploadedFile(name='evil.svg', content=svg_content, content_type='image/svg+xml')
+        )
+        with self.assertRaises(ValidationError) as cm:
+            ia.clean()
+        self.assertIn('svg', str(cm.exception).lower())
+
+    def test_clean_skips_reencode_for_existing_stored_image(self):
+        """clean() does not re-encode already-committed (stored) images."""
+        ia = self._make_attachment(self._uploaded_png('existing.png'))
+        # Mark as committed (simulates an already-saved instance fetched from DB)
+        ia.image._committed = True
+        original_file = ia.image.file
+        # Call model clean() directly; the re-encoding block must be skipped
+        ia.clean()
+        self.assertIs(ia.image.file, original_file)
 
 
 class TableConfigTestCase(TestCase):