unpack-sdk.py - mozsearch

Enable keyboard shortcuts

# 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 http://mozilla.org/MPL/2.0/.

import hashlib

import os

import shutil

import stat

import sys

import tempfile

from io import BytesIO

from urllib.request import urlopen

from mozpack.macpkg import Pbzx, uncpio, unxar

def unpack_sdk(url, sha256, extract_prefix, out_dir="."):

    with tempfile.TemporaryFile() as pkg:

        hash = hashlib.sha256()

        for attempt in range(3):

            if attempt != 0:

                print(f"Failed to download from {url}. Retrying", file=sys.stderr)

            with urlopen(url) as fh:

                # Equivalent to shutil.copyfileobj, but computes sha256 at the same time.

                while True:

                    buf = fh.read(1024 * 1024)

                    if not buf:

                        break

                    hash.update(buf)

                    pkg.write(buf)

            digest = hash.hexdigest()

            if digest == sha256:

                break

        else:

            raise Exception(f"(actual) {digest} != (expected) {sha256}")

        pkg.seek(0, os.SEEK_SET)

        for name, content in unxar(pkg):

            if name in ("Payload", "Content"):

                extract_payload(content, extract_prefix, out_dir)

def extract_payload(fileobj, extract_prefix, out_dir="."):

    hardlinks = {}

    for path, st, content in uncpio(Pbzx(fileobj)):

        if not path:

            continue

        path = path.decode()

        matches = path.startswith(extract_prefix)

        if matches:

            path = os.path.join(out_dir, path[len(extract_prefix) :].lstrip("/"))

        # When there are hardlinks, normally a cpio stream is supposed to

        # contain the data for all of them, but, even with compression, that

        # can be a waste of space, so in some cpio streams (*cough* *cough*,

        # Apple's, e.g. in Xcode), the files after the first one have dummy

        # data.

        # As we may be filtering the first file out (if it doesn't match

        # extract_prefix), we need to keep its data around (we're not going

        # to be able to rewind).

        if stat.S_ISREG(st.mode) and st.nlink > 1:

            key = (st.dev, st.ino)

            hardlink = hardlinks.get(key)

            if hardlink:

                hardlink[0] -= 1

                if hardlink[0] == 0:

                    del hardlinks[key]

                content = hardlink[1]

                if isinstance(content, BytesIO):

                    content.seek(0)

                    if matches:

                        hardlink[1] = path

            elif matches:

                hardlink = hardlinks[key] = [st.nlink - 1, path]

            else:

                hardlink = hardlinks[key] = [st.nlink - 1, BytesIO(content.read())]

                content = hardlink[1]

        if not matches:

            continue

        if stat.S_ISDIR(st.mode):

            os.makedirs(path, exist_ok=True)

        else:

            parent = os.path.dirname(path)

            if parent:

                os.makedirs(parent, exist_ok=True)

            if stat.S_ISLNK(st.mode):

                os.symlink(content.read(), path)

            elif stat.S_ISREG(st.mode):

                if isinstance(content, str):

                    os.link(content, path)

                else:

                    with open(path, "wb") as out:

                        shutil.copyfileobj(content, out)

            else:

                raise Exception(f"File mode {st.mode:o} is not supported")

if __name__ == "__main__":

    unpack_sdk(*sys.argv[1:])