check file contents; work with file extensions

htdocs/lib/settings/ImageUpload.js

@@ -24,8 +24,8 @@ $.fn.imageUpload = function() {
                         processData: false,
                         contentType: 'application/octet-stream',
                     }).done(function(data){
-                        $input.val(data.uuid);
-                        $img.prop('src', "/imageupload?id=" + id + "&uuid=" + data.uuid);
+                        $input.val(data.file);
+                        $img.prop('src', '/imageupload?file=' + data.file);
                     }).always(function(){
                         $uploadButton.prop('disabled', false);
                     });

owrx/controllers/assets.py

@@ -78,7 +78,7 @@ class AssetsController(GzipMixin, ModificationAwareController, metaclass=ABCMeta
             f.close()
 
             if content_type is None:
-                (content_type, encoding) = mimetypes.MimeTypes().guess_type(self.getFilePath(file))
+                (content_type, encoding) = mimetypes.guess_type(self.getFilePath(file))
             self.send_response(data, content_type=content_type, last_modified=modified, max_age=3600)
         except FileNotFoundError:
             self.send_response("file not found", code=404)
@@ -96,7 +96,8 @@ class OwrxAssetsController(AssetsController):
         }
         if file in mappedFiles and ("mapped" not in self.request.query or self.request.query["mapped"][0] != "false"):
             config = CoreConfig()
-            user_file = config.get_data_directory() + "/" + mappedFiles[file]
+            for ext in ["png", "jpg"]:
+                user_file = "{}/{}.{}".format(config.get_data_directory(), mappedFiles[file], ext)
                 if os.path.exists(user_file) and os.path.isfile(user_file):
                     return user_file
         return pkg_resources.resource_filename("htdocs", file)
@@ -168,7 +169,7 @@ class CompiledAssetsController(GzipMixin, ModificationAwareController):
 
         contents = [self.getContents(f) for f in files]
 
-        (content_type, encoding) = mimetypes.MimeTypes().guess_type(profileName)
+        (content_type, encoding) = mimetypes.guess_type(profileName)
         self.send_response("\n".join(contents), content_type=content_type, last_modified=modified, max_age=3600)
 
     def getContents(self, file):

owrx/controllers/imageupload.py

@@ -8,28 +8,43 @@ import json
 class ImageUploadController(AuthorizationMixin, AssetsController):
     def __init__(self, handler, request, options):
         super().__init__(handler, request, options)
-        self.uuid = request.query["uuid"][0] if "uuid" in request.query else None
-        self.id = request.query["id"][0] if "id" in request.query else None
+        self.file = request.query["file"][0] if "file" in request.query else None
 
     def getFilePath(self, file=None):
-        if self.uuid is None:
-            raise FileNotFoundError("missing uuid")
-        if self.id is None:
-            raise FileNotFoundError("missing id")
-        return "{tmp}/{file}-{uuid}".format(
+        if self.file is None:
+            raise FileNotFoundError("missing filename")
+        return "{tmp}/{file}".format(
             tmp=CoreConfig().get_temporary_directory(),
-            file=self.id,
-            uuid=self.uuid,
+            file=self.file
         )
 
     def indexAction(self):
         self.serve_file(None)
 
+    def _is_png(self, contents):
+        return contents[0:8] == bytes([0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A])
+
+    def _is_jpg(self, contents):
+        return contents[0:3] == bytes([0xFF, 0xD8, 0xFF])
+
     def processImage(self):
-        self.uuid = uuid.uuid4().hex
+        if "id" not in self.request.query:
+            self.send_response("{}", content_type="application/json", code=400)
         # TODO: limit file size
-        # TODO: check image mime type, if possible
         contents = self.get_body()
-        with open(self.getFilePath(), 'wb') as f:
+        filetype = None
+        if self._is_png(contents):
+            filetype = "png"
+        if self._is_jpg(contents):
+            filetype = "jpg"
+        if filetype is None:
+            self.send_response("{}", content_type="application/json", code=400)
+        else:
+            self.file = "{id}-{uuid}.{ext}".format(
+                id=self.request.query["id"][0],
+                uuid=uuid.uuid4().hex,
+                ext=filetype,
+            )
+            with open(self.getFilePath(), "wb") as f:
                 f.write(contents)
-        self.send_response(json.dumps({"uuid": self.uuid}), content_type="application/json")
+            self.send_response(json.dumps({"file": self.file}), content_type="application/json")

owrx/controllers/settings.py

@@ -374,11 +374,21 @@ class GeneralSettingsController(AuthorizationMixin, WebpageController):
     def handle_image(self, data, image_id):
         if image_id in data:
             config = CoreConfig()
-            data_file = config.get_data_directory() + "/" + image_id
             if data[image_id] == "restore":
-                os.unlink(data_file)
+                # remove all possible file extensions
+                for ext in ["png", "jpg"]:
+                    try:
+                        os.unlink("{}/{}.{}".format(config.get_data_directory(), image_id, ext))
+                    except FileNotFoundError:
+                        pass
             elif data[image_id]:
-                temporary_file = "{}/{}-{}".format(config.get_temporary_directory(), image_id, data[image_id])
+                if not data[image_id].startswith(image_id):
+                    logger.warning("invalid file name: %s", data[image_id])
+                else:
+                    # get file extension (luckily, all options are three characters long)
+                    ext = data[image_id][-3:]
+                    data_file = "{}/{}.{}".format(config.get_data_directory(), image_id, ext)
+                    temporary_file = "{}/{}".format(config.get_temporary_directory(), data[image_id])
                     shutil.copy(temporary_file, data_file)
             del data[image_id]
             # remove any accumulated temporary files on save