Commit c9132d8d authored by Lars Kruse's avatar Lars Kruse Committed by Robert

FileManager: determine suitable unique filename for storing attachments

preserve filename extension and add prefix
parent c24348b6
from os import path
import io
import os
import tempfile
from typing import List
import django
from django.db import models
import django.core.files
from features.contributions.signals import ParsedMailAttachment
from imagekit.models import ImageSpecField
from imagekit.processors import ResizeToFit, Transpose
......@@ -14,14 +17,41 @@ from features.contributions import models as contributions
IGNORE_CONTENT_TYPES = {'application/pgp-signature'}
def get_unique_storage_filename(name_template: str, base_dir: str,
default_prefix: str='attachment-') -> str:
""" determine a suitable name for a file to be stored in a directory
The file may not overwrite an existing file and it should keep its original extension.
"""
temp_kwargs = {'dir': base_dir, 'prefix': default_prefix, 'delete': False}
if name_template:
basename, extension = os.path.splitext(os.path.basename(name_template))
if extension:
temp_kwargs['suffix'] = '.' + extension
if basename:
temp_kwargs['prefix'] = basename + '-'
storage_file = tempfile.NamedTemporaryFile(**temp_kwargs)
storage_file.close()
return storage_file.name
class FileManager(models.Manager):
def create_from_message_attachments(self, attachments: List[ParsedMailAttachment],
attached_to):
for attachment in attachments:
if attachment.content_type in IGNORE_CONTENT_TYPES:
continue
file_source = attachment.model_obj if attachment.data is None else attachment.data
f = self.create(file=file_source, filename=attachment.filename)
if attachment.data is not None:
# create the file and reference it
filename = get_unique_storage_filename(attachment.filename,
File.file.field.storage.base_location)
f = self.create()
f.file.save(os.path.basename(filename),
django.core.files.File(io.BytesIO(attachment.data)))
else:
file_source = attachment.model_obj
f = self.create(file=file_source, filename=attachment.filename)
contributions.Contribution.objects.create(
container_id=attached_to.container_id,
container_type=attached_to.container_type,
......@@ -57,7 +87,7 @@ class File(core.models.Model):
@property
def display_name(self):
return self.filename or path.basename(self.file.name)
return self.filename or os.path.basename(self.file.name)
def __str__(self):
return self.display_name
......@@ -66,13 +96,13 @@ class File(core.models.Model):
return self.file.url
def is_image(self):
root, ext = path.splitext(self.file.name.lower())
root, ext = os.path.splitext(self.file.name.lower())
return ext in ('.svg', '.apng', '.png', '.gif', '.jpg', '.jpeg')
def is_video(self):
root, ext = path.splitext(self.file.name.lower())
root, ext = os.path.splitext(self.file.name.lower())
return ext in ('.mp4', '.ogg', '.webm')
def is_audio(self):
root, ext = path.splitext(self.file.name.lower())
root, ext = os.path.splitext(self.file.name.lower())
return ext in ('.opus', '.ogg', '.aac', '.mp3', '.flac', '.m4a')
import contextlib
import os
import shutil
import unittest
import tempfile
import django
import django_mailbox
import core
import core.tests
from features.associations import models as associations
from features.files.models import get_unique_storage_filename
from features.images import tests as images
from features.memberships import test_mixins as memberships
......@@ -93,6 +97,42 @@ class Gestalt(images.ImageMixin, memberships.AuthenticatedMemberMixin, core.test
self.assertOk(url=self.get_group_file_url())
class FilenameGenerator(unittest.TestCase):
def setUp(self):
self._base_dir = tempfile.mkdtemp()
# Manage a set of created filenames. Otherwise the 'tearDown' method preempts the context
# cleanup and thus fails to remove the temporary directory.
self._filenames = set()
super().setUp()
def tearDown(self):
for filename in self._filenames:
os.unlink(filename)
os.rmdir(self._base_dir)
@contextlib.contextmanager
def get_unique_filename(self, name_template, default_prefix=None):
filename = get_unique_storage_filename(name_template, self._base_dir,
default_prefix=default_prefix)
self._filenames.add(filename)
yield filename
os.unlink(filename)
self._filenames.remove(filename)
def test_unique_name_generator(self):
with self.get_unique_filename('foo.bar.baz', 'nom') as filename:
self.assertEqual(os.path.dirname(filename), self._base_dir)
self.assertTrue(os.path.basename(filename))
self.assertTrue(os.path.basename(filename).startswith('foo.bar-'), filename)
self.assertTrue(os.path.basename(filename).endswith('.baz'), filename)
with self.get_unique_filename(None, 'foo-') as filename:
self.assertEqual(os.path.dirname(filename), self._base_dir)
self.assertTrue(os.path.basename(filename))
self.assertTrue(os.path.basename(filename).startswith('foo-'), filename)
self.assertNotIn('.', os.path.basename(filename))
class TestUrls(core.tests.Test):
def test_files_404(self):
r = self.client.get(self.get_url('create-group-file', 'non-existent'))
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment