Skip to content

Commit

Permalink
Add support for downloading a single file and putting it in an archive
Browse files Browse the repository at this point in the history
  • Loading branch information
glandium committed Dec 26, 2024
1 parent 5a03d0f commit 28b3d03
Showing 1 changed file with 46 additions and 5 deletions.
51 changes: 46 additions & 5 deletions src/taskgraph/run-task/fetch-content
Original file line number Diff line number Diff line change
Expand Up @@ -480,12 +480,31 @@ def should_repack_archive(
return True


EXECUTABLE_SIGNATURES = [
b"\xFE\xED\xFA\xCE", # mach-o 32-bits big endian
b"\xCE\xFA\xED\xFE", # mach-o 32-bits little endian
b"\xFE\xED\xFA\xCF", # mach-o 64-bits big endian
b"\xCF\xFA\xED\xFE", # mach-o 64-bits little endian
b"\xCA\xFE\xBA\xBE", # mach-o FAT binary
b"\x7F\x45\x4C\x46", # Elf binary
]


def repack_archive(
orig: pathlib.Path, dest: pathlib.Path, strip_components=0, prefix=""
orig: pathlib.Path,
dest: pathlib.Path,
strip_components=0,
prefix="",
single_file=False,
):
assert orig != dest
log(f"Repacking {orig} as {dest}")
orig_typ, ifh = open_stream(orig)
if single_file:
ifh = io.BufferedReader(orig.open(mode="rb"))
signature = ifh.peek(4)[:4]
orig_typ = "exec" if signature in EXECUTABLE_SIGNATURES else None
else:
orig_typ, ifh = open_stream(orig)
typ = archive_type(dest)
if not typ:
raise Exception("Archive type not supported for %s" % dest.name)
Expand All @@ -510,7 +529,20 @@ def repack_archive(

with rename_after_close(dest, "wb") as fh:
ctx = ZstdCompressor()
if orig_typ == "zip":
if single_file:
with ctx.stream_writer(fh) as compressor, tarfile.open(
fileobj=compressor,
mode="w:",
) as tar:
tarinfo = tarfile.TarInfo()
tarinfo.name = filter(orig.name) if filter else orig.name
stat = orig.stat()
tarinfo.size = stat.st_size
tarinfo.mtime = stat.st_mtime
tarinfo.mode = 0o0755 if orig_typ == "exec" else 0o0644
tar.addfile(tarinfo, ifh)

elif orig_typ == "zip":
assert typ == "tar"
zip = zipfile.ZipFile(ifh)
# Convert the zip stream to a tar on the fly.
Expand Down Expand Up @@ -824,8 +856,12 @@ def command_static_url(args):
if gpg_sig_url:
gpg_verify_path(dl_dest, gpg_key, gpg_signature)

if should_repack_archive(dl_dest, dest, args.strip_components, args.add_prefix):
repack_archive(dl_dest, dest, args.strip_components, args.add_prefix)
if args.single_file or should_repack_archive(
dl_dest, dest, args.strip_components, args.add_prefix
):
repack_archive(
dl_dest, dest, args.strip_components, args.add_prefix, args.single_file
)
elif dl_dest != dest:
log(f"Renaming {dl_dest} to {dest}")
dl_dest.rename(dest)
Expand Down Expand Up @@ -960,6 +996,11 @@ def main():
dest="headers",
help="Header to send as part of the request, can be passed " "multiple times",
)
url.add_argument(
"--single-file",
action="store_true",
help="The downloaded file is not an archive and should be put inside one",
)
url.add_argument("url", help="URL to fetch")
url.add_argument("dest", help="Destination path")

Expand Down

0 comments on commit 28b3d03

Please sign in to comment.