From 08c17c83d440c416fdf87704e9b678b5f6291556 Mon Sep 17 00:00:00 2001
From: Fabian Dill <Berserker66@users.noreply.github.com>
Date: Sun, 5 Mar 2023 14:10:05 +0100
Subject: [PATCH] Setup: auto download SNI (#1312)

Co-authored-by: black-sliver <59490463+black-sliver@users.noreply.github.com>
---
 setup.py | 74 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 74 insertions(+)

diff --git a/setup.py b/setup.py
index 509981da..df2d43ec 100644
--- a/setup.py
+++ b/setup.py
@@ -7,6 +7,11 @@ import sys
 import sysconfig
 import typing
 import zipfile
+import urllib.request
+import io
+import json
+import threading
+import platform
 from collections.abc import Iterable
 from hashlib import sha3_512
 from pathlib import Path
@@ -49,6 +54,68 @@ apworlds: set = {
     "Timespinner",
 }
 
+
+def download_SNI():
+    print("Updating SNI")
+    machine_to_go = {
+        "x86_64": "amd64",
+        "aarch64": "arm64",
+        "armv7l": "arm"
+    }
+    platform_name = platform.system().lower()
+    machine_name = platform.machine().lower()
+    # force amd64 on macos until we have universal2 sni, otherwise resolve to GOARCH
+    machine_name = "amd64" if platform_name == "darwin" else machine_to_go.get(machine_name, machine_name)
+    with urllib.request.urlopen("https://api.github.com/repos/alttpo/sni/releases/latest") as request:
+        data = json.load(request)
+    files = data["assets"]
+
+    source_url = None
+
+    for file in files:
+        download_url: str = file["browser_download_url"]
+        machine_match = download_url.rsplit("-", 1)[1].split(".", 1)[0] == machine_name
+        if platform_name in download_url and machine_match:
+            # prefer "many" builds
+            if "many" in download_url:
+                source_url = download_url
+                break
+            source_url = download_url
+
+    if source_url and source_url.endswith(".zip"):
+        with urllib.request.urlopen(source_url) as download:
+            with zipfile.ZipFile(io.BytesIO(download.read()), "r") as zf:
+                for member in zf.infolist():
+                    zf.extract(member, path="SNI")
+        print(f"Downloaded SNI from {source_url}")
+
+    elif source_url and (source_url.endswith(".tar.xz") or source_url.endswith(".tar.gz")):
+        import tarfile
+        mode = "r:xz" if source_url.endswith(".tar.xz") else "r:gz"
+        with urllib.request.urlopen(source_url) as download:
+            sni_dir = None
+            with tarfile.open(fileobj=io.BytesIO(download.read()), mode=mode) as tf:
+                for member in tf.getmembers():
+                    if member.name.startswith("/") or "../" in member.name:
+                        raise ValueError(f"Unexpected file '{member.name}' in {source_url}")
+                    elif member.isdir() and not sni_dir:
+                        sni_dir = member.name
+                    elif member.isfile() and not sni_dir or not member.name.startswith(sni_dir):
+                        raise ValueError(f"Expected folder before '{member.name}' in {source_url}")
+                    elif member.isfile() and sni_dir:
+                        tf.extract(member)
+            # sadly SNI is in its own folder on non-windows, so we need to rename
+            shutil.rmtree("SNI", True)
+            os.rename(sni_dir, "SNI")
+        print(f"Downloaded SNI from {source_url}")
+
+    elif source_url:
+        print(f"Don't know how to extract SNI from {source_url}")
+
+    else:
+        print(f"No SNI found for system spec {platform_name} {machine_name}")
+
+
 if os.path.exists("X:/pw.txt"):
     print("Using signtool")
     with open("X:/pw.txt", encoding="utf-8-sig") as f:
@@ -174,6 +241,10 @@ class BuildExeCommand(cx_Freeze.command.build_exe.BuildEXE):
         print("Created Manifest")
 
     def run(self):
+        # start downloading sni asap
+        sni_thread = threading.Thread(target=download_SNI, name="SNI Downloader")
+        sni_thread.start()
+
         # pre build steps
         print(f"Outputting to: {self.buildfolder}")
         os.makedirs(self.buildfolder, exist_ok=True)
@@ -185,6 +256,9 @@ class BuildExeCommand(cx_Freeze.command.build_exe.BuildEXE):
         self.buildtime = datetime.datetime.utcnow()
         super().run()
 
+        # need to finish download before copying
+        sni_thread.join()
+
         # include_files seems to not be done automatically. implement here
         for src, dst in self.include_files:
             print(f"copying {src} -> {self.buildfolder / dst}")