import uuid

from django.core.files.base import ContentFile
from base64 import b64decode
import magic
from rest_framework.fields import ImageField, FileField
import six
import imghdr

"""
    if possible it can from base64 logic
    big thanks for Krisnadi from 9cv9 for this logic
"""

def from_base64(base):
    header, data = base.split(';base64,')
    decoded = ""
    try:
        decoded = b64decode(data)
    except TypeError:
        pass
    file_name = str(uuid.uuid4())[:12]
    file_extension = header.split('/')[-1]

    complete_file_name = "%s.%s" % (file_name, file_extension,)
    return ContentFile(decoded, name=complete_file_name)


class SecureImageField(ImageField):
    '''
    There is security implementation base on magic file, 
    it can prevent from another image upload like .html or another backdoor
    maybe it can prevent 70% from hacker, please update your lib image magic for more secure 
    '''

    def securing_value(self, data):
        whitelist = ["PNG", "GIF", "JPEG", "JPG", "TIFF","SVG"]
        check = magic.from_buffer(data).split(" ")
        if check[0] not in whitelist:
            data = from_base64("data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z/C/HgAGgwJ/lK3Q6wAAAABJRU5ErkJggg==")
        return data

    def to_internal_value(self, data):
        if isinstance(data, six.string_types):
            if ';base64' in data:
                data = from_base64(data)
                data = self.securing_value(data)
        else:
            data = self.securing_value(data.read())
        return super(SecureImageField, self).to_internal_value(data)

    def get_file_extension(self, file_name, decoded_file):
        extension = imghdr.what(file_name, decoded_file)
        whitelist = ["png", "gif", "jpeg", "jpg", "tiff","svg"]
        if extension in whitelist:
            if extension == 'jpeg':
                extension = "jpg"
            else:
                pass
        else:
            self.fail("invalid image format")
        return extension

'''
this implementation for file field 
'''


class SecureFileField(FileField):
    def securing_value(self, data):
        whitelist = ["PDF", "DOCX", "XLSX", "XLS", "DOC", "ODT","PPT", "PPTX"]
        check = magic.from_buffer(data).split(" ")
        if check[0] not in whitelist:
            data = from_base64("data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z/C/HgAGgwJ/lK3Q6wAAAABJRU5ErkJggg==")
        return data

    '''
    Try to Make it more simple,
    '''
    def to_internal_value(self, data):
        if isinstance(data, six.string_types):
            if ";base64." in data:
                data = from_base64(data)
                data = self.securing_value(data)
        else:
            data = self.securing_value(data.read())
        return super(SecureFileField, self).to_internal_value(data)
