From 0ae66a2608ba1dd2ad896fa906583360251d3b0c Mon Sep 17 00:00:00 2001
From: grossmj <grossmj@gns3.net>
Date: Thu, 7 Mar 2019 17:05:32 +0700
Subject: [PATCH] Support selecting a compression type when exporting a
 project.

---
 gns3server/controller/snapshot.py             |  1 +
 .../api/controller/appliance_handler.py       |  2 +-
 .../api/controller/project_handler.py         | 39 +++++++++----------
 3 files changed, 20 insertions(+), 22 deletions(-)

diff --git a/gns3server/controller/snapshot.py b/gns3server/controller/snapshot.py
index 20ed326c..72be81fd 100644
--- a/gns3server/controller/snapshot.py
+++ b/gns3server/controller/snapshot.py
@@ -96,6 +96,7 @@ class Snapshot:
         try:
             begin = time.time()
             with tempfile.TemporaryDirectory() as tmpdir:
+                # Do not compress the snapshots
                 with aiozipstream.ZipFile(compression=zipfile.ZIP_STORED) as zstream:
                     await export_project(zstream, self._project, tmpdir, keep_compute_id=True, allow_all_nodes=True)
                     async with aiofiles.open(self.path, 'wb') as f:
diff --git a/gns3server/handlers/api/controller/appliance_handler.py b/gns3server/handlers/api/controller/appliance_handler.py
index 20eaeb36..178279e6 100644
--- a/gns3server/handlers/api/controller/appliance_handler.py
+++ b/gns3server/handlers/api/controller/appliance_handler.py
@@ -36,7 +36,7 @@ class ApplianceHandler:
     async def list_appliances(request, response):
 
         controller = Controller.instance()
-        if request.query.get("update", "no") == "yes":
+        if request.query.get("update", "no").lower() == "yes":
             await controller.appliance_manager.download_appliances()
         controller.appliance_manager.load_appliances()
         response.json([c for c in controller.appliance_manager.appliances.values()])
diff --git a/gns3server/handlers/api/controller/project_handler.py b/gns3server/handlers/api/controller/project_handler.py
index ab703fa0..8ba36697 100644
--- a/gns3server/handlers/api/controller/project_handler.py
+++ b/gns3server/handlers/api/controller/project_handler.py
@@ -305,12 +305,25 @@ class ProjectHandler:
 
         controller = Controller.instance()
         project = await controller.get_loaded_project(request.match_info["project_id"])
+        if request.query.get("include_images", "no").lower() == "yes":
+            include_images = True
+        else:
+            include_images = False
+        compression_query = request.query.get("compression", "zip").lower()
+        if compression_query == "zip":
+            compression = zipfile.ZIP_DEFLATED
+        elif compression_query == "none":
+            compression = zipfile.ZIP_STORED
+        elif compression_query == "bzip2":
+            compression = zipfile.ZIP_BZIP2
+        elif compression_query == "lzma":
+            compression = zipfile.ZIP_LZMA
 
         try:
             begin = time.time()
             with tempfile.TemporaryDirectory() as tmp_dir:
-                with aiozipstream.ZipFile(compression=zipfile.ZIP_DEFLATED) as zstream:
-                    await export_project(zstream, project, tmp_dir, include_images=bool(int(request.query.get("include_images", "0"))))
+                with aiozipstream.ZipFile(compression=compression) as zstream:
+                    await export_project(zstream, project, tmp_dir, include_images=include_images)
 
                     # We need to do that now because export could failed and raise an HTTP error
                     # that why response start need to be the later possible
@@ -430,23 +443,7 @@ class ProjectHandler:
             raise aiohttp.web.HTTPForbidden()
         path = os.path.join(project.path, path)
 
-        response.content_type = "application/octet-stream"
-        response.set_status(200)
-        response.enable_chunked_encoding()
-
-        try:
-            with open(path, "rb") as f:
-                await response.prepare(request)
-                while True:
-                    data = f.read(CHUNK_SIZE)
-                    if not data:
-                        break
-                    await response.write(data)
-
-        except FileNotFoundError:
-            raise aiohttp.web.HTTPNotFound()
-        except PermissionError:
-            raise aiohttp.web.HTTPForbidden()
+        await response.stream_file(path)
 
     @Route.post(
         r"/projects/{project_id}/files/{path:.+}",
@@ -475,7 +472,7 @@ class ProjectHandler:
         response.set_status(200)
 
         try:
-            with open(path, 'wb+') as f:
+            async with aiofiles.open(path, 'wb+') as f:
                 while True:
                     try:
                         chunk = await request.content.read(CHUNK_SIZE)
@@ -483,7 +480,7 @@ class ProjectHandler:
                         raise aiohttp.web.HTTPRequestTimeout(text="Timeout when writing to file '{}'".format(path))
                     if not chunk:
                         break
-                    f.write(chunk)
+                    await f.write(chunk)
         except FileNotFoundError:
             raise aiohttp.web.HTTPNotFound()
         except PermissionError: