c3lf-system-3/core/files/models.py

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