diff --git a/core/files/models.py b/core/files/models.py index 33a6265..a8eb775 100644 --- a/core/files/models.py +++ b/core/files/models.py @@ -1,6 +1,5 @@ from django.core.files.base import ContentFile from django.db import models, IntegrityError -from django_softdelete.models import SoftDeleteModel from inventory.models import Item @@ -10,7 +9,8 @@ def hash_upload(instance, filename): class FileManager(models.Manager): - def get_or_create(self, **kwargs): + + def __file_data_helper(self, **kwargs): if 'data' in kwargs and type(kwargs['data']) == str: import base64 from hashlib import sha256 @@ -31,6 +31,10 @@ class FileManager(models.Manager): pass else: raise ValueError('data must be a base64 encoded string or file and hash must be provided') + return kwargs + + def get_or_create(self, **kwargs): + kwargs = self.__file_data_helper(**kwargs) try: return self.get(hash=kwargs['hash']), False except self.model.DoesNotExist: @@ -39,26 +43,7 @@ class FileManager(models.Manager): return obj, True def create(self, **kwargs): - if 'data' in kwargs and type(kwargs['data']) == str: - import base64 - from hashlib import sha256 - raw = kwargs['data'] - if not raw.startswith('data:'): - raise ValueError('data must be a base64 encoded string or file and hash must be provided') - raw = raw.split(';base64,') - if len(raw) != 2: - raise ValueError('data must be a base64 encoded string or file and hash must be provided') - mime_type = raw[0].split(':')[1] - content = base64.b64decode(raw[1], validate=True) - kwargs.pop('data') - content_hash = sha256(content).hexdigest() - kwargs['file'] = ContentFile(content, content_hash) - kwargs['hash'] = content_hash - kwargs['mime_type'] = mime_type - elif 'file' in kwargs and 'hash' in kwargs and type(kwargs['file']) == ContentFile and 'mime_type' in kwargs: - pass - else: - raise ValueError('data must be a base64 encoded string or file and hash must be provided') + kwargs = self.__file_data_helper(**kwargs) if not self.filter(hash=kwargs['hash']).exists(): obj = super().create(**kwargs) obj.file.save(content=kwargs['file'], name=kwargs['hash'])