scripts/project-tools/update_addons.py

@@ -1,99 +1,96 @@
 #!/usr/bin/env python3
 
 import hashlib
-import os
-import pathlib
 from pathlib import Path
 from urllib.request import urlretrieve
-import tempfile
 import sys
-import zipfile
-import shutil
+import requests
+import glob
+import os
 
 
-def main():
-    download_path = Path(tempfile.mkdtemp(prefix="blender_studio_addons_download" + "_"))
-    zip_filepath = download_path.joinpath("blender_studio_add-ons_latest.zip")
-    checksum_path = download_path.joinpath("blender_studio_add-ons_latest.zip.sha256")
-    extract_path = download_path.joinpath("extract")
-    extract_path.mkdir()
-    print(f"Downloading {zip_filepath.name}......", end="")
-    urlretrieve(
-        "https://projects.blender.org/studio/blender-studio-pipeline/releases/download/latest/blender_studio_add-ons_latest.zip",
-        str(zip_filepath),
-    )
-    print("Complete")
-    print(f"Downloading {checksum_path.name}......", end="")
-    urlretrieve(
-        "https://projects.blender.org/studio/blender-studio-pipeline/releases/download/latest/blender_studio_add-ons_latest.zip.sha256",
-        str(checksum_path),
-    )
-    print("Complete")
-
-    checksum = generate_checksum(zip_filepath)
-    with open(checksum_path, "r") as f:
-        lines = f.readlines()
-        if lines[0].split(" ")[0].strip() != checksum:
-            print("Download Error: checksum is invalid, please try again")
-            sys.exit(1)
-
-    current_file_folder_path = pathlib.Path(__file__).parent
-    download_folder_path = (current_file_folder_path / "../../shared/artifacts/addons/").resolve()
+def update_blender_studio_addons(download_folder_path: Path):
     if not download_folder_path.exists():
         print(
             f"Ensure script is run out of Project Tools directory {str(download_folder_path)} does not exist"
         )
         sys.exit(1)
-    extract_release_zip(zip_filepath, extract_path, download_folder_path)
-    shutil.rmtree(download_path)
+
+    sha_file = download_folder_path.joinpath("blender_studio_add-ons_latest.zip.sha256")
+    zip_file = download_folder_path.joinpath("blender_studio_add-ons_latest.zip")
+
+    url_sha = "https://projects.blender.org/studio/blender-studio-pipeline/releases/download/latest/blender_studio_add-ons_latest.zip.sha256"
+    url_zip = "https://projects.blender.org/studio/blender-studio-pipeline/releases/download/latest/blender_studio_add-ons_latest.zip"
+
+    # Check current sha and early return if match
+    web_sha = requests.get(url_sha).text.strip().lower()
+    if sha_file.exists() & zip_file.exists():
+        if shasum_matches(zip_file, web_sha):
+            print(f"{zip_file.name} already up to date, canceling update")
+            return
+        else:
+            # Remove current files
+            if sha_file.exists():
+                sha_file.unlink()
+            if zip_file.exists():
+                zip_file.unlink()
+
+    print(f"Downloading {zip_file.name}......", end="")
+    urlretrieve(url_zip, str(zip_file))
+    print("Complete")
+    print(f"Downloading {sha_file.name}......", end="")
+    urlretrieve(url_sha, str(sha_file))
+    print("Complete")
+
+    if not shasum_matches(zip_file, web_sha):
+        print(f"Downloaded file {zip_file.name} does not match its shasum, exiting!")
+        exit(1)
+
     print("Blender Studio Add-Ons Successfully Updated for Current Project")
     print(
-        "These Add-Ons will be copied to your local directory next time you launch Blender via Projet Tools"
+        "Blender Studio Add-Ons will be copied to your local directory next time you launch Blender via Projet Tools"
     )
 
 
-def extract_release_zip(file_path: Path, extract_path: Path, dst_path: Path):
-    """Extracts the contents of a zip file to a destination folder.
-
-    Args:
-        file_path (Path): Path to the zip file to extract.
-        extract_path (Path): Path to extract the contents of the zip file to.
-        dst_path (Path): Destination path to copy the extracted files to.
-    """
-
-    with zipfile.ZipFile(file_path, 'r') as zip_ref:
-        zip_ref.extractall(extract_path)
-
-    try:
-        src_path = [subdir for subdir in Path(extract_path).iterdir()][0]
-    except IndexError:
-        print("The archive %s does not contain any directory" % file_path.name)
-        sys.exit(1)
-
-    for file in src_path.iterdir():
-        print(f"Extacting: {file.name}")
-        original_file = dst_path / file.name
-        if original_file.exists():
-            os.remove(original_file)
-        shutil.move(file, dst_path)
-
-    shutil.rmtree(extract_path)
+def shasum_matches(file, sha_sum):
+    with open(file, "rb") as f:
+        digest = hashlib.file_digest(f, "sha256")
+        return sha_sum.startswith(digest.hexdigest())
 
 
-def generate_checksum(archive_path: Path) -> str:
-    """Generate checksum for archive file.
+def download_file(url, out_folder, filename):
+    print("Downloading: " + url)
+    local_filename = out_folder / filename
 
-    Args:
-        archive_path (Path): Path to archive file to generate checksum for.
+    # TODO Can't check any shasums before downloading so always remove and redownload everything for now
+    prev_downloaded_files = glob.glob(f"{local_filename}*")
+    for file in prev_downloaded_files:
+        os.remove(file)
 
-    Returns:
-        str: Hex digest string of checksum.
-    """
+    # NOTE the stream=True parameter below
+    with requests.get(url, stream=True) as r:
+        r.raise_for_status()
+        with open(local_filename, 'wb') as f:
+            for chunk in r.iter_content(chunk_size=None):
+                if chunk:
+                    f.write(chunk)
 
-    with open(archive_path, 'rb') as file:
-        digest = hashlib.file_digest(file, "sha256")
-    return digest.hexdigest()
+    local_hash_filename = local_filename.with_suffix(".zip.sha256")
+    with open(local_filename, "rb") as f:
+        digest = hashlib.file_digest(f, "sha256")
+        with open(local_hash_filename, "w") as hash_file:
+            hash_file.write(digest.hexdigest())
+
+    return local_filename
 
 
-if __name__ == "__main__":
-    main()
+current_file_folder_path = Path(__file__).parent
+download_folder_path = (current_file_folder_path / "../../shared/artifacts/addons/").resolve()
+update_blender_studio_addons(Path('/data/my_project/shared/artifacts/addons'))
+
+# Customize this script to download add-ons from other sources
+# download_file(
+#    "https://projects.blender.org/studio/blender-studio-pipeline/archive/main.zip",
+#    download_folder_path,
+#    "blender-studio-pipeline-main.zip",
+# )