1 files changed, 171 insertions, 0 deletions
diff --git a/comm/taskcluster/docker/tb-atn/atn_langpack.py b/comm/taskcluster/docker/tb-atn/atn_langpack.py
new file mode 100644
index 0000000000..fc66c09b2f
--- /dev/null
+++ b/comm/taskcluster/docker/tb-atn/atn_langpack.py
@@ -0,0 +1,171 @@
+#!python3
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+import datetime
+import json
+import logging
+import os
+import random
+import string
+import sys
+from enum import Enum
+from pathlib import Path
+from pprint import pprint as pp
+from typing import List, Literal, NamedTuple, Tuple, Union
+
+import jwt
+import requests
+from redo import retry
+
+logging.getLogger("requests").setLevel(logging.DEBUG)
+
+ATN_UPLOAD_URL = "https://addons.thunderbird.net/api/v3/addons/langpack-{langcode}@thunderbird.mozilla.org/versions/{version}/"
+CHUNK_SIZE = 128 * 1024
+
+
+class ATNChannel(Enum):
+    LISTED = "listed"
+    UNLISTED = "unlisted"
+
+
+Locales = List[str]
+Version = str
+ApiParam = str
+EnvVars = NamedTuple(
+    "EnvVars",
+    [
+        ("langpack_version", Version),
+        ("locales", Locales),
+        ("langpack_dir", Path),
+        ("langpack_channel", Literal[ATNChannel.LISTED, ATNChannel.UNLISTED]),
+        ("api_key", ApiParam),
+        ("api_secret", ApiParam),
+    ],
+)
+Result = Tuple[str, Union[object, None]]
+
+
+def print_line(message):
+    msg_bytes = message.encode("utf-8")
+    written = 0
+    while written < len(msg_bytes):
+        written += sys.stdout.buffer.write(msg_bytes[written:]) or 0
+    sys.stdout.buffer.flush()
+
+
+class ATNUploader:
+    def __init__(self, options: EnvVars):
+        self.api_key = options.api_key
+        self.api_secret = options.api_secret
+        self.langpack_dir = options.langpack_dir
+        self.langpack_version = options.langpack_version
+        self.langpack_channel = options.langpack_channel
+        self.locales = options.locales
+
+    def mk_headers(self) -> dict:
+        now = datetime.datetime.utcnow()
+        payload = {
+            "iss": self.api_key,
+            "jti": "".join(
+                random.choice(string.ascii_uppercase + string.digits) for _ in range(64)
+            ),
+            "exp": now + datetime.timedelta(seconds=60),
+            "iat": now,
+        }
+        headers = {
+            "Authorization": "JWT {0}".format(
+                jwt.encode(payload, self.api_secret, algorithm="HS256")
+            )
+        }
+        return headers
+
+    def upload_langpack(self, locale: str) -> Result:
+        langpack_path = self.langpack_dir / locale / "target.langpack.xpi"
+        headers = self.mk_headers()
+        langpack_fd = open(langpack_path, "rb")
+        file = {"upload": ("upload", langpack_fd)}
+        data = {"channel": self.langpack_channel}
+
+        url = ATN_UPLOAD_URL.format(version=self.langpack_version, langcode=locale)
+        with requests.put(url, files=file, data=data, headers=headers, verify=False) as resp:
+            if not resp.ok:
+                print_line(f"Failed {locale}")
+                return resp.json()
+            else:
+                return resp.json()
+
+    def upload_all_locales(self) -> Tuple[List[Result], List[Result]]:
+        failed = []
+        success = []
+        for locale in self.locales:
+            try:
+                rv = retry(self.upload_langpack, args=(locale,), attempts=3, sleeptime=10)
+                if "error" not in rv:
+                    success.append((locale, rv))
+                else:
+                    failed.append((locale, rv))
+            except requests.HTTPError as e:
+                print_line(e)
+                failed.append((locale, None))
+        return success, failed
+
+
+def get_secret(name: str) -> Tuple[ApiParam, ApiParam]:
+    secret = {}
+    if "MOZ_SCM_LEVEL" in os.environ:
+        level = os.environ.get("MOZ_SCM_LEVEL", "1")
+        taskcluster_url = os.environ.get("TASKCLUSTER_PROXY_URL") or os.environ.get(
+            "TASKCLUSTER_ROOT_URL", ""
+        )
+        secrets_url = (
+            f"{taskcluster_url}/secrets/v1/secret/project/comm/thunderbird/releng"
+            f"/build/level-{level}/{name}"
+        )
+        res = requests.get(secrets_url)
+        res.raise_for_status()
+        secret = res.json()
+    elif "SECRET_FILE" in os.environ:  # For local dev/debug
+        with open(os.environ["SECRET_FILE"]) as fp:
+            secret = json.load(fp)["secret"]
+    secret = secret.get("secret")
+    api_key = secret["api_key"] if "api_key" in secret else None
+    api_secret = secret["api_secret"] if "api_secret" in secret else None
+    if api_key is None or api_secret is None:
+        raise Exception(f"Unable to get secret. {secret.keys()}")
+
+    return api_key, api_secret
+
+
+def read_env_vars() -> EnvVars:
+    try:
+        langpack_version = os.environ["LANGPACK_VERSION"]
+        locales_json = os.environ["LOCALES"]
+        langpack_dir = Path(os.environ["MOZ_FETCHES_DIR"]).resolve()
+        langpack_channel = os.environ["ATN_CHANNEL"]
+    except KeyError:
+        raise Exception("Missing environment variable(s)")
+
+    locales = json.loads(locales_json)
+    api_key, api_secret = get_secret("atn_langpack")
+
+    return EnvVars(
+        langpack_version, locales, langpack_dir, ATNChannel(langpack_channel), api_key, api_secret
+    )
+
+
+def main():
+    options = read_env_vars()
+
+    atn_uploader = ATNUploader(options)
+    success, failed = atn_uploader.upload_all_locales()
+
+    pp(success)
+    if failed:
+        pp(failed)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()