From ef6d8336a1dc478666607f500218e44918a99017 Mon Sep 17 00:00:00 2001
From: Heitor Neiva <hneiva@mozilla.com>
Date: Tue, 2 Apr 2024 13:14:42 -0700
Subject: [PATCH 1/7] Bug 1889223 - Submit all files to notarization before
 polling

This is the first behavior that consumes all files at once, so it's a bit odd how it's being handled in script.py
---
 signingscript/src/signingscript/script.py | 18 ++++-
 signingscript/src/signingscript/sign.py   | 85 +++++++++++++++++++++--
 signingscript/src/signingscript/task.py   |  8 +++
 signingscript/src/signingscript/utils.py  |  7 +-
 signingscript/tests/test_script.py        | 17 +++++
 signingscript/tests/test_sign.py          | 27 +++++++
 signingscript/tests/test_task.py          |  1 +
 7 files changed, 154 insertions(+), 9 deletions(-)

diff --git a/signingscript/src/signingscript/script.py b/signingscript/src/signingscript/script.py
index 647b0cc79..6ea05093d 100755
--- a/signingscript/src/signingscript/script.py
+++ b/signingscript/src/signingscript/script.py
@@ -10,7 +10,7 @@
 import scriptworker.client
 
 from signingscript.exceptions import SigningScriptError
-from signingscript.task import build_filelist_dict, sign, task_cert_type, task_signing_formats
+from signingscript.task import apple_notarize_stacked, build_filelist_dict, sign, task_cert_type, task_signing_formats
 from signingscript.utils import copy_to_dir, load_apple_notarization_configs, load_autograph_configs
 
 log = logging.getLogger(__name__)
@@ -24,6 +24,7 @@ async def async_main(context):
         context (Context): the signing context.
 
     """
+    work_dir = context.config["work_dir"]
     async with aiohttp.ClientSession() as session:
         all_signing_formats = task_signing_formats(context)
         if "gpg" in all_signing_formats or "autograph_gpg" in all_signing_formats:
@@ -36,7 +37,7 @@ async def async_main(context):
             if not context.config.get("widevine_cert"):
                 raise Exception("Widevine format is enabled, but widevine_cert is not defined")
 
-        if "apple_notarization" in all_signing_formats or "apple_notarization_geckodriver" in all_signing_formats:
+        if {"apple_notarization", "apple_notarization_geckodriver", "apple_notarization_stacked"}.intersection(all_signing_formats):
             if not context.config.get("apple_notarization_configs", False):
                 raise Exception("Apple notarization is enabled but apple_notarization_configs is not defined")
             setup_apple_notarization_credentials(context)
@@ -44,9 +45,11 @@ async def async_main(context):
         context.session = session
         context.autograph_configs = load_autograph_configs(context.config["autograph_configs"])
 
-        work_dir = context.config["work_dir"]
         filelist_dict = build_filelist_dict(context)
         for path, path_dict in filelist_dict.items():
+            if path_dict["formats"] == ["apple_notarization_stacked"]:
+                # Skip if only format is notarization_stacked - handled below
+                continue
             copy_to_dir(path_dict["full_path"], context.config["work_dir"], target=path)
             log.info("signing %s", path)
             output_files = await sign(context, os.path.join(work_dir, path), path_dict["formats"], authenticode_comment=path_dict.get("comment"))
@@ -55,6 +58,15 @@ async def async_main(context):
                 copy_to_dir(os.path.join(work_dir, source), context.config["artifact_dir"], target=source)
             if "gpg" in path_dict["formats"] or "autograph_gpg" in path_dict["formats"]:
                 copy_to_dir(context.config["gpg_pubkey"], context.config["artifact_dir"], target="public/build/KEY")
+
+        # notarization_stacked is a special format that takes in all files at once instead of sequentially like other formats
+        notarization_dict = {path: path_dict for path, path_dict in filelist_dict.items() if "apple_notarization_stacked" in path_dict["formats"]}
+        if notarization_dict:
+            output_files = await apple_notarize_stacked(context, notarization_dict)
+            for source in output_files:
+                source = os.path.relpath(source, work_dir)
+                copy_to_dir(os.path.join(work_dir, source), context.config["artifact_dir"], target=source)
+
     log.info("Done!")
 
 
diff --git a/signingscript/src/signingscript/sign.py b/signingscript/src/signingscript/sign.py
index f0888ddc1..33e0fd375 100644
--- a/signingscript/src/signingscript/sign.py
+++ b/signingscript/src/signingscript/sign.py
@@ -1637,8 +1637,7 @@ async def apple_notarize(context, path, *args, **kwargs):
     """
     # Setup workdir
     notarization_workdir = os.path.join(context.config["work_dir"], "apple_notarize")
-    shutil.rmtree(notarization_workdir, ignore_errors=True)
-    utils.mkdir(notarization_workdir)
+    utils.mkdir(notarization_workdir, delete_before_create=True)
 
     _, extension = os.path.splitext(path)
     if extension == ".pkg":
@@ -1654,7 +1653,85 @@ async def apple_notarize_geckodriver(context, path, *args, **kwargs):
     """
     # Setup workdir
     notarization_workdir = os.path.join(context.config["work_dir"], "apple_notarize")
-    shutil.rmtree(notarization_workdir, ignore_errors=True)
-    utils.mkdir(notarization_workdir)
+    utils.mkdir(notarization_workdir, delete_before_create=True)
 
     return await _notarize_geckodriver(context, path, notarization_workdir)
+
+
+@time_async_function
+async def apple_notarize_stacked(context, filelist_dict):
+    """
+    Notarizes multiple packages using rcodesign.
+    Submits everything before polling for status.
+    """
+    ATTEMPTS = 5
+
+    # notarization submissions map (path -> submission_id)
+    submissions_map = {}
+    relpath_index_map = {}
+    task_index = 0
+    # Submit to notarization one by one
+    for relpath, path_dict in filelist_dict.items():
+        task_index += 1
+        relpath_index_map[relpath] = task_index
+        notarization_workdir = os.path.join(context.config["work_dir"], f"apple_notarize-{task_index}")
+        utils.mkdir(notarization_workdir, delete_before_create=True)
+        _, extension = os.path.splitext(relpath)
+        if extension == ".pkg":
+            path = os.path.join(notarization_workdir, relpath)
+            utils.copy_to_dir(path_dict["full_path"], notarization_workdir, target=relpath)
+            submissions_map[path] = await retry_async(
+                func=rcodesign_notarize,
+                args=(path, context.apple_credentials_path),
+                attempts=ATTEMPTS,
+                retry_exceptions=RCodesignError,
+            )
+        else:
+            await _extract_tarfile(context, path_dict["full_path"], extension, notarization_workdir)
+            workdir_files = os.listdir(notarization_workdir)
+            supported_files = [filename for filename in workdir_files if _can_notarize(filename, (".app", ".pkg"))]
+            if not supported_files:
+                raise SigningScriptError("No supported files found")
+            for file in supported_files:
+                path = os.path.join(notarization_workdir, file)
+                submissions_map[path] = await retry_async(
+                    func=rcodesign_notarize,
+                    args=(path, context.apple_credentials_path),
+                    attempts=ATTEMPTS,
+                    retry_exceptions=RCodesignError,
+                )
+
+    # Notary wait all files
+    for path, submission_id in submissions_map.items():
+        await retry_async(
+            func=rcodesign_notary_wait,
+            args=(submission_id, context.apple_credentials_path),
+            attempts=ATTEMPTS,
+            retry_exceptions=RCodesignError,
+        )
+
+    for path in submissions_map.keys():
+        await retry_async(
+            func=rcodesign_staple,
+            args=[path],
+            attempts=ATTEMPTS,
+            retry_exceptions=RCodesignError,
+        )
+
+    # Staple + create tarball where necessary
+    stapled_files = []
+    for relpath, path_dict in filelist_dict.items():
+        task_index = relpath_index_map[relpath]
+        notarization_workdir = os.path.join(context.config["work_dir"], f"apple_notarize-{task_index}")
+        target_path = os.path.join(context.config["work_dir"], relpath)
+        _, extension = os.path.splitext(relpath)
+        if extension == ".pkg":
+            utils.copy_to_dir(os.path.join(notarization_workdir, relpath), os.path.dirname(target_path))
+        else:
+            all_files = []
+            for root, _, files in os.walk(notarization_workdir):
+                for f in files:
+                    all_files.append(os.path.join(root, f))
+            await _create_tarfile(context, target_path, all_files, extension, notarization_workdir)
+        stapled_files.append(target_path)
+    return stapled_files
diff --git a/signingscript/src/signingscript/task.py b/signingscript/src/signingscript/task.py
index 7a8571ebd..1bed35cdf 100644
--- a/signingscript/src/signingscript/task.py
+++ b/signingscript/src/signingscript/task.py
@@ -18,6 +18,7 @@
 from signingscript.sign import (
     apple_notarize,
     apple_notarize_geckodriver,
+    apple_notarize_stacked,  # noqa: F401
     sign_authenticode,
     sign_debian_pkg,
     sign_file,
@@ -33,6 +34,10 @@
 
 log = logging.getLogger(__name__)
 
+
+async def noop_sign(*args, **kwargs):
+    return []
+
 FORMAT_TO_SIGNING_FUNCTION = immutabledict(
     {
         "autograph_hash_only_mar384": sign_mar384_with_autograph_hash,
@@ -58,6 +63,9 @@
         "autograph_rsa": sign_file_detached,
         "apple_notarization": apple_notarize,
         "apple_notarization_geckodriver": apple_notarize_geckodriver,
+        # This format is handled in script.py
+        # "apple_notarization_stacked": apple_notarize_stacked,
+        "apple_notarization_stacked": noop_sign,
         "default": sign_file,
     }
 )
diff --git a/signingscript/src/signingscript/utils.py b/signingscript/src/signingscript/utils.py
index 651581cd2..abd3ed865 100644
--- a/signingscript/src/signingscript/utils.py
+++ b/signingscript/src/signingscript/utils.py
@@ -6,6 +6,7 @@
 import json
 import logging
 import os
+import shutil
 from asyncio.subprocess import PIPE, STDOUT
 from dataclasses import dataclass
 from shutil import copyfile
@@ -35,13 +36,15 @@ class AppleNotarization:
     private_key: str
 
 
-def mkdir(path):
+def mkdir(path, delete_before_create=False):
     """Equivalent to `mkdir -p`.
 
     Args:
         path (str): the path to mkdir
-
+        delete_before_create (bool, optional): whether to delete the path before creating it. Defaults to False
     """
+    if delete_before_create:
+        shutil.rmtree(path, ignore_errors=True)
     try:
         os.makedirs(path)
         log.info("mkdir {}".format(path))
diff --git a/signingscript/tests/test_script.py b/signingscript/tests/test_script.py
index 7f2ce68b2..567a059dd 100644
--- a/signingscript/tests/test_script.py
+++ b/signingscript/tests/test_script.py
@@ -30,6 +30,9 @@ async def fake_sign(_, val, *args, authenticode_comment=None):
             assert authenticode_comment == "Some authenticode comment"
         return [val]
 
+    async def fake_notarize_stacked(_, filelist_dict, *args, **kwargs):
+        return filelist_dict.keys()
+
     mocker.patch.object(script, "load_autograph_configs", new=noop_sync)
     mocker.patch.object(script, "load_apple_notarization_configs", new=noop_sync)
     mocker.patch.object(script, "setup_apple_notarization_credentials", new=noop_sync)
@@ -37,6 +40,7 @@ async def fake_sign(_, val, *args, authenticode_comment=None):
     mocker.patch.object(script, "task_signing_formats", return_value=formats)
     mocker.patch.object(script, "build_filelist_dict", new=fake_filelist_dict)
     mocker.patch.object(script, "sign", new=fake_sign)
+    mocker.patch.object(script, "apple_notarize_stacked", new=fake_notarize_stacked)
     context = mock.MagicMock()
     context.config = {"work_dir": tmpdir, "artifact_dir": tmpdir, "autograph_configs": {}, "apple_notarization_configs": "fake"}
     context.config.update(extra_config)
@@ -99,6 +103,19 @@ async def test_async_main_apple_notarization(tmpdir, mocker):
     await async_main_helper(tmpdir, mocker, formats)
 
 
+@pytest.mark.asyncio
+async def test_async_main_apple_notarization_stacked(tmpdir, mocker):
+    formats = ["apple_notarization_stacked"]
+    mocker.patch.object(script, "copy_to_dir", new=noop_sync)
+    await async_main_helper(tmpdir, mocker, formats)
+
+
+@pytest.mark.asyncio
+async def test_async_main_apple_notarization_stacked_multi(tmpdir, mocker):
+    formats = ["autograph_mar", "apple_notarization_stacked"]
+    mocker.patch.object(script, "copy_to_dir", new=noop_sync)
+    await async_main_helper(tmpdir, mocker, formats)
+
 @pytest.mark.asyncio
 async def test_async_main_apple_notarization_no_config(tmpdir, mocker):
     formats = ["apple_notarization"]
diff --git a/signingscript/tests/test_sign.py b/signingscript/tests/test_sign.py
index c5b24b0b2..6e0afefab 100644
--- a/signingscript/tests/test_sign.py
+++ b/signingscript/tests/test_sign.py
@@ -23,6 +23,7 @@
 
 import signingscript.sign as sign
 import signingscript.utils as utils
+import signingscript.rcodesign as rcodesign
 from signingscript.exceptions import SigningScriptError
 from signingscript.utils import get_hash
 
@@ -1481,3 +1482,29 @@ async def test_apple_notarize_geckodriver(mocker, context):
 
     await sign.apple_notarize_geckodriver(context, "/foo/bar.pkg")
     notarize_geckodriver.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_apple_notarize_stacked(mocker, context):
+    notarize = mock.AsyncMock()
+    mocker.patch.object(sign, "rcodesign_notarize", notarize)
+    wait = mock.AsyncMock()
+    mocker.patch.object(sign, "rcodesign_notary_wait", wait)
+    staple = mock.AsyncMock()
+    mocker.patch.object(sign, "rcodesign_staple", staple)
+
+    mocker.patch.object(sign, "_extract_tarfile", noop_async)
+    mocker.patch.object(sign, "_create_tarfile", noop_async)
+    mocker.patch.object(sign.os, "listdir", lambda *_: ["/foo.pkg", "/baz.app", "/foobar"])
+    mocker.patch.object(sign.shutil, "rmtree", noop_sync)
+    mocker.patch.object(sign.utils, "mkdir", noop_sync)
+    mocker.patch.object(sign.utils, "copy_to_dir", noop_sync)
+
+    await sign.apple_notarize_stacked(
+        context,
+        {"/app.tar.gz": {"full_path": "/app.tar.gz", "formats": ["apple_notarize_stacked"]}}
+        )
+    # one for each file format
+    assert notarize.await_count == 2
+    assert wait.await_count == 2
+    assert staple.await_count == 2
diff --git a/signingscript/tests/test_task.py b/signingscript/tests/test_task.py
index 10793fd76..b505c43d1 100644
--- a/signingscript/tests/test_task.py
+++ b/signingscript/tests/test_task.py
@@ -151,6 +151,7 @@ def fake_log(context, new_files, *args):
         ("widevine", stask.sign_widevine),
         ("autograph_authenticode", stask.sign_authenticode),
         ("autograph_authenticode_stub", stask.sign_authenticode),
+        ("apple_notarization", stask.apple_notarize),
         ("default", stask.sign_file),
         # Key id cases
         ("autograph_hash_only_mar384:firefox_20190321_dev", stask.sign_mar384_with_autograph_hash),

From f69f61e3fa8bc738e27624a4c67f31a8c241ec64 Mon Sep 17 00:00:00 2001
From: Heitor Neiva <hneiva@mozilla.com>
Date: Thu, 11 Apr 2024 12:26:29 -0700
Subject: [PATCH 2/7] ci(signingscript): PR changes + Prevent mixed formats

Removed noop sign sign in tasks.py since we don't allow for mixed formats anymore
---
 signingscript/src/signingscript/script.py |  5 +++++
 signingscript/src/signingscript/sign.py   | 12 ++++++++++--
 signingscript/src/signingscript/task.py   |  4 ----
 signingscript/tests/test_script.py        |  6 ++++--
 4 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/signingscript/src/signingscript/script.py b/signingscript/src/signingscript/script.py
index 6ea05093d..8b3f2ef3f 100755
--- a/signingscript/src/signingscript/script.py
+++ b/signingscript/src/signingscript/script.py
@@ -45,11 +45,16 @@ async def async_main(context):
         context.session = session
         context.autograph_configs = load_autograph_configs(context.config["autograph_configs"])
 
+        # TODO: Make task.sign take in the whole filelist_dict and return a dict of output files.
+        #       That would likely mean changing all behaviors to accept and deal with multiple files at once.
+
         filelist_dict = build_filelist_dict(context)
         for path, path_dict in filelist_dict.items():
             if path_dict["formats"] == ["apple_notarization_stacked"]:
                 # Skip if only format is notarization_stacked - handled below
                 continue
+            if "apple_notarization_stacked" in path_dict["formats"]:
+                raise SigningScriptError("apple_notarization_stacked cannot be mixed with other signing types")
             copy_to_dir(path_dict["full_path"], context.config["work_dir"], target=path)
             log.info("signing %s", path)
             output_files = await sign(context, os.path.join(work_dir, path), path_dict["formats"], authenticode_comment=path_dict.get("comment"))
diff --git a/signingscript/src/signingscript/sign.py b/signingscript/src/signingscript/sign.py
index 33e0fd375..64a639866 100644
--- a/signingscript/src/signingscript/sign.py
+++ b/signingscript/src/signingscript/sign.py
@@ -1604,7 +1604,11 @@ async def _notarize_geckodriver(context, path, workdir):
 
 
 async def _notarize_all(context, path, workdir):
-    """Notarizes all files in a tarball"""
+    """
+    Notarizes all files in a tarball
+
+    @Deprecated: This function is deprecated and will be removed in the future. Use apple_notarize_stacked instead.
+    """
     _, extension = os.path.splitext(path)
     # Attempt extracting
     await _extract_tarfile(context, path, extension, tmp_dir=workdir)
@@ -1634,6 +1638,8 @@ async def _notarize_all(context, path, workdir):
 async def apple_notarize(context, path, *args, **kwargs):
     """
     Notarizes given package(s) using rcodesign.
+
+    @Deprecated: This function is deprecated and will be removed in the future. Use apple_notarize_stacked instead.
     """
     # Setup workdir
     notarization_workdir = os.path.join(context.config["work_dir"], "apple_notarize")
@@ -1710,6 +1716,7 @@ async def apple_notarize_stacked(context, filelist_dict):
             retry_exceptions=RCodesignError,
         )
 
+    # Staple files
     for path in submissions_map.keys():
         await retry_async(
             func=rcodesign_staple,
@@ -1718,13 +1725,14 @@ async def apple_notarize_stacked(context, filelist_dict):
             retry_exceptions=RCodesignError,
         )
 
-    # Staple + create tarball where necessary
+    # Wrap up
     stapled_files = []
     for relpath, path_dict in filelist_dict.items():
         task_index = relpath_index_map[relpath]
         notarization_workdir = os.path.join(context.config["work_dir"], f"apple_notarize-{task_index}")
         target_path = os.path.join(context.config["work_dir"], relpath)
         _, extension = os.path.splitext(relpath)
+        # Pkgs don't need to be tarred
         if extension == ".pkg":
             utils.copy_to_dir(os.path.join(notarization_workdir, relpath), os.path.dirname(target_path))
         else:
diff --git a/signingscript/src/signingscript/task.py b/signingscript/src/signingscript/task.py
index 1bed35cdf..306d103f1 100644
--- a/signingscript/src/signingscript/task.py
+++ b/signingscript/src/signingscript/task.py
@@ -35,9 +35,6 @@
 log = logging.getLogger(__name__)
 
 
-async def noop_sign(*args, **kwargs):
-    return []
-
 FORMAT_TO_SIGNING_FUNCTION = immutabledict(
     {
         "autograph_hash_only_mar384": sign_mar384_with_autograph_hash,
@@ -65,7 +62,6 @@ async def noop_sign(*args, **kwargs):
         "apple_notarization_geckodriver": apple_notarize_geckodriver,
         # This format is handled in script.py
         # "apple_notarization_stacked": apple_notarize_stacked,
-        "apple_notarization_stacked": noop_sign,
         "default": sign_file,
     }
 )
diff --git a/signingscript/tests/test_script.py b/signingscript/tests/test_script.py
index 567a059dd..f6429a0ed 100644
--- a/signingscript/tests/test_script.py
+++ b/signingscript/tests/test_script.py
@@ -111,10 +111,12 @@ async def test_async_main_apple_notarization_stacked(tmpdir, mocker):
 
 
 @pytest.mark.asyncio
-async def test_async_main_apple_notarization_stacked_multi(tmpdir, mocker):
+async def test_async_main_apple_notarization_stacked_mixed_fail(tmpdir, mocker):
     formats = ["autograph_mar", "apple_notarization_stacked"]
     mocker.patch.object(script, "copy_to_dir", new=noop_sync)
-    await async_main_helper(tmpdir, mocker, formats)
+    with pytest.raises(SigningScriptError):
+        await async_main_helper(tmpdir, mocker, formats)
+
 
 @pytest.mark.asyncio
 async def test_async_main_apple_notarization_no_config(tmpdir, mocker):

From a1e00096b98616724e2514936e3f376e4470e4b5 Mon Sep 17 00:00:00 2001
From: Heitor Neiva <hneiva@mozilla.com>
Date: Tue, 23 Apr 2024 11:30:26 -0700
Subject: [PATCH 3/7] nit(signingscript): Add reference issue comments for
 stacked behavior

---
 signingscript/src/signingscript/script.py | 1 +
 signingscript/src/signingscript/task.py   | 1 +
 2 files changed, 2 insertions(+)

diff --git a/signingscript/src/signingscript/script.py b/signingscript/src/signingscript/script.py
index 8b3f2ef3f..323670613 100755
--- a/signingscript/src/signingscript/script.py
+++ b/signingscript/src/signingscript/script.py
@@ -65,6 +65,7 @@ async def async_main(context):
                 copy_to_dir(context.config["gpg_pubkey"], context.config["artifact_dir"], target="public/build/KEY")
 
         # notarization_stacked is a special format that takes in all files at once instead of sequentially like other formats
+        # Should be fixed in https://github.com/mozilla-releng/scriptworker-scripts/issues/980
         notarization_dict = {path: path_dict for path, path_dict in filelist_dict.items() if "apple_notarization_stacked" in path_dict["formats"]}
         if notarization_dict:
             output_files = await apple_notarize_stacked(context, notarization_dict)
diff --git a/signingscript/src/signingscript/task.py b/signingscript/src/signingscript/task.py
index 306d103f1..ae82b96d5 100644
--- a/signingscript/src/signingscript/task.py
+++ b/signingscript/src/signingscript/task.py
@@ -61,6 +61,7 @@
         "apple_notarization": apple_notarize,
         "apple_notarization_geckodriver": apple_notarize_geckodriver,
         # This format is handled in script.py
+        # Should be refactored in https://github.com/mozilla-releng/scriptworker-scripts/issues/980
         # "apple_notarization_stacked": apple_notarize_stacked,
         "default": sign_file,
     }

From c26bafa03741833b9bcc352e3e1a1fa60fcde2fd Mon Sep 17 00:00:00 2001
From: Heitor Neiva <hneiva@mozilla.com>
Date: Tue, 23 Apr 2024 11:34:37 -0700
Subject: [PATCH 4/7] nit(signingscript): Lint fix

---
 signingscript/tests/test_sign.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/signingscript/tests/test_sign.py b/signingscript/tests/test_sign.py
index 6e0afefab..398559e2e 100644
--- a/signingscript/tests/test_sign.py
+++ b/signingscript/tests/test_sign.py
@@ -23,7 +23,6 @@
 
 import signingscript.sign as sign
 import signingscript.utils as utils
-import signingscript.rcodesign as rcodesign
 from signingscript.exceptions import SigningScriptError
 from signingscript.utils import get_hash
 
@@ -1500,10 +1499,7 @@ async def test_apple_notarize_stacked(mocker, context):
     mocker.patch.object(sign.utils, "mkdir", noop_sync)
     mocker.patch.object(sign.utils, "copy_to_dir", noop_sync)
 
-    await sign.apple_notarize_stacked(
-        context,
-        {"/app.tar.gz": {"full_path": "/app.tar.gz", "formats": ["apple_notarize_stacked"]}}
-        )
+    await sign.apple_notarize_stacked(context, {"/app.tar.gz": {"full_path": "/app.tar.gz", "formats": ["apple_notarize_stacked"]}})
     # one for each file format
     assert notarize.await_count == 2
     assert wait.await_count == 2

From df02a5c32125550ec34f852c777f4656735ff8a5 Mon Sep 17 00:00:00 2001
From: Heitor Neiva <hneiva@mozilla.com>
Date: Thu, 25 Apr 2024 10:12:19 -0700
Subject: [PATCH 5/7] nit(signingscript): Revert utils.mkdir parameter for
 deleting before creating

---
 signingscript/src/signingscript/sign.py  | 8 +++++---
 signingscript/src/signingscript/utils.py | 6 +-----
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/signingscript/src/signingscript/sign.py b/signingscript/src/signingscript/sign.py
index 64a639866..f899edc40 100644
--- a/signingscript/src/signingscript/sign.py
+++ b/signingscript/src/signingscript/sign.py
@@ -1643,7 +1643,7 @@ async def apple_notarize(context, path, *args, **kwargs):
     """
     # Setup workdir
     notarization_workdir = os.path.join(context.config["work_dir"], "apple_notarize")
-    utils.mkdir(notarization_workdir, delete_before_create=True)
+    utils.mkdir(notarization_workdir)
 
     _, extension = os.path.splitext(path)
     if extension == ".pkg":
@@ -1659,7 +1659,8 @@ async def apple_notarize_geckodriver(context, path, *args, **kwargs):
     """
     # Setup workdir
     notarization_workdir = os.path.join(context.config["work_dir"], "apple_notarize")
-    utils.mkdir(notarization_workdir, delete_before_create=True)
+    shutil.rmtree(notarization_workdir, ignore_errors=True)
+    utils.mkdir(notarization_workdir)
 
     return await _notarize_geckodriver(context, path, notarization_workdir)
 
@@ -1681,7 +1682,8 @@ async def apple_notarize_stacked(context, filelist_dict):
         task_index += 1
         relpath_index_map[relpath] = task_index
         notarization_workdir = os.path.join(context.config["work_dir"], f"apple_notarize-{task_index}")
-        utils.mkdir(notarization_workdir, delete_before_create=True)
+        shutil.rmtree(notarization_workdir, ignore_errors=True)
+        utils.mkdir(notarization_workdir)
         _, extension = os.path.splitext(relpath)
         if extension == ".pkg":
             path = os.path.join(notarization_workdir, relpath)
diff --git a/signingscript/src/signingscript/utils.py b/signingscript/src/signingscript/utils.py
index abd3ed865..977b1d6be 100644
--- a/signingscript/src/signingscript/utils.py
+++ b/signingscript/src/signingscript/utils.py
@@ -6,7 +6,6 @@
 import json
 import logging
 import os
-import shutil
 from asyncio.subprocess import PIPE, STDOUT
 from dataclasses import dataclass
 from shutil import copyfile
@@ -36,15 +35,12 @@ class AppleNotarization:
     private_key: str
 
 
-def mkdir(path, delete_before_create=False):
+def mkdir(path):
     """Equivalent to `mkdir -p`.
 
     Args:
         path (str): the path to mkdir
-        delete_before_create (bool, optional): whether to delete the path before creating it. Defaults to False
     """
-    if delete_before_create:
-        shutil.rmtree(path, ignore_errors=True)
     try:
         os.makedirs(path)
         log.info("mkdir {}".format(path))

From 7925b96895a1fb9f1fe6134559e4e89ce052c96a Mon Sep 17 00:00:00 2001
From: Heitor Neiva <hneiva@mozilla.com>
Date: Thu, 25 Apr 2024 10:50:23 -0700
Subject: [PATCH 6/7] fix(signingscript): Check all files before submitting for
 notarization

---
 signingscript/src/signingscript/sign.py | 37 +++++++++++++------------
 1 file changed, 20 insertions(+), 17 deletions(-)

diff --git a/signingscript/src/signingscript/sign.py b/signingscript/src/signingscript/sign.py
index f899edc40..1c36687cb 100644
--- a/signingscript/src/signingscript/sign.py
+++ b/signingscript/src/signingscript/sign.py
@@ -1673,11 +1673,11 @@ async def apple_notarize_stacked(context, filelist_dict):
     """
     ATTEMPTS = 5
 
-    # notarization submissions map (path -> submission_id)
-    submissions_map = {}
     relpath_index_map = {}
+    paths_to_notarize = []
     task_index = 0
-    # Submit to notarization one by one
+
+    # Create list of files to be notarized + check for potential problems
     for relpath, path_dict in filelist_dict.items():
         task_index += 1
         relpath_index_map[relpath] = task_index
@@ -1688,26 +1688,29 @@ async def apple_notarize_stacked(context, filelist_dict):
         if extension == ".pkg":
             path = os.path.join(notarization_workdir, relpath)
             utils.copy_to_dir(path_dict["full_path"], notarization_workdir, target=relpath)
-            submissions_map[path] = await retry_async(
-                func=rcodesign_notarize,
-                args=(path, context.apple_credentials_path),
-                attempts=ATTEMPTS,
-                retry_exceptions=RCodesignError,
-            )
-        else:
+            paths_to_notarize.append(path)
+        elif extension == ".gz":
             await _extract_tarfile(context, path_dict["full_path"], extension, notarization_workdir)
             workdir_files = os.listdir(notarization_workdir)
             supported_files = [filename for filename in workdir_files if _can_notarize(filename, (".app", ".pkg"))]
             if not supported_files:
-                raise SigningScriptError("No supported files found")
+                raise SigningScriptError(f"No supported files found for file {relpath}")
             for file in supported_files:
                 path = os.path.join(notarization_workdir, file)
-                submissions_map[path] = await retry_async(
-                    func=rcodesign_notarize,
-                    args=(path, context.apple_credentials_path),
-                    attempts=ATTEMPTS,
-                    retry_exceptions=RCodesignError,
-                )
+                paths_to_notarize.append(path)
+        else:
+            raise SigningScriptError(f"Unsupported file extension: {extension} for file {relpath}")
+
+    # notarization submissions map (path -> submission_id)
+    submissions_map = {}
+    # Submit to notarization one by one
+    for path in paths_to_notarize:
+        submissions_map[path] = await retry_async(
+            func=rcodesign_notarize,
+            args=(path, context.apple_credentials_path),
+            attempts=ATTEMPTS,
+            retry_exceptions=RCodesignError,
+        )
 
     # Notary wait all files
     for path, submission_id in submissions_map.items():

From 12d3b778f78797d1a6d0222e6fbd6f6a224f61ed Mon Sep 17 00:00:00 2001
From: Heitor Neiva <hneiva@mozilla.com>
Date: Thu, 25 Apr 2024 14:53:09 -0700
Subject: [PATCH 7/7] chore(signingscript): Add tests for
 apple_notarize_stacked

---
 signingscript/tests/test_sign.py | 46 +++++++++++++++++++++++++++++---
 1 file changed, 42 insertions(+), 4 deletions(-)

diff --git a/signingscript/tests/test_sign.py b/signingscript/tests/test_sign.py
index 398559e2e..2e5a57c39 100644
--- a/signingscript/tests/test_sign.py
+++ b/signingscript/tests/test_sign.py
@@ -1495,12 +1495,50 @@ async def test_apple_notarize_stacked(mocker, context):
     mocker.patch.object(sign, "_extract_tarfile", noop_async)
     mocker.patch.object(sign, "_create_tarfile", noop_async)
     mocker.patch.object(sign.os, "listdir", lambda *_: ["/foo.pkg", "/baz.app", "/foobar"])
+    mocker.patch.object(sign.os, "walk", lambda *_: [("/", None, ["foo.pkg", "baz.app"])])
     mocker.patch.object(sign.shutil, "rmtree", noop_sync)
     mocker.patch.object(sign.utils, "mkdir", noop_sync)
     mocker.patch.object(sign.utils, "copy_to_dir", noop_sync)
 
-    await sign.apple_notarize_stacked(context, {"/app.tar.gz": {"full_path": "/app.tar.gz", "formats": ["apple_notarize_stacked"]}})
+    await sign.apple_notarize_stacked(
+        context,
+        {
+            "/app.tar.gz": {"full_path": "/app.tar.gz", "formats": ["apple_notarize_stacked"]},
+            "/app2.pkg": {"full_path": "/app2.pkg", "formats": ["apple_notarize_stacked"]},
+        },
+    )
     # one for each file format
-    assert notarize.await_count == 2
-    assert wait.await_count == 2
-    assert staple.await_count == 2
+    assert notarize.await_count == 3
+    assert wait.await_count == 3
+    assert staple.await_count == 3
+
+
+@pytest.mark.asyncio
+async def test_apple_notarize_stacked_unsupported(mocker, context):
+    """Test unsupported file extensions"""
+
+    mocker.patch.object(sign, "_extract_tarfile", noop_async)
+    mocker.patch.object(sign.shutil, "rmtree", noop_sync)
+    mocker.patch.object(sign.utils, "mkdir", noop_sync)
+    mocker.patch.object(sign.utils, "copy_to_dir", noop_sync)
+
+    # Returns unsupported file formats
+    mocker.patch.object(sign.os, "listdir", lambda *_: ["/foo.aaa", "/baz.bbb", "/foobar"])
+
+    with pytest.raises(SigningScriptError):
+        # Main file is supported, contents uses the above os.listdir
+        await sign.apple_notarize_stacked(
+            context,
+            {
+                "/app.tar.gz": {"full_path": "/app.tar.gz", "formats": ["apple_notarize_stacked"]},
+            },
+        )
+
+    with pytest.raises(SigningScriptError):
+        # Main file extension is unsupported
+        await sign.apple_notarize_stacked(
+            context,
+            {
+                "/app.bbb": {"full_path": "/app.bbb", "formats": ["apple_notarize_stacked"]},
+            },
+        )