From 54a2176d7cdeb45405a30481e8ec92ff15d5cdd7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@stuvel.eu>
Date: Tue, 3 May 2016 10:57:56 +0200
Subject: [PATCH] Override browser-given content type for certain types.

If the browser gives us audio/xxx, video/xxx or image/xxx, we just use
it. Otherwise the mimetypes package is used to guess the mime type
based on the file's extension.

Also ensures the content type of the file is updated in MongoDB. Does
NOT update any nodes that link to this file.
---
 pillar/application/modules/file_storage.py | 41 +++++++++++++++++++++-
 1 file changed, 40 insertions(+), 1 deletion(-)

diff --git a/pillar/application/modules/file_storage.py b/pillar/application/modules/file_storage.py
index 17c41885..b1b03952 100644
--- a/pillar/application/modules/file_storage.py
+++ b/pillar/application/modules/file_storage.py
@@ -456,6 +456,42 @@ def create_file_doc(name, filename, content_type, length, project, backend='gcs'
     return file_doc
 
 
+def override_content_type(uploaded_file):
+    """Overrides the content type based on file extensions.
+
+    :param uploaded_file: file from request.files['form-key']
+    :type uploaded_file: werkzeug.datastructures.FileStorage
+    """
+
+    import mimetypes
+
+    # Possibly use the browser-provided mime type
+    mimetype = uploaded_file.mimetype
+    if '/' in mimetype:
+        mimecat = mimetype.split('/')[0]
+        if mimecat in {'video', 'audio', 'image'}:
+            # The browser's mime type is probably ok, just use it.
+            return
+
+    # Add our own extensions to the mimetypes package
+    if not mimetypes.inited:
+        mimetypes.init()
+        mimetypes.add_type('application/x-blend', '.blend')
+
+    # And then use it to set the mime type.
+    (mimetype, encoding) = mimetypes.guess_type(uploaded_file.filename)
+
+    # Only override the mime type if we can detect it, otherwise just
+    # keep whatever the browser gave us.
+    if mimetype:
+        # content_type property can't be set directly
+        uploaded_file.headers['content-type'] = mimetype
+
+        # It has this, because we used uploaded_file.mimetype earlier this function.
+        del uploaded_file._parsed_content_type
+
+
+
 @file_storage.route('/stream/<string:project_id>', methods=['POST', 'OPTIONS'])
 @require_login(require_roles={u'subscriber', u'admin', u'demo'})
 def stream_to_gcs(project_id):
@@ -473,6 +509,8 @@ def stream_to_gcs(project_id):
 
     file_id, internal_fname, status = create_file_doc_for_upload(project['_id'], uploaded_file)
 
+    override_content_type(uploaded_file)
+
     if uploaded_file.content_type.startswith('image/'):
         # We need to do local thumbnailing, so we have to write the stream
         # both to Google Cloud Storage and to local storage.
@@ -502,7 +540,8 @@ def stream_to_gcs(project_id):
     update_file_doc(file_id,
                     status='queued_for_processing',
                     file_path=internal_fname,
-                    length=blob.size)
+                    length=blob.size,
+                    content_type=uploaded_file.mimetype)
 
     process_file(gcs, file_id, local_file)