Skip to content

Commit 4350267

Browse files
committed
Add support for downloading a single file and putting it in an archive
1 parent 5a03d0f commit 4350267

File tree

1 file changed

+49
-5
lines changed

1 file changed

+49
-5
lines changed

src/taskgraph/run-task/fetch-content

Lines changed: 49 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -480,12 +480,34 @@ def should_repack_archive(
480480
return True
481481

482482

483+
EXECUTABLE_SIGNATURES = set([
484+
b"\xFE\xED\xFA\xCE", # mach-o 32-bits big endian
485+
b"\xCE\xFA\xED\xFE", # mach-o 32-bits little endian
486+
b"\xFE\xED\xFA\xCF", # mach-o 64-bits big endian
487+
b"\xCF\xFA\xED\xFE", # mach-o 64-bits little endian
488+
b"\xCA\xFE\xBA\xBE", # mach-o FAT binary
489+
b"\x7F\x45\x4C\x46", # Elf binary
490+
])
491+
492+
483493
def repack_archive(
484-
orig: pathlib.Path, dest: pathlib.Path, strip_components=0, prefix=""
494+
orig: pathlib.Path,
495+
dest: pathlib.Path,
496+
strip_components=0,
497+
prefix="",
498+
force_archive=False,
485499
):
486500
assert orig != dest
487501
log(f"Repacking {orig} as {dest}")
488-
orig_typ, ifh = open_stream(orig)
502+
try:
503+
orig_typ, ifh = open_stream(orig)
504+
except ArchiveTypeNotSupported:
505+
if force_archive:
506+
ifh = io.BufferedReader(orig.open(mode="rb"))
507+
signature = ifh.peek(4)[:4]
508+
orig_typ = "exec" if signature in EXECUTABLE_SIGNATURES else None
509+
else:
510+
raise
489511
typ = archive_type(dest)
490512
if not typ:
491513
raise Exception("Archive type not supported for %s" % dest.name)
@@ -510,7 +532,20 @@ def repack_archive(
510532

511533
with rename_after_close(dest, "wb") as fh:
512534
ctx = ZstdCompressor()
513-
if orig_typ == "zip":
535+
if orig_typ in ("exec", None):
536+
with ctx.stream_writer(fh) as compressor, tarfile.open(
537+
fileobj=compressor,
538+
mode="w:",
539+
) as tar:
540+
tarinfo = tarfile.TarInfo()
541+
tarinfo.name = filter(orig.name) if filter else orig.name
542+
st = orig.stat()
543+
tarinfo.size = st.st_size
544+
tarinfo.mtime = st.st_mtime
545+
tarinfo.mode = 0o0755 if orig_typ == "exec" else 0o0644
546+
tar.addfile(tarinfo, ifh)
547+
548+
elif orig_typ == "zip":
514549
assert typ == "tar"
515550
zip = zipfile.ZipFile(ifh)
516551
# Convert the zip stream to a tar on the fly.
@@ -824,8 +859,12 @@ def command_static_url(args):
824859
if gpg_sig_url:
825860
gpg_verify_path(dl_dest, gpg_key, gpg_signature)
826861

827-
if should_repack_archive(dl_dest, dest, args.strip_components, args.add_prefix):
828-
repack_archive(dl_dest, dest, args.strip_components, args.add_prefix)
862+
if args.force_archive or should_repack_archive(
863+
dl_dest, dest, args.strip_components, args.add_prefix
864+
):
865+
repack_archive(
866+
dl_dest, dest, args.strip_components, args.add_prefix, args.force_archive
867+
)
829868
elif dl_dest != dest:
830869
log(f"Renaming {dl_dest} to {dest}")
831870
dl_dest.rename(dest)
@@ -960,6 +999,11 @@ def main():
960999
dest="headers",
9611000
help="Header to send as part of the request, can be passed " "multiple times",
9621001
)
1002+
url.add_argument(
1003+
"--force-archive",
1004+
action="store_true",
1005+
help="Create an archive even when the downloaded file is not an archive",
1006+
)
9631007
url.add_argument("url", help="URL to fetch")
9641008
url.add_argument("dest", help="Destination path")
9651009

0 commit comments

Comments
 (0)