from django.core.files.base import ContentFile from django.db import models, IntegrityError from django_softdelete.models import SoftDeleteModel from inventory.models import Item def hash_upload(instance, filename): return f"{instance.hash[:2]}/{instance.hash[2:4]}/{instance.hash[4:6]}/{instance.hash[6:]}" class FileManager(models.Manager): def get_or_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') try: return self.get(hash=kwargs['hash']), False except self.model.DoesNotExist: obj = super().create(**kwargs) obj.file.save(content=kwargs['file'], name=kwargs['hash']) 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') if not self.filter(hash=kwargs['hash']).exists(): obj = super().create(**kwargs) obj.file.save(content=kwargs['file'], name=kwargs['hash']) return obj else: raise IntegrityError('File with this hash already exists') class AbstractFile(models.Model): created_at = models.DateTimeField(blank=True, null=True) updated_at = models.DateTimeField(blank=True, null=True) deleted_at = models.DateTimeField(blank=True, null=True) file = models.FileField(upload_to=hash_upload) mime_type = models.CharField(max_length=255, null=False, blank=False) hash = models.CharField(max_length=64, null=False, blank=False, unique=True) objects = FileManager() def save(self, *args, **kwargs): from django.utils import timezone if not self.created_at: self.created_at = timezone.now() self.updated_at = timezone.now() super().save(*args, **kwargs) class Meta: abstract = True class File(AbstractFile): item = models.ForeignKey(Item, models.CASCADE, db_column='iid', null=True, blank=True, related_name='files') def __str__(self): return self.hash