From 30bd679de8e7efdbeae234613f8cf5c8c48cf5d4 Mon Sep 17 00:00:00 2001
From: Brett Williams <blw1138@mac.com>
Date: Mon, 5 Jun 2023 15:52:49 -0500
Subject: [PATCH] Process and unzip uploaded zip files

---
 lib/server/job_server.py | 34 +++++++++++++++++++++++++++++++---
 1 file changed, 31 insertions(+), 3 deletions(-)

diff --git a/lib/server/job_server.py b/lib/server/job_server.py
index d99d7bb..f506a90 100755
--- a/lib/server/job_server.py
+++ b/lib/server/job_server.py
@@ -7,6 +7,7 @@ import shutil
 import socket
 import threading
 import time
+import zipfile
 from datetime import datetime
 from zipfile import ZipFile
 
@@ -297,7 +298,7 @@ def add_job_handler():
             job_dir = os.path.join(server.config['UPLOAD_FOLDER'], '_'.join(
                 [datetime.now().strftime("%Y.%m.%d_%H.%M.%S"),
                  jobs_list[0]['renderer'],
-                 uploaded_file.filename]))
+                 os.path.splitext(uploaded_file.filename)[0]]))
             os.makedirs(job_dir, exist_ok=True)
 
             upload_dir = os.path.join(job_dir, 'source')
@@ -350,8 +351,7 @@ def add_job(job_params, remove_job_dir_on_failure=False):
     args = job_params.get('args', {})
     client = job_params.get('client', None) or RenderQueue.hostname
     force_start = job_params.get('force_start', False)
-    custom_id = None
-    job_dir = None
+    job_dir = None  # todo: I dont think this gets set anywhere
 
     # check for minimum render requirements
     if None in [renderer, input_path, output_path]:
@@ -359,6 +359,34 @@ def add_job(job_params, remove_job_dir_on_failure=False):
         logger.error(err_msg)
         return {'error': err_msg, 'code': 400}
 
+    # process uploaded zip files
+    if '.zip' in input_path:
+        zip_path = input_path
+        work_path = os.path.dirname(zip_path)
+        try:
+            with zipfile.ZipFile(zip_path, 'r') as myzip:
+                myzip.extractall(os.path.dirname(zip_path))
+
+            project_files = [x for x in os.listdir(work_path) if os.path.isfile(os.path.join(work_path, x))]
+            project_files = [x for x in project_files if '.zip' not in x]
+            supported_exts = RenderWorkerFactory.class_for_name(renderer).engine.supported_extensions
+            if supported_exts:
+                project_files = [file for file in project_files if any(file.endswith(ext) for ext in supported_exts)]
+
+            if len(project_files) != 1:  # we have to narrow down to 1 main project file, otherwise error
+                return {'error': f'Cannot find valid project file in {os.path.basename(zip_path)}'}, 400
+
+            extracted_project_path = os.path.join(work_path, project_files[0])
+            logger.info(f"Extracted zip file to {extracted_project_path}")
+            input_path = extracted_project_path
+        except (zipfile.BadZipFile, zipfile.LargeZipFile) as e:
+            err_msg = f"Error processing zip file: {e}"
+            logger.error(err_msg)
+            return {'error': err_msg, 'code': 400}
+        finally:
+            os.remove(zip_path)
+            logger.info(f"Removed zip: {zip_path}")
+
     # local renders
     if client == RenderQueue.hostname:
         logger.info(f"Creating job locally - {name if name else input_path}")